{"slug": "show-hn-tpu-accelerated-quantum-circuit-simulation-in-jax", "title": "Show HN: TPU-accelerated quantum circuit simulation in Jax", "summary": "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.", "body_md": "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.\n\nA 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.\n\n-\n**100% Pure JAX:** Zero dependencies on heavy frameworks (Qiskit, Cirq, Pennylane). Compiled natively into a single monolithic XLA kernel for bare-metal execution speeds. -\n**Multi-Device Sharding:** Scale up to 36 qubits (549 GB state-vector footprint) using distributed JAX`PositionalSharding`\n\nacross a 64-chip Cloud TPU v6e mesh. -\n**Reverse-Mode Auto-Differentiation:** Compute exact gradients in a single backward pass via`jax.grad`\n\nfor fast training of variational algorithms (VQE, QAOA, QNNs). -\n**Hardware-Level Optimizations:** Structured loop primitives (`jax.lax.fori_loop`\n\n) prevent XLA graph bloat, while`jax.checkpoint`\n\n(gradient rematerialization) keeps memory complexity at$\\mathcal{O}(1)$ . -\n**Stochastic Noise Support:** Built-in Monte Carlo trajectory simulations for open systems and depolarizing NISQ gate noise.\n\n```\n.\n├── gpu/                     # GPU Modular Simulator & Research scripts\n│   ├── jax_qsim/            # Core contraction engine (tensordot + transpose)\n│   └── quantum_research/    # VQE, QAOA, GHZ state prep, noise trajectories\n├── tpu/                     # TPU Scaling Suite (experiments and runners)\n├── shors/                   # TPU-sharded Shor's Algorithm (33 qubits)\n├── grover_simulation/       # Grover's Search (up to 36 qubits on 64 TPU chips)\n├── tests/                   # Pytest verification suite\n└── requirements.txt         # Core dependencies\n```\n\nEnsure you have CUDA 12 installed, then set up the environment:\n\n```\npython3 -m venv venv && source venv/bin/activate\npip install --upgrade \"jax[cuda12]\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\npip install matplotlib pytest numpy\n```\n\nVerify the JAX device setup:\n\n``` python\npython3 -c \"import jax; print('Backend:', jax.default_backend()); print('Devices:', jax.devices())\"\n```\n\nRun the local GPU benchmarks:\n\n```\npython benchmarks/benchmark_27q.py\n```\n\nIn your TPU VM cluster SSH session:\n\n```\npython3 -m venv tpu_env && source tpu_env/bin/activate\npip install \"jax[tpu]\" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html\npip install matplotlib numpy\n```\n\nRun the scaling suite:\n\n```\npython tpu/tpu_quantum_scale.py\n```\n\n| Environment | Hardware | Max Qubits | State-Vector Footprint | Gate Speed (10-q) |\n|---|---|---|---|---|\nLocal GPU |\nNVIDIA RTX 2050 (4 GB VRAM) | 29 | ~4.29 GB | ~0.01 ms |\nTPU Mesh (v5e-16) |\n16x TPU v5e (256 GB aggregate HBM2e) | 33 | 64.00 GB | ~0.01 ms |\nTPU Mesh (v6e-64) |\n64x TPU v6e (2.0 TB aggregate HBM3) | 36 | 549.76 GB | ~0.01 ms |\n\nWe 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.\n\nLicensed under the [Apache License 2.0](/AshiteshSingh/Tpu-Accelerated-Quantum-JAX/blob/main/LICENSE).", "url": "https://wpnews.pro/news/show-hn-tpu-accelerated-quantum-circuit-simulation-in-jax", "canonical_source": "https://github.com/AshiteshSingh/Tpu-Accelerated-Quantum-JAX", "published_at": "2026-07-18 14:43:18+00:00", "updated_at": "2026-07-18 15:21:32.458983+00:00", "lang": "en", "topics": ["machine-learning", "ai-research", "ai-infrastructure", "ai-tools", "developer-tools"], "entities": ["Google", "NVIDIA", "JAX", "TPU Research Cloud", "Cloud TPU v6e", "Cloud TPU v5e", "AshiteshSingh"], "alternates": {"html": "https://wpnews.pro/news/show-hn-tpu-accelerated-quantum-circuit-simulation-in-jax", "markdown": "https://wpnews.pro/news/show-hn-tpu-accelerated-quantum-circuit-simulation-in-jax.md", "text": "https://wpnews.pro/news/show-hn-tpu-accelerated-quantum-circuit-simulation-in-jax.txt", "jsonld": "https://wpnews.pro/news/show-hn-tpu-accelerated-quantum-circuit-simulation-in-jax.jsonld"}}