Scaling Quantum State-Vector Simulation to 36 Qubits on Google Cloud TPU: A Pure-JAX Architecture Without Out-of-Memory Failures

Classical simulation of quantum circuits scales exponentially: an n-qubit state vector requires 2ⁿ complex amplitudes, consuming 2ⁿ×8 bytes of memory for single-precision complex numbers. On a single accelerator this wall is typically hit at 29-30 qubits (≈4.3 GB). We describe an architecture built entirely in pure JAX with zero dependency on Qiskit, PennyLane, Cirq, or any classical quantum framework that scales quantum state-vector simulation to 36 qubits (549.76 GB state-vector footprint) across a 64-chip Google Cloud TPU v6e mesh, and to 33 qubits (64 GB) across a 16-chip TPU v5e cluster. The key insight is that JAX's PositionalSharding distributes the 2ⁿ-amplitude state vector uniformly across chip HBM pools as a first-class XLA buffer, while jax.lax.fori_loop keeps the XLA computation graph at O(1) size regardless of circuit depth, preventing the host-side Out-of-Memory (OOM) errors that plague compiler-unrolled deep circuits. We demonstrate five experiments (GHZ state preparation, VQC, VQE H2, QAOA MaxCut, Monte Carlo noise simulation) running at sub-millisecond gate latency after JIT compilation, with full reverse-mode automatic differentiation through the entire circuit.