AshiteshSingh / Tpu-Accelerated-Quantum-JAX
High-performance, differentiable quantum state-vector & tensor network simulator in 100% pure JAX (no classical framework overhead). Accelerated on NVIDIA GPUs and Google Cloud TPU v6e-64/v5e VM clusters up to 37 qubits! Supported by Google's TPU Research Cloud (TRC) program.
Why this job: works in the areas we cover, but does not say which job it does
What it does
- 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.
Get started
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 - Maintenance
- Active
- Last commit 55 days ago
- Stars
- 132
- +12 this week
- Licence
- Apache-2.0
- Allows business use
- Downloads for
- None listed
- No release files
- Latest version
- None yet
- No releases
- Built with
- Jupyter Notebook
- 72% of the code
Is it looked after?
- Contributors
- 1
- Last 90 days
- Issues
- None
- No issues filed
- Pull requests
- 6 merged
- 0 waiting
- Releases
- None
- Never published
- No contributing guide
- No security policy
- No code of conduct
Similar repos
| Repo | Stars | Last commit | Licence | Downloads for |
|---|---|---|---|---|
| LuxDL / Lux.jl | 738 | today | MIT | No downloads |
| mit-han-lab / torchquantum | 1,681 | 4 months ago | MIT | No downloads |
| stfc / qiskit-machine-learning | 1,116 | today | Apache-2.0 | No downloads |
More: languages, topics and project record
Languages
- Jupyter Notebook 72%
- Python 27%
- Shell 1%
Topics
- google-cloud
- grovers-algorithm
- jax
- machine-learning
- nvidia-gpu
- python
- quantum-computing
- quantum-computing-research
- quantum-ml
- quantum-simulation
- shors-algorithm
- tpu
Project record
- JobOther AI tools
- StatusListed
- Commits266
- Forks13
- Watching2
- Stars in 28 days+34
- CreatedMay 24, 2026
- First seen by usNov 7, 2026
- CheckedNov 12, 2026