Simulating 36-qubit quantum circuits at 549 GB scale. ~0.01ms per gate. 100% pure JAX.Accelerated on NVIDIA GPUs and Google Cloud TPU v6e-64 / v5e clusters. Supported by the Google TPU Research Cloud (TRC) program.
A high-performance, research-grade quantum state-vector simulator built purely in JAX. Run differentiable, noise-resilient, and large-scale quantum circuits accelerated on local NVIDIA GPUs and multi-worker Google Cloud TPU clusters.
100% Pure JAX: Zero dependencies on heavy frameworks (Qiskit, Cirq, Pennylane). Compiled natively into a single monolithic XLA kernel for bare-metal execution speeds. -
Multi-Device Sharding: Scale up to 36 qubits (549 GB state-vector footprint) using distributed JAXPositionalSharding
across a 64-chip Cloud TPU v6e mesh. -
Reverse-Mode Auto-Differentiation: Compute exact gradients in a single backward pass viajax.grad
for fast training of variational algorithms (VQE, QAOA, QNNs). -
Hardware-Level Optimizations: Structured loop primitives (jax.lax.fori_loop
) prevent XLA graph bloat, whilejax.checkpoint
(gradient rematerialization) keeps memory complexity at$\mathcal{O}(1)$ . - Stochastic Noise Support: Built-in Monte Carlo trajectory simulations for open systems and depolarizing NISQ gate noise.
.
├── gpu/ # GPU Modular Simulator & Research scripts
│ ├── jax_qsim/ # Core contraction engine (tensordot + transpose)
│ └── quantum_research/ # VQE, QAOA, GHZ state prep, noise trajectories
├── tpu/ # TPU Scaling Suite (experiments and runners)
├── shors/ # TPU-sharded Shor's Algorithm (33 qubits)
├── grover_simulation/ # Grover's Search (up to 36 qubits on 64 TPU chips)
├── tests/ # Pytest verification suite
└── requirements.txt # Core dependencies
Ensure you have CUDA 12 installed, then set up the environment:
python3 -m venv venv && source venv/bin/activate
pip install --upgrade "jax[cuda12]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
pip install matplotlib pytest numpy
Verify the JAX device setup:
python3 -c "import jax; print('Backend:', jax.default_backend()); print('Devices:', jax.devices())"
Run the local GPU benchmarks:
python benchmarks/benchmark_27q.py
In your TPU VM cluster SSH session:
python3 -m venv tpu_env && source tpu_env/bin/activate
pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
pip install matplotlib numpy
Run the scaling suite:
python tpu/tpu_quantum_scale.py
| Environment | Hardware | Max Qubits | State-Vector Footprint | Gate Speed (10-q) |
|---|---|---|---|---|
| Local GPU | ||||
| NVIDIA RTX 2050 (4 GB VRAM) | 29 | ~4.29 GB | ~0.01 ms | |
| TPU Mesh (v5e-16) | ||||
| 16x TPU v5e (256 GB aggregate HBM2e) | 33 | 64.00 GB | ~0.01 ms | |
| TPU Mesh (v6e-64) | ||||
| 64x TPU v6e (2.0 TB aggregate HBM3) | 36 | 549.76 GB | ~0.01 ms |
We are extremely grateful to the TPU Research Cloud (TRC) program by Google for providing access to Cloud TPU v6e and v5e VM clusters that enabled this scale of research.
Licensed under the Apache License 2.0.