Show HN: TPU-accelerated quantum circuit simulation in Jax

A developer released a high-performance quantum circuit simulator built entirely in JAX, capable of simulating 36-qubit circuits with a 549 GB state-vector footprint on Google Cloud TPU v6e-64 clusters. The open-source project achieves ~0.01ms per gate and supports reverse-mode auto-differentiation, making it suitable for variational quantum algorithms. It is supported by Google's TPU Research Cloud program.

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 JAX PositionalSharding across a 64-chip Cloud TPU v6e mesh. - Reverse-Mode Auto-Differentiation: Compute exact gradients in a single backward pass via jax.grad for fast training of variational algorithms VQE, QAOA, QNNs . - Hardware-Level Optimizations: Structured loop primitives jax.lax.fori loop prevent XLA graph bloat, while jax.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: python 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 /AshiteshSingh/Tpu-Accelerated-Quantum-JAX/blob/main/LICENSE .