Quantum Computing in PyTorch

Faster, Scalable, Easy Debugging, Easy Deployment on Real Machine

Community fork notice

This repository is a community-maintained fork of mit-han-lab/torchquantum, kept alive while the upstream is on pause. The PyPI distribution name is torchquantum-ng, but the Python import name is unchanged — code that already uses import torchquantum keeps working.

# Recommended: remove the upstream first so the two don't overlay each
# other in site-packages (they install into the same `torchquantum/` dir).
pip uninstall -y torchquantum
pip install torchquantum-ng

import torchquantum as tq  # unchanged

Compared to upstream torchquantum==0.1.8 this fork additionally:

• tests on Python 3.8 – 3.13 (and an experimental 3.14-dev job);
• drops the unused dill pin and moves tensorflow / pyscf to optional extras (pip install 'torchquantum-ng[quantumnas]', pip install 'torchquantum-ng[chem]') so plain installs work on new Python releases without dragging in heavy native wheels.

👋 Welcome

What it is doing

Simulate quantum computations on classical hardware using PyTorch. It supports statevector simulation and pulse simulation on GPUs. It can scale up to the simulation of 30+ qubits with multiple GPUs.

Who will benefit

Researchers on quantum algorithm design, parameterized quantum circuit training, quantum optimal control, quantum machine learning, quantum neural networks.

Differences from Qiskit/Pennylane

Dynamic computation graph, automatic gradient computation, fast GPU support, batch model tersorized processing.

News

• Torchquantum is used in the winning team for ACM Quantum Computing for Drug Discovery Challenge.
• Torchquantum is highlighted in unitaryHACK.
• TorchQuantum received a Unitary Foundation microgrant.
• TorchQuantum is integrated to IBM Qiskit Ecosystem.
• TorchQuantum is integrated to PyTorch Ecosystem.
• v0.1.8 Available!
• Check the dev branch for new latest features on quantum layers and quantum algorithms.
• Join our Slack for real time support!
• Welcome to contribute! Please contact us or post in the Github Issues if you want to have new examples implemented by TorchQuantum or any other questions.
• Qmlsys website goes online: qmlsys.mit.edu and torchquantum.org

Features

• Easy construction and simulation of quantum circuits in PyTorch
• Dynamic computation graph for easy debugging
• Gradient support via autograd
• Batch mode inference and training on CPU/GPU
• Easy deployment on real quantum devices such as IBMQ
• Easy hybrid classical-quantum model construction
• (coming soon) pulse-level simulation

Installation

git clone https://github.com/mit-han-lab/torchquantum.git
cd torchquantum
pip install --editable .

Basic Usage

import torchquantum as tq
import torchquantum.functional as tqf

qdev = tq.QuantumDevice(n_wires=2, bsz=5, device="cpu", record_op=True) # use device='cuda' for GPU

# use qdev.op
qdev.h(wires=0)
qdev.cnot(wires=[0, 1])

# use tqf
tqf.h(qdev, wires=1)
tqf.x(qdev, wires=1)

# use tq.Operator
op = tq.RX(has_params=True, trainable=True, init_params=0.5)
op(qdev, wires=0)

# print the current state (dynamic computation graph supported)
print(qdev)

# obtain the qasm string
from torchquantum.plugin import op_history2qasm
print(op_history2qasm(qdev.n_wires, qdev.op_history))

# measure the state on z basis
print(tq.measure(qdev, n_shots=1024))

# obtain the expval on a observable by stochastic sampling (doable on simulator and real quantum hardware)
from torchquantum.measurement import expval_joint_sampling
expval_sampling = expval_joint_sampling(qdev, 'ZX', n_shots=1024)
print(expval_sampling)

# obtain the expval on a observable by analytical computation (only doable on classical simulator)
from torchquantum.measurement import expval_joint_analytical
expval = expval_joint_analytical(qdev, 'ZX')
print(expval)

# obtain gradients of expval w.r.t. trainable parameters
expval[0].backward()
print(op.params.grad)


# Apply gates to qdev with tq.QuantumModule
ops = [
    {'name': 'hadamard', 'wires': 0}, 
    {'name': 'cnot', 'wires': [0, 1]},
    {'name': 'rx', 'wires': 0, 'params': 0.5, 'trainable': True},
    {'name': 'u3', 'wires': 0, 'params': [0.1, 0.2, 0.3], 'trainable': True},
    {'name': 'h', 'wires': 1, 'inverse': True}
]

qmodule = tq.QuantumModule.from_op_history(ops)
qmodule(qdev)

Guide to the examples

We also prepare many example and tutorials using TorchQuantum.

For beginning level, you may check QNN for MNIST, Quantum Convolution (Quanvolution) and Quantum Kernel Method, and Quantum Regression.

For intermediate level, you may check Amplitude Encoding for MNIST, Clifford gate QNN, Save and Load QNN models, PauliSum Operation, How to convert tq to Qiskit.

