Tpu-Accelerated-Quantum-JAX

Tpu-Accelerated-Quantum-JAX offers a high-performance, differentiable quantum state-vector and tensor network simulator that operates exclusively with JAX, eliminating dependencies on classical frameworks. It is optimized for execution on NVIDIA GPUs and Google Cloud TPU v6e-64/v5e clusters, enabling the simulation of up to 37-qubit systems with remarkable efficiency.

Key Features

• 100% Pure JAX: Built without heavy dependencies such as Qiskit, Cirq, or Pennylane. The project compiles into a single monolithic XLA kernel, allowing for bare-metal execution speeds.
• Multi-Device Sharding: Supports scaling to 36 qubits with a state-vector footprint of 549 GB, utilizing JAX’s PositionalSharding across a 64-chip Cloud TPU v6e mesh.
• Reverse-Mode Auto-Differentiation: Enables precise gradient computations using jax.grad, facilitating fast training for various quantum algorithms such as Variational Quantum Eigensolver (VQE), Quantum Approximate Optimization Algorithm (QAOA), and Quantum Neural Networks (QNNs).
• Hardware-Level Optimizations: Implements structured loop primitives to prevent XLA graph bloat and gradient rematerialization techniques to maintain memory complexity at O(1).
• Stochastic Noise Support: Features built-in Monte Carlo trajectory simulations to manage open system dynamics and depolarizing noise for near-term quantum gates.

Directory Layout

The project is organized as follows:

.
├── gpu/                     # GPU Modular Simulator & Research scripts
│   ├── jax_qsim/            # Core contraction engine (tensordot + transpose)
│   └── quantum_research/    # VQE, QAOA, GHZ state preparation, 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

Performance Summary

Acknowledgements

Gratitude is extended to the TPU Research Cloud (TRC) program by Google for providing access to Cloud TPU v6e and v5e VM clusters, which facilitated this scale of research.

For further insights into the research, benchmarks, and in-depth technical details, visit ashitesh.me.