JAX Quantum Research Suite: A Unified, Hardware-Accelerated, Differentiable Simulator for NISQ-Era Algorithm Research Across GPU and Cloud TPU Clusters

We present the JAX Quantum Research Suite, a high-performance, differentiable quantum circuit simulator spanning two co-existing hardware acceleration layers: a GPU division targeting NVIDIA RTX-class consumer GPUs via CUDA and a Cloud TPU division distributed across Google Cloud TPU v5e-16 and v6e-64 VM clusters. The suite is implemented in pure JAX (jax.numpy, jax.lax) for all statevector primitives, enabling seamless composition with jax.grad, jax.vmap, and jax.sharding.Positional Sharding.