For expert, you may check Parameter Shift on-chip Training, VQA Gradient Pruning, VQE, VQA for State Prepration, QAOA (Quantum Approximate Optimization Algorithm).

Usage

Construct parameterized quantum circuit models as simple as constructing a normal pytorch model.

import torch.nn as nn
import torch.nn.functional as F
import torchquantum as tq
import torchquantum.functional as tqf

class QFCModel(nn.Module):
  def __init__(self):
    super().__init__()
    self.n_wires = 4
    self.measure = tq.MeasureAll(tq.PauliZ)

    self.encoder_gates = [tqf.rx] * 4 + [tqf.ry] * 4 + \
                         [tqf.rz] * 4 + [tqf.rx] * 4
    self.rx0 = tq.RX(has_params=True, trainable=True)
    self.ry0 = tq.RY(has_params=True, trainable=True)
    self.rz0 = tq.RZ(has_params=True, trainable=True)
    self.crx0 = tq.CRX(has_params=True, trainable=True)

  def forward(self, x):
    bsz = x.shape[0]
    # down-sample the image
    x = F.avg_pool2d(x, 6).view(bsz, 16)

    # create a quantum device to run the gates
    qdev = tq.QuantumDevice(n_wires=self.n_wires, bsz=bsz, device=x.device)

    # encode the classical image to quantum domain
    for k, gate in enumerate(self.encoder_gates):
      gate(qdev, wires=k % self.n_wires, params=x[:, k])

    # add some trainable gates (need to instantiate ahead of time)
    self.rx0(qdev, wires=0)
    self.ry0(qdev, wires=1)
    self.rz0(qdev, wires=3)
    self.crx0(qdev, wires=[0, 2])

    # add some more non-parameterized gates (add on-the-fly)
    qdev.h(wires=3)
    qdev.sx(wires=2)
    qdev.cnot(wires=[3, 0])
    qdev.qubitunitary(wires=[1, 2], params=[[1, 0, 0, 0],
                                            [0, 1, 0, 0],
                                            [0, 0, 0, 1j],
                                            [0, 0, -1j, 0]])

    # perform measurement to get expectations (back to classical domain)
    x = self.measure(qdev).reshape(bsz, 2, 2)

    # classification
    x = x.sum(-1).squeeze()
    x = F.log_softmax(x, dim=1)

    return x

VQE Example

Train a quantum circuit to perform VQE task. Quito quantum computer as in simple_vqe.py script:

cd examples/vqe
python vqe.py

MNIST Example

Train a quantum circuit to perform MNIST classification task and deploy on the real IBM Quito quantum computer as in mnist_example.py script:

cd examples/mnist
python mnist.py

Files

Coding Style

torchquantum uses pre-commit hooks to ensure Python style consistency and prevent common mistakes in its codebase.

To enable it pre-commit hooks please reproduce:

pip install pre-commit
pre-commit install

Papers using TorchQuantum

• [HPCA'22] Wang et al., "QuantumNAS: Noise-Adaptive Search for Robust Quantum Circuits"
• [DAC'22] Wang et al., "QuantumNAT: Quantum Noise-Aware Training with Noise Injection, Quantization and Normalization"
• [DAC'22] Wang et al., "QOC: Quantum On-Chip Training with Parameter Shift and Gradient Pruning"
• [QCE'22] Liang et al., "Variational Quantum Pulse Learning"
• [ICCAD'22] Hu et al., "Quantum Neural Network Compression"
• [ICCAD'22] Wang et al., "QuEst: Graph Transformer for Quantum Circuit Reliability Estimation"
• [ICML Workshop] Yun et al., "Slimmable Quantum Federated Learning"
• [IEEE ICDCS] Yun et al., "Quantum Multi-Agent Reinforcement Learning via Variational Quantum Circuit Design"
• [QCE'23] Zhan et al., "Quantum Sensor Network Algorithms for Transmitter Localization"

Dependencies

• 3.9 >= Python >= 3.7 (Python 3.10 may have the concurrent package issue for Qiskit)
• PyTorch >= 1.8.0
• configargparse >= 0.14
• GPU model training requires NVIDIA GPUs

Contact

TorchQuantum Forum

Hanrui Wang hanrui@mit.edu

Contributors

Jiannan Cao, Jessica Ding, Jiai Gu, Song Han, Zhirui Hu, Zirui Li, Zhiding Liang, Pengyu Liu, Yilian Liu, Mohammadreza Tavasoli, Hanrui Wang, Zhepeng Wang, Zhuoyang Ye

Citation

@inproceedings{hanruiwang2022quantumnas,
    title     = {Quantumnas: Noise-adaptive search for robust quantum circuits},
    author    = {Wang, Hanrui and Ding, Yongshan and Gu, Jiaqi and Li, Zirui and Lin, Yujun and Pan, David Z and Chong, Frederic T and Han, Song},
    booktitle = {The 28th IEEE International Symposium on High-Performance Computer Architecture (HPCA-28)},
    year      = {2022}
}