Author: Ashitesh Kumar Singh (Class 12th Student & Independent Quantum Computing & Cryptography Researcher • GitHub)
Category: Mobile Security / Reverse Engineering / Cloud Infrastructure / Responsible Disclosure
Status: Remediated & Officially Acknowledged by CERT-In

📌 Executive Summary

In July 2026, while in Class 12th, I was reverse engineering and analyzing the client-side architecture of a widely adopted Indian school enterprise resource planning (ERP) platform used by my own school. During this research, I uncovered an escalating vulnerability chain.

What began as an inspection of local caching on my personal student account quickly escalated into direct access to central cloud infrastructure, production databases, and master administrative accounts. The vulnerability exposed an estimated 5 million+ citizen entities (a conservative baseline calculation factoring in all enrolled students, parents, and faculty across thousands of schools nationwide, with the actual real-world total likely being even higher), spanning over 12 TB+ of sensitive records—including:

• Detailed personally identifiable information (PII) of minor students

• Parent contact data & financial records

• Faculty dossiers & master administrative credentials

• Enterprise employee payroll and salary registers

• Defense-affiliated school registries

If an adversary had exploited this architecture, the private data of millions of students, parents, and teachers would have leaked across the country. Furthermore, the elevated access technically granted full write and update privileges—meaning I had the capability to manipulate exam marks, alter academic records, falsify attendance, and modify fee payment ledgers.

Rather than tampering with any records or exploiting this access, I immediately halted all database operations, prepared a reproducible Proof of Concept (PoC) for validation, and coordinated with CERT-In (Indian Computer Emergency Response Team). Following formal verification and vendor remediation, CERT-In issued an official acknowledgment recognizing the responsible disclosure.

🔍 The Reality Check: An 8 GB Slice of My Own School

To verify the scope of the vulnerability safely without intruding into unauthorized organizations, I restricted testing strictly to the tenant slice of my own enrolled school.

As a Class 12th student analyzing an application I used every day, this was no longer an abstract dataset or synthetic test entries—these were real-world production records belonging to my own classmates, friends, and teachers:

• Complete Academic Profiles: Unrestricted access to classmates' full academic histories, raw exam marksheets, internal performance remarks, attendance logs, and disciplinary notes.

• Arbitrary Record Manipulation (Untouched): Elevated access carried write privileges that made it technically possible to arbitrarily alter exam scores, modify report cards, clear pending tuition fee dues, and tamper with student profiles. Adhering to strict ethical principles, not a single record was ever modified, falsified, or tampered with.

• Intimate Identity Records: Direct visibility into student profile photos, residential home addresses, and linked national identity documents (including student and parent Aadhaar records and PAN details).

• Financial & Banking Records: Parent bank account numbers, fee payment transactions, tuition billing histories, and saved card transaction tokens.

• Device Telemetry & Expired Token Replay: Hardware identifiers, device models, OS versions, and push notification tokens. Critically, backend verification failed to enforce proper invalidation—allowing state queries and data retrieval through expired tokens and orphaned device sessions.

⚠️ Scope & Scale: Querying the partition for just my own school accounted for approximately 8 GB of sensitive data. Extrapolating this verified footprint across the thousands of educational institutions served by the platform nationwide—including several schools catering to families of Indian defense personnel—established an aggregate central data footprint exceeding 12 TB+ and upwards of 5 million+ citizen records (calculated as a conservative lower bound across students, parents, and teachers, with the actual total likely to be even higher). If breached by a malicious threat actor, the personal and financial data of millions of students, parents, and educators would have been catastrophically leaked.

🌐 The Full Scale of Exposure

Beyond student and institutional records, deeper infrastructure routes revealed unrestricted access to central administrative systems:

[Student Mobile App]
        │
        ▼ (Static APK Analysis / JADX)
[Insecure SQLite Persistence (CWE-312)]
        │
        ▼ (Plaintext Operational AWS Keys Extracted)
[Internal 12-Hour Rotating Service APIs]
        │
        ▼ (Unrestricted IAM Escalation)
[Central Production Clusters (12 TB+ / 5M+ Records)]
   ├── Student & Parent PII (Aadhaar, Marks, Bank Data)
   ├── Defense School Registries
   ├── Cleartext Master Admin Credentials
   └── Corporate HR, Payroll & Employee Salaries

• Master Admin Portals: Cleartext administrative passwords that allowed direct logins into multi-tenant master consoles across thousands of schools.

• Corporate & Employee Data: Internal company payroll databases, staff salaries, HR files, and vendor configuration secrets.

• Defense School Assets: Confidential student and parent rosters for educational institutions operating under defense family welfare boards.

⚙️ Technical Breakdown

1. Static Reverse Engineering & APK Logic Inspection

Applying static reverse engineering to the vendor's production Android APK (using JADX-GUI and APKTool) revealed that permanent root AWS keys were not hardcoded in the string pool. However, decompiling and reversing the client session management logic exposed a critical architectural flaw: upon successful authentication, operational credentials were written directly into local app cache files.

2. Dynamic Execution & Insecure Local Persistence (CWE-312)

Accessing an Android application's sandboxed internal database directory (/data/data/) normally requires root permissions on a device. Instead of rooting a physical smartphone, I utilized LDPlayer (with root mode enabled) as a controlled testing environment.

Authenticating through my legitimate personal student account inside the emulator, I inspected the app's internal database directory via ADB:

adb shell
su
cd /data/data/<package_name>/databases/
sqlite3 local_cache.db

Querying the SQLite database confirmed the critical failure: The app stored active operational AWS Access Key IDs and Secret Access Keys in plaintext on client-side storage.

3. Traversing the 12-Hour Rotating API

The recovered credentials carried overly permissive IAM policies across the vendor's cloud environment:

1. The keys granted direct interaction with an internal service routing gateway.

2. Within this subsystem, a secondary service API rotated access tokens every 12 hours.

3. Inspecting routing tables and endpoint responses exposed direct connection paths leading straight into the vendor's central production database clusters.

4. Verification and Ethical Cessation

Once the 8 GB institutional slice and the existence of multi-terabyte production tables were confirmed, all probing was terminated immediately:

• 🚫 No data dumping: No tables were dumped, mirrored, or exfiltrated.

• 🚫 Zero data tampering: Marks, fees, attendance, and student profiles were never modified or altered.

• 🚫 Zero retention: No personal records of classmates, parents, or staff were retained.

• 🧹 Complete sanitization: All temporary debugging databases, tokens, and test logs were permanently wiped from the testing machine.

⏱️ Disclosure Timeline

🛡️ Verification Proofs & CERT-In Correspondence

1. PoC Request with Active Timestamp (July 09, 2026)

Following initial reporting, CERT-In's Incident Response team verified receipt and requested an active-timestamp validation proof:

2. Official Incident Registration (July 10, 2026)

Upon receiving the step-by-step PoC, CERT-In formally registered the incident and assigned a case reference ID:

3. Formal Acknowledgment of Responsible Disclosure (September 24, 2026)

Following vendor remediation and verification, CERT-In issued an official acknowledgment letter:

🏆 The Outcome: Safeguarding Millions & Enforcing DPDP Compliance

• Zero Records Breached, Everything Fixed: Because this disclosure was conducted in strict good faith, not a single record was leaked, breached, or sold. All vulnerable endpoints, credential persistence flaws, and permissive IAM roles were thoroughly remediated and patched.

• Protecting 5 Million+ Students, Parents & Teachers: As a Class 12th student, I am glad and proud to have helped secure the sensitive personal data of approximately 5 million+ students, parents, and teachers (calculated as a conservative baseline across thousands of schools, though the true figure is likely even higher) before any malicious actors could discover and weaponize this flaw.

• Mandated DPDP Act (Digital Personal Data Protection) Consent: Because this massive dataset involved the PII of minor students, the vulnerability exposed significant compliance vulnerabilities. Following this disclosure and CERT-In's intervention, the vendor was required to enforce explicit DPDP Act consent workflows and parental consent verification directly into the platform's login and onboarding architecture.

💡 Key Takeaways for EdTech & Mobile Security

1. Zero Trust on Client Storage
   Mobile devices must never store cloud infrastructure keys or IAM credentials. All database interactions must terminate at an authenticated backend API implementing strict role-based access control (RBAC).

2. Proper Token Lifecycle Management
   Session tokens and authentication artifacts must be validated server-side on every request, with strict expiration enforcement to prevent token replay and orphaned sessions.

3. Strict Cloud IAM Scoping (Principle of Least Privilege)
   Cloud identities must adhere to the Principle of Least Privilege (PoLP). A client-facing configuration must never have network paths or IAM permissions to query master production clusters, corporate HR payroll, or cross-tenant datasets.

4. Compliance by Design (DPDP Act & Minor PII)
   Platforms managing students' and minors' records must treat data protection as a legal baseline. Integrating transparent consent, parental verification, and automated audit logs is essential for compliance under modern regulations like India's DPDP Act.

Responsible Disclosure Notice:
This vulnerability was researched and disclosed entirely in good faith in compliance with responsible security disclosure standards. All vendor names, sensitive endpoints, school identifiers, and exploit scripts have been omitted or redacted.

Simulating quantum circuits on a classical computer is exponentially hard. An n-qubit state vector holds 2ⁿ complex amplitudes — at 30 qubits that's 8 GB, at 33 qubits 64 GB, at 36 qubits 549 GB. No single GPU or TPU chip can hold that.

Most frameworks (Qiskit, PennyLane) give up around 29–30 qubits for two reasons:

1. Memory wall — the full state vector doesn't fit on one device
2. XLA graph explosion — deep circuits compiled with Python for-loops unroll into millions of XLA graph nodes, crashing the compiler host before a single gate runs

I built a simulator in 100% pure JAX that solves both — scaling to 36 qubits (549 GB) across a 64-chip Google Cloud TPU v6e cluster, with full reverse-mode autodiff and sub-millisecond gate latency.

Solution 1: Distributed State Vector with PositionalSharding

Instead of storing all 2ⁿ amplitudes on one chip, I split the state vector across the entire TPU mesh using JAX's native sharding API:

mesh          = Mesh(mesh_utils.create_device_mesh((16,)), axis_names=("chips",))
state_sharding = NamedSharding(mesh, P("chips", None))

# 33-qubit state: shape [16, 2^29] — each chip holds 2^29 amplitudes
state = jax.jit(init_ground_state, out_shardings=state_sharding)()

Each chip holds exactly its slice of the state vector in its own HBM pool. JAX's Inter-Chip Interconnects (ICI) handle cross-chip gate operations automatically. The host CPU never holds the full array.

Solution 2: lax.fori_loop for O(1) Compiler Graph Size

A Python for-loop inside @jax.jit statically unrolls every iteration into a separate XLA graph node. At 100 circuit layers × 30 qubits, the HLO graph has millions of nodes — the XLA compiler host OOMs during compilation, before execution even starts.

jax.lax.fori_loop compiles the loop body once as a single While HLO node:

# ✅ O(1) XLA graph — compiles once, runs D times at hardware speed
state = lax.fori_loop(0, D, circuit_body, state)

# ❌ O(D) XLA graph — explodes at 100+ layers
for i in range(D):
    state = apply_gate(state, ...)

Graph size stays constant regardless of circuit depth. No compiler OOM.

The Full Compilation Stack

Python (JAX / jnp)
      ↓  @jax.jit traces once
XLA High-Level Operations (HLO)
      ↓  XLA compiler: fusion, tiling, layout
TPU HBM3 machine instructions
      ↓  ICI for cross-chip ops
Result: 36-qubit state vector across 64 HBM3 pools

The entire simulation — state init, gate application, observables, and gradient accumulation — runs inside one monolithic XLA kernel. No Python re-entry between gates. No CPU↔device copies mid-circuit.

Key Results

• Gate latency (after JIT warm-up): ~0.01ms
• Gradient of 100-parameter circuit: ~8ms (one jax.grad backward pass vs ~10s with PennyLane parameter-shift)
• VQE H₂ ground state: converged to −1.13627 Hartree — within 1.6×10⁻³ Hartree of the FCI reference (quantum chemical accuracy)
• Grover's at 36 qubits: marked state probability peaks at ~0.998 after 205,887 iterations
• Barren plateau: gradient variance confirmed to decay as 2⁻ⁿ up to 24 qubits (consistent with McClean et al. 2018)

Try It for Free

The repo includes a Colab notebook that runs on a free Colab TPU v5e-1 — no Google Cloud account needed.

Six self-contained experiments: GHZ state prep, VQC XOR classifier, VQE H₂, QAOA MaxCut, Monte Carlo noise simulation, and noisy NISQ fidelity decay.

Why Pure JAX — Not PennyLane, Qiskit, or PyTorch?

JAX is XLA. There is no bridge. Python code is the source language for a JIT compiler that targets TPU silicon directly.

References

1. Bradbury et al. (2018). JAX: composable transformations of Python+NumPy programs.
2. McClean et al. (2018). Barren plateaus in quantum neural network training landscapes. Nature Communications.
3. Peruzzo et al. (2014). A variational eigenvalue solver on a photonic quantum processor. Nature Communications.
4. Farhi et al. (2014). A quantum approximate optimization algorithm. arXiv:1411.7308.
5. Vidal (2003). Efficient classical simulation of slightly entangled quantum computations. PRL.

Supported by the Google TPU Research Cloud (TRC) program.
Apache 2.0 License · github.com/AshiteshSingh/Tpu-Accelerated-Quantum-JAX

Classical simulation of quantum systems is the bedrock of modern quantum algorithm research. However, simulating variational quantum algorithms (VQAs) has historically run into a massive wall: gradient computation scaling.

In this post, I’ll introduce the JAX Quantum Research Suite—a hardware-accelerated, fully differentiable quantum simulator designed to run seamlessly on consumer GPUs and scale up to Google Cloud TPU clusters.

We’ll explore how it overcomes traditional limitations, achieves up to a 48.7× speedup in gradient calculations, and scales simulations up to 36 qubits on TPU v6e and 1,024 qubits using Matrix Product States (MPS).

This research and scaling development was supported by the Google TPU Research Cloud (TRC) program.

The Core Challenge: Differentiating Quantum Circuits

For variational algorithms like VQE (Variational Quantum Eigensolver) or QAOA (Quantum Approximate Optimization Algorithm), we need to optimize parameterized quantum circuits.

In standard frameworks, computing gradients is done using the Parameter-Shift Rule (PSR). For a circuit with \(P\) parameters, PSR requires \(2P\) full circuit evaluations. When you scale up to hundreds of parameters, this linear \(O(P)\) overhead becomes a major bottleneck.

By using JAX, our simulator implements reverse-mode automatic differentiation (backpropagation). This allows us to compute all \(P\) gradients in a single backward pass—an \(O(1)\) scaling cost relative to the parameter count.

🚀 Key Highlights & Benchmarks

1. 48.7× Faster Gradients (Reverse-Mode AD vs. Parameter-Shift)

On a 15-qubit Hardware-Efficient Ansatz circuit with 120 parameters, we benchmarked gradient calculation times:

• Parameter-Shift Rule (PSR): 1,826 ms
• JAX Reverse-Mode Autodiff (jax.grad): 37.5 ms
• Performance Leap: A 48.7× speedup, which grows even larger as the number of parameters increases.
• Vs. PennyLane JAX Backend: Our custom statevector simulator ran ~4× faster (2ms vs 8ms for 50 params on GPU) due to monolithic XLA kernel fusion.

2. Grover's Algorithm at 36 Qubits (549 GB Statevector)

Full statevector simulation scales exponentially in memory. A 36-qubit simulation requires 549.76 GB of memory to store the amplitudes. Using a 64-chip Google Cloud TPU v6e mesh, we successfully ran Grover's algorithm at 36 qubits, distributing the massive statevector across TPU HBM3.

3. Shor's Algorithm at 33 Qubits on TPU v5e-16

We demonstrated a distributed QFT (Quantum Fourier Transform) pipeline for Shor's order-finding at a 33-qubit scale (64 GB statevector) on a 16-chip Cloud TPU v5e mesh, using network-optimized collective communication (ppermute) to reduce inter-chip transfer spikes from 8 GB to 128 MB.

4. 1,024-Qubit MPS VQE & Numerical Stability Breakthroughs

To simulate beyond statevector limits (36+ qubits), we built a differentiable Matrix Product State (MPS) simulator in pure JAX. When differentiating through Singular Value Decomposition (SVD), JAX frequently crashes due to Wirtinger gradient singularities. We introduced three key stabilization techniques:

1. SVD Epsilon Floor (1e-7) to prevent division-by-zero.
2. Site-level Normalization to stop amplitude drift.
3. SVD Jitter & Momentum SGD to low-pass filter and damp "V-bounce" oscillations caused by bond-dimension entanglement ceilings.

🛠️ Three Engineering Tricks for Cloud TPU Scaling

Tracing deep quantum circuits into XLA bytecode can crash compile hosts or run out of memory. We solved these issues with three specific JAX primitives:

A. Multi-Device PositionalSharding

We partition the statevector's leading dimension across physical TPU chips:

sharding = PositionalSharding(jax.devices()).reshape(NUM_DEV, 1)
state = jax.device_put(state, sharding)

XLA handles local gate execution automatically, routing cross-shard gates via TPU Inter-Chip Interconnects (ICI).

B. lax.fori_loop for O(1) Compiler Graph Size

Unrolled Python loops create millions of HLO nodes, causing compile-time memory leaks. We wrap our circuit execution using JAX's loop control:

state_new = jax.lax.fori_loop(0, depth, body_fn, state_init)

This forces XLA to compile the gate block once, keeping the compiler graph size constant (\(O(1)\)) regardless of circuit depth.

C. jax.checkpoint for O(1) Backpropagation Memory

Reverse-mode AD normally stores all intermediate states in memory. For deep circuits, this causes HBM OOM errors. We apply gradient checkpointing to recompute intermediate states on-the-fly during the backward pass:

@jax.checkpoint
def circuit_layer(state, params):
    return apply_gates(state, params)

🎯 Summary of Experiments

• GHZ State Prep: Learned to prepare a 3-qubit GHZ state with \(\mathcal{F} > 0.9999\) in 200 epochs.
• XOR VQC Classifier: Batched classification using jax.vmap achieving 97%+ accuracy.
• VQE (H₂ Molecule): Solved ground state energy to within 0.1 mHa of the FCI reference (chemical accuracy).
• QAOA MaxCut: Achieved a 99.7% approximation ratio at depth \(p=5\) on a 6-node weighted graph.
• Monte Carlo Noise Trajectories: Simulated noisy NISQ channels (amplitude/phase damping) showing exact statistical convergence matching analytical curves.
• Barren Plateaus: Confirmed McClean's exponential gradient variance decay (\(\mathcal{O}(2^{-n})\)) up to 24 qubits on TPU.

🙏 Acknowledgements

This research was made possible through the Google TPU Research Cloud (TRC) program, which provided the TPU v5e and v6e compute resources necessary to scale our simulations to 33 and 36 qubits.

• Codebase & Experiments: GitHub Repository
• Author: Ashitesh Singh

For the 37-qubit random circuit sampling (RCS) benchmark, we leverage TensorCircuit with a JAX backend — the only experiment in the suite using an external quantum framework — enabling tensor-network-based amplitude sampling at extreme scale on TPU v6e-64. All other experiments use our custom pure-JAX statevector engine with zero external quantum framework dependencies.

On a consumer NVIDIA GeForce RTX 2050 (4 GB VRAM) — selected specifically to test the framework's baseline performance on an entry-level consumer-grade GPU — our JAX statevector simulator establishes a baseline for single-GPU execution. Gradient computation scales favorably: on a 120-parameter CPU circuit (N=10 rigorous runs, 9 stable runs post-JIT warm-up), jax.grad computes all gradients in a single reverse-mode backward pass at 37.5ms vs. 1,826ms for PennyLane's parameter-shift rule — a 48.7× improvement (Section 5.3a). On a smaller 50-parameter GPU circuit (N=3 preliminary), the ratio is ~75×, consistent with PSR overhead scaling linearly with parameter count. Against PennyLane's own JAX reverse-mode backend, the advantage is ~4× (2ms vs ~8ms). On the Cloud TPU v5e-16 mesh (256 GB aggregate HBM2e), full state-vector simulation scales to 33 qubits (64 GB), and on the TPU v6e-64 cluster, Grover's algorithm is evaluated at 36 qubits (549 GB). The 37-qubit RCS result (Section 5.7) uses tensor-network amplitude sampling via TensorCircuit and yields F_XEB ≈ 0 (preliminary, N=5 runs) — indicating the sampled distribution is near-uniform, a null result expected for deep chaotic circuits evaluated under finite bond-dimension approximation.

We demonstrate the suite through nine experiment subsections (Sections 4.1–4.9): GHZ state preparation, Variational Quantum Classification (XOR), Variational Quantum Eigensolver (H₂), QAOA MaxCut, Monte Carlo quantum noise trajectories, noisy NISQ fidelity decay, barren plateau gradient variance scaling, Shor's 33-qubit order-finding demonstration, and Grover's 36-qubit full statevector simulation — plus MPS VQE at 512–1024 qubits with novel SVD numerical stability analysis.

Keywords: Quantum simulation, JAX, XLA, automatic differentiation, TPU, variational quantum eigensolver, QAOA, Shor's algorithm, Grover's algorithm, matrix product states, barren plateaus, NISQ

1. Introduction

Classical simulation of quantum systems remains the primary tool for algorithm development and hardware validation in the Noisy Intermediate-Scale Quantum (NISQ) era [1]. While the broader machine learning community heavily utilizes GPUs, Tensor Processing Units (TPUs) are fundamentally massive matrix-multiplication engines connected by ultra-high-speed Inter-Chip Interconnects (ICI). This architecture makes TPUs exceptionally well-suited for quantum simulation: statevector updates and tensor-network contractions are overwhelmingly dominated by large-scale matrix multiplications. TPUs provide seamless out-of-the-box distributed memory scaling via JAX, generally offering faster and more stable execution for these specific distributed workloads without the need for researchers to write complex C++ or CUDA code. While dedicated C++/CUDA simulators — qsim, Qiskit-Aer with cuQuantum — achieve high raw gate-application throughput, they are architecturally incompatible with the core requirement of modern variational quantum algorithm research: end-to-end, hardware-accelerated automatic differentiation.

The standard hardware-compatible gradient method in tools like PennyLane is the Parameter-Shift Rule (PSR) [2]. For a circuit with \(P\) parameters, PSR requires \(2P\) circuit evaluations to obtain all gradients, creating an \(O(P)\) bottleneck prohibitive for large variational circuits. Reverse-mode automatic differentiation (backpropagation) computes all \(P\) gradients in a single backward pass — O(1) evaluations — but requires a simulator that is transparently differentiable at the mathematical level.

The JAX ecosystem [3] uniquely resolves this tension. Because JAX traces Python code directly into XLA High-Level Operations (HLO) bytecode, a quantum circuit expressed as jnp.tensordot and jnp.transpose calls can be:

1. JIT-compiled (@jax.jit) into a monolithic XLA kernel — zero Python overhead after first compilation

2. Automatically differentiated (jax.grad) in a single reverse-mode backward pass — O(1) evaluations for all parameters

3. Vectorized (jax.vmap) over batches of parameters or states without code changes

4. Distributed (jax.sharding.PositionalSharding) across multi-chip TPU meshes via the same API

No other framework achieves this seamless compilation path to TPU hardware without additional tooling: PennyLane stops at Python-level device dispatch; Qiskit-Aer is a CUDA C++ library incompatible with TPU HBM; TensorFlow Quantum's Cirq backend runs circuit simulation on CPU (only classical post-processing hits TPU cores); PyTorch's torch_xla adds a Python-to-XLA bridge at every tensor boundary. JAX traces Python code directly into XLA HLO — quantum gate operations are first-class XLA nodes with no additional bridging layer. Note: Google's own qsim simulator (C++/CUDA) achieves higher raw gate throughput than our JAX engine at small qubit counts, but is not differentiable and has no TPU backend; cuQuantum (NVIDIA) offers gate fusion and multi-GPU state-vector simulation but is similarly non-differentiable and GPU-only.

1.1 Contributions

This paper makes the following contributions:

1. A complete, modular pure-JAX statevector simulator (jax_qsim) with 20+ gate types, density matrix mode, and Kraus noise channels — all composable with jax.grad

2. O(1) XLA graph size for deep circuits via jax.lax.fori_loop — enabling 33-qubit Shor's algorithm on TPU without compiler OOM

3. O(1) backpropagation memory via jax.checkpoint rematerialization for multi-layer variational training at large scale

4. Novel MPS numerical analysis: documentation and resolution of SVD Wirtinger gradient singularities, \(\log_2(\chi)\) entanglement ceiling effects, and V-bounce oscillation in the bond-dimension-limited regime

5. Comprehensive cross-hardware benchmarks with raw data from consumer GPU through TPU cluster

1.2 Related Work

qujax (Duffield et al., JOSS 2023) [4] independently derived the same (2,)*N statetensor functional design. However, the primary novelty and central claim of our work lies in the native multi-chip Cloud TPU distribution layer (scaling to 64 chips seamlessly via JAX's PositionalSharding) enabling extreme-scale memory scaling, which qujax lacks. Furthermore, researchers can use our distributed code to simulate massive systems immediately without having to write any complex C++/CUDA code, as the distributed systems architecture on TPU is already solved for them. Our work also extends the ecosystem with a full circuit-builder abstraction and density matrix mode.

TensorCircuit (Zhang et al., Quantum 2023) [5] uses jnp.einsum tensor contractions as its core primitive and supports JAX, TF, and PyTorch backends. We use TensorCircuit specifically for the 37-qubit RCS benchmark where tensor-network contraction is more efficient than full statevector for random circuit amplitude sampling.

qsim (Google, 2020) is a high-performance C++/CUDA statevector simulator using gate fusion and AVX/GPU parallelism. It achieves higher raw gate throughput than our JAX engine at small qubit counts (particularly on NVIDIA GPUs with cuQuantum integration). Unlike our work, qsim is not differentiable and has no TPU backend — it cannot compute jax.grad through a circuit.

cuQuantum (NVIDIA, 2022) [15] provides CUDA-accelerated statevector and tensor-network simulation with multi-GPU state-vector support and a gate-fusion pass. It outperforms our simulator in raw single-gate throughput but is GPU-only, non-differentiable, and cannot run on TPU hardware.

PennyLane Catalyst (Xanaду, JOSS 2024) [6] uses MLIR/LLVM compilation (not XLA) to produce native binaries. Unlike Catalyst, our approach does not require a separate compilation toolchain beyond JAX.

Quimb [7] is a widely-used Python library for tensor network simulation (MPS, DMRG) that could serve as an alternative backend for large-qubit simulations. Our MPS engine differs by being fully JAX-differentiable, enabling jax.grad through SVD operations.

PennyLane (Bergholm et al., 2018) [8] is the primary framework we compare against for gradient benchmarks. Unlike PennyLane's default parameter-shift rule, our simulator uses JAX reverse-mode autodiff for O(1) gradient computation.

2. Architecture

2.1 Dual Hardware Division Design

JAX Quantum Research Suite
├── GPU Division  (gpu/)
│   ├── jax_qsim/
│   │   ├── core.py         ← Tensor contraction gate engine (tensordot + transpose)
│   │   ├── circuit.py      ← High-level Circuit builder with run() + compile()
│   │   ├── ops.py          ← Gate library (20+ gate types, parameterized)
│   │   ├── observables.py  ← PauliString, Hamiltonian expectations (used in VQE/QAOA)
│   │   └── noise.py        ← Kraus channel density matrix simulator (used in Noise trajectories)
│   └── quantum_research/   ← 8 GPU research scripts (VQE, QAOA, barren plateau...)
│
├── tpu/
│   └── tpu_quantum_scale.py  ← Self-contained 5-experiment suite (GHZ, VQC, VQE, QAOA, Noise)
│
├── shors/
│   └── shors_algorithm_33q.py  ← 33-qubit distributed Shor's on TPU v5e-16
│
└── grover_simulation/       ← 20/30/36-qubit Grover statevector on TPU v6e-64
    ├── 20qubits.py
    ├── 30qubits.py
    └── 36qubits.py

2.2 GPU Division: Statevector Gate Engine

An \(n\)-qubit state is represented as a tensor of shape \((2,)^n\) and all gate applications compile to pure XLA tensor operations.

Single-qubit gate application — 3D transpose method:

For gate \(U \in \mathbb{C}^{2\times 2}\) acting on qubit \(q\) of a flat \(2^n\)-amplitude state vector:

$$|\psi'\rangle = \text{reshape}\bigl(\mathbf{U} \cdot \text{reshape}(\text{transpose}(\text{reshape}(|\psi\rangle, [2^q, 2, 2^{n-q-1}]), [1,0,2]), [2,-1]), [2, 2^q, 2^{n-q-1}]\bigr)$$

This access pattern maps to a single XLA Dot instruction with contiguous strides, producing better memory coalescing than index-arithmetic approaches on GPU HBM.

Two-qubit gate application — 5D transpose method:

For gate \(U \in \mathbb{C}^{4\times 4}\) on qubits \((q_1, q_2)\), \(q_1 < q_2\):

$$\text{state} \xrightarrow{\text{reshape}} (2^{q_1}, 2, 2^{q_2-q_1-1}, 2, 2^{n-q_2-1}) \xrightarrow{\text{transpose}(1,3,0,2,4)} \xrightarrow{\text{apply } U} \xrightarrow{\text{transpose back}} |\psi'\rangle$$

Zhang et al. (2023) show this approach generates better GPU memory coalescing than tensordot for 2-qubit gates because the contracted dimensions (size 2) are contiguous after the initial transpose.

JIT-compiled circuit execution via static circuit structure:

@functools.partial(jax.jit, static_argnums=(2, 3, 4))
def _run_circuit_functional(params, initial_state, num_qubits, ops, state_type):
    for op_name, qubits, p_val in ops:    # fully unrolled at trace time
        ...
    return state

Marking ops as static causes JAX to fully unroll and fuse the gate sequence into a flat XLA graph. Compilation happens once; subsequent calls execute the cached XLA binary with ~0ms overhead.

2.3 TPU Division: Three Engineering Contributions

Contribution A — Multi-Device PositionalSharding:

A 33-qubit state-vector (\(2^{33}\) complex64 amplitudes = 64.00 GB) exceeds any single device. Across a 16-chip TPU v5e mesh:

sharding = PositionalSharding(jax.devices()).reshape(NUM_DEV, 1)
state = jax.device_put(state, sharding)

JAX partitions the leading state-vector dimension across physical chips. Intra-chip gate operations execute without communication; cross-shard gates use XLA collective ops over TPU Inter-Chip Interconnects (ICI).

Contribution B — lax.fori_loop for O(1) Compiler Graph Size:

Python for loops inside @jax.jit unroll fully, making the XLA DAG size \(O(\text{depth})\). A 100-layer circuit produces millions of HLO nodes, causing the compiler host CPU to OOM before any computation begins. We use:

$$\text{state}_\text{new} = \texttt{jax.lax.fori_loop}(0, d, \texttt{body_fn}, \text{state}_0)$$

The XLA compiler compiles body_fn once, representing the loop as a single hardware instruction block. Graph size is \(O(1)\) for any circuit depth \(d\).

Contribution C — jax.checkpoint for O(1) Backpropagation Memory:

Reverse-mode differentiation through a \(d\)-layer circuit naively stores all \(d\) intermediate states in HBM. We wrap layers with jax.checkpoint:

@jax.checkpoint
def circuit_layer(state, params):
    return apply_gates(state, params)

Intermediate states are discarded during the forward pass and recomputed on-the-fly during the backward pass, reducing memory from \(O(d)\) to \(O(1)\) at the cost of one additional forward pass.

2.4 37-Qubit RCS: TensorCircuit + JAX pmap

The 37-qubit random circuit sampling experiment uses TensorCircuit (with tc.set_backend("jax")) for amplitude computation rather than full statevector simulation. TensorCircuit's tensor-network contractor evaluates individual circuit amplitudes via circuit.amplitude(bitstring), which is then vmapped over batches of bitstrings and pmap'd across TPU chips:

tc.set_backend("jax")
tc.set_contractor("auto")

single_chip_batcher = jax.vmap(get_amplitude_probability, in_axes=(None, 0))
parallel_tpu_driver = jax.pmap(single_chip_batcher, in_axes=(None, 0))

This is the only experiment in the suite using an external quantum framework. For all statevector experiments (Shor's, Grover's, VQE, QAOA, GHZ, noise), we use exclusively our pure-JAX engine.

2.5 MPS Tensor Network Engine

For systems exceeding statevector memory limits, the suite implements a differentiable MPS simulator in pure JAX (no Quimb, no tensornetwork library). Each \(n\)-qubit state:

$$|\psi_\text{MPS}\rangle = \sum_{i_1,\dots,i_n} A^{i_1}{[1]} A^{i_2}{[2]} \cdots A^{i_n}{[n]} |i_1 \cdots i_n\rangle, \quad A^{i_k}{[k]} \in \mathbb{C}^{\chi \times \chi}$$

Two-qubit gates are applied via local contraction and SVD truncation:

fused = jnp.einsum("ijk,klm->ijlm", site1, site2)
transformed = jnp.einsum("abcd,ibcj->iadj", gate_u, fused)
mat = mat + 1e-9 * noise   # SVD jitter (see Section 4.9)
u, s, vh = jnp.linalg.svd(mat, full_matrices=False)
s_trunc = s[:CHI] / (jnp.linalg.norm(s[:CHI]) + EPS)

3. Experimental Setup

3.1 Hardware Platforms

3.2 Software Stack

• JAX 0.4.x with CUDA 12 backend (GPU) / libtpu (TPU)

• Python 3.10+, NumPy, Matplotlib

• TensorCircuit: used exclusively in the 37-qubit RCS experiment

• PennyLane default.qubit (JAX interface) and Cirq Simulator: used in GPU benchmarks for comparison

3.3 Statistical Methodology

⚠ Important for reproducibility: Benchmarks involving GPU/TPU execution measure wall-clock time after explicit .block_until_ready() synchronization. All tables report mean ± standard deviation over the stated number of timed runs, excluding the first compilation/warmup call. For time-critical results (gate speed, gradient step), we recommend reproducers run ≥10 timed iterations. Current benchmark files record 3 timed execution runs; tables note this where applicable. We flag data points where N < 10 as preliminary.

3.4 Standard Benchmark Circuit

Unless otherwise specified, benchmarks use a Hardware-Efficient Ansatz (HEA):

• \(L = 3\) layers of \([RY, RZ]\) on each qubit + nearest-neighbor CNOT chain

• Final \([RY, RZ]\) layer

• Total parameters: \(n \times 2 \times (L+1)\) for \(n\) qubits

4. Experiments and Results

4.1 GHZ State Preparation (Entanglement Learning)

Objective: Optimize a 9-parameter, 3-layer ansatz \(U(\vec{\theta})\) to prepare the 3-qubit GHZ state:

$$|\text{GHZ}\rangle = \frac{|000\rangle + |111\rangle}{\sqrt{2}}$$

Loss function: $$\mathcal{L}(\vec{\theta}) = 1 - \left|\langle\text{GHZ}|U(\vec{\theta})|000\rangle\right|^2$$

Optimizer: Adam (\(\text{lr} = 0.05\), \(\beta_1 = 0.9\), \(\beta_2 = 0.999\)), 200 epochs.

Results: Both GPU and TPU divisions converge to fidelity \(\mathcal{F} > 0.9999\) within 200 epochs with smooth monotonic convergence. Wall-clock time per gradient step: ~0.4ms (GPU, post-JIT warmup).

Figure 1: GHZ State Preparation — Fidelity and loss convergence over 200 Adam epochs on GPU (RTX 2050). Orange: infidelity (loss, 1 − F); Green: state fidelity F. Both GPU and TPU executions achieve identical final fidelity.

4.2 Variational Quantum Classifier (XOR Boundary Learning)

Objective: Resolve the XOR classification boundary using a 2-qubit, 8-parameter VQC with angle-encoded inputs:

$$P(y_i = 1|\vec{x}_i) = \frac{1 + \langle\psi(\vec{x}_i)|V^\dagger(\vec{\theta}) Z_0 V(\vec{\theta})|\psi(\vec{x}_i)\rangle}{2}$$

Key implementation feature: Batch evaluation over 200 training points uses jax.vmap(predict, in_axes=(None, 0)), compiling a single batched XLA kernel. This is theoretically expected to achieve ~200× throughput over PennyLane's sequential per-sample circuit calls (one XLA kernel vs. 200 Python dispatches); a direct wall-clock comparison was not measured in this work.

Results: Achieves 97%+ classification accuracy within 150 Adam epochs.

Figure 2: Variational Quantum Classifier. Left: learned decision boundary correctly separates XOR classes. Right: MSE loss convergence curve. Results shown from GPU execution; TPU produces identical final decision boundary.

4.3 Variational Quantum Eigensolver: H₂ Ground State

Objective: Find the ground-state energy of molecular hydrogen (H₂) in the STO-3G basis, mapped to a 4-qubit Hamiltonian via Jordan-Wigner transformation:

$$H = g_0 I + g_1 Z_0 + g_2 Z_1 + g_3 Z_2 + g_4 Z_3 + g_5 Z_0 Z_1 + \cdots + g_{14}(Y_0 Y_1 X_2 X_3)$$

with coefficients \((g_0 = -0.81054, g_1 = 0.17120, \dots)\) at equilibrium bond length \(R = 0.735\) Å.

Ansatz: 3-layer HEA initialized near the Hartree-Fock reference \(|0011\rangle\) (40 parameters). Gradient: Reverse-mode jax.grad — all 40 parameter gradients in one backward pass.

Results:

Figure 3: VQE H₂ molecular simulation. Top-left: energy convergence toward FCI reference (dashed red). Green shaded band marks chemical accuracy (±1.6 mHa). Top-right: gradient norm decay. Bottom-right: H₂ potential energy surface with FCI curve (green) and VQE result (star marker).

Figure 4: VQE convergence on TPU v5e-16. Energy (blue) converges to the FCI value of −1.1372 Ha within ~280 epochs. Gradient norm (purple) decays monotonically.

4.4 QAOA MaxCut (Combinatorial Optimization)

Circuit: QAOA at depths \(p \in {1,2,3,4,5}\) on a 6-node weighted graph (\(|E|=9\)):

$$|\vec{\gamma},\vec{\beta}\rangle = \prod_{k=1}^p e^{-i\beta_k H_M} e^{-i\gamma_k H_C}|+\rangle^{\otimes 6}$$

Results:

Classical MaxCut optimum: 9.0. QAOA at \(p=5\) achieves 99.7% of the classical optimum.

Figure 5: QAOA MaxCut on a 6-node weighted graph. Top-left: cut-value convergence for QAOA depths p=1…5; color corresponds to depth. Top-right: approximation ratio approaches 1.0 (classical optimum, dashed) as p increases. Bottom-right: graph topology with edge weights.

Figure 6: QAOA MaxCut on TPU v5e-16. Convergence curves per depth p confirm identical results to GPU execution, validating hardware-portability of the JAX implementation.

4.5 Quantum Noise Simulation (Monte Carlo Trajectories)

We implement stochastic quantum trajectory simulation for three standard noise channels:

Amplitude Damping (T1 relaxation, initial state \(|1\rangle\)): $$K_0 = \begin{pmatrix}1&0\0&\sqrt{1-\gamma}\end{pmatrix}, \quad K_1 = \begin{pmatrix}0&\sqrt{\gamma}\0&0\end{pmatrix}$$

Phase Damping (T2 dephasing, initial state \(|+\rangle\)): $$K_0 = \begin{pmatrix}1&0\0&\sqrt{1-\gamma}\end{pmatrix}, \quad K_1 = \begin{pmatrix}0&0\0&\sqrt{\gamma}\end{pmatrix}$$

Depolarizing Channel (initial state \(|+\rangle\)): $$\mathcal{E}(\rho) = (1-p)\rho + \frac{p}{3}(X\rho X + Y\rho Y + Z\rho Z)$$

Batch evaluation uses jax.vmap(simulate_trajectory, in_axes=(0, None)) over trajectory random keys — all trajectories execute as a single GPU kernel.

Results: Monte Carlo averages at \(N_\text{traj} \in {10, 100, 500}\) converge to exact analytical solutions with expected \(1/\sqrt{N_\text{traj}}\) statistical convergence rates. The jax.vmap parallelism is theoretically expected to yield ~\(N_\text{traj}\times\) throughput relative to sequential trajectory simulation (single fused XLA kernel vs. \(N_\text{traj}\) sequential dispatches); this ratio was not directly benchmarked against a sequential baseline.

Figure 7: Monte Carlo quantum noise simulation. Left: amplitude damping |1⟩ relaxation vs damping rate γ. Center: phase damping ⟨X⟩ decay of |+⟩. Right: depolarizing channel ⟨X⟩ decay. Yellow curve: exact analytical solution. Colored markers: 10/100/500-trajectory Monte Carlo averages converging to exact solutions.

4.6 Barren Plateau Gradient Variance Scaling

For a depth-\(d\), \(n\)-qubit Haar-random circuit with observable \(H\), the McClean et al. (2018) prediction [9]:

$$\text{Var}{\vec{\theta}}\left[\partial{\theta_k}\langle H\rangle\right] \in \mathcal{O}(2^{-n})$$

GPU findings: Gradient variance scaling benchmarks up to 15 qubits confirm the exponential decay predicted by theory. The gradient signal vanishes to XLA compiler precision noise (~\(10^{-7}\)) around 12 qubits for depth-4 random circuits.

TPU scaling advantage: The v5e-16 mesh enables gradient variance evaluation up to 24 qubits, extending the observable scaling region by an additional decade and confirming the McClean bound into the memory-limited regime for consumer GPU hardware.

Figure 8: Barren plateau study. Gradient variance ⟨(∂E/∂θ)²⟩ (y-axis, log scale) vs qubit count (x-axis). The exponential fit (dashed) confirms the O(2^{-n}) theoretical prediction from McClean et al. (2018). Gradient signal approaches compiler precision (~10⁻⁷) near 12 qubits.

4.7 Shor's Algorithm: 33-Qubit Distributed State-Vector Demonstration

[!IMPORTANT] Scope clarification: The factoring instances (N=15, 21, 35) are trivially solvable classically. The 33-qubit scale is an engineering demonstration of the distributed shard_map + ppermute QFT scheme at 64 GB state-vector scale — not a contribution to factoring capability.

Circuit pipeline (22 counting + 11 work qubits = 33 total):

1. Initialize: \(|0\rangle^{\otimes 22} \otimes |1\rangle_w\)

2. Hadamard superposition on counting register

3. Controlled modular exponentiation via shard_map (network spikes reduced from 8 GB to 128 MB via chunked ppermute)

4. Inverse QFT on counting register

5. Period extraction via continued fractions on measurement peaks \(s/2^{22} \approx j/r\)

Result: The 33-qubit simulation correctly extracts periods \(r \in {4, 6, 12}\) for \(N \in {15, 21, 35}\) respectively, and phase peaks at \(s \cdot 2^{22}/r\) match theoretical positions — validating the distributed QFT implementation at 64 GB state-vector scale.

4.8 Grover's Algorithm at 36 Qubits (TPU v6e-64)

Search space: \(N = 2^{36} \approx 6.87 \times 10^{10}\) states; state-vector requires 549.76 GB, distributed across 64 TPU v6e chips (~8.59 GB per chip).

Optimal iteration count: $$k_\text{opt} = \left\lfloor\frac{\pi}{4}\sqrt{2^{36}}\right\rfloor = 205{,}887$$

MPS entanglement analysis: Grover's diffusion operator \(R_s = 2|s\rangle\langle s| - I\) is globally non-local. Bipartite entanglement entropy \(S(A:B)\) grows rapidly with each oracle call. For MPS with \(\chi = 64\): \(S_\text{max} = \log_2 64 = 6\) bits — saturated within ~50 Grover iterations (< 0.025% of the required \(k_\text{opt}\)). Full statevector simulation is therefore necessary; MPS is not viable for Grover's algorithm above ~15 qubits.

4.9 MPS VQE at 512–1024 Qubits: Numerical Stability Analysis

While Sections 4.1–4.8 focused exclusively on full state-vector simulation (which is strictly memory-bound to ~36 qubits), this section transitions to Matrix Product States (MPS) to push into the 512–1024 qubit regime. This represents a distinct computational approach designed to handle scale by capping entanglement. Here, we investigate the numerical stability of deep MPS optimization.

512-qubit stability breakthrough: Differentiating through SVD in JAX causes gradient explosions due to Wirtinger calculus singularities at degenerate singular values. We document three fixes achieving stable convergence:

1. SVD epsilon floor (EPS = 1e-7): Prevents \(s_i \rightarrow 0\) division-by-zero in normalization

2. Site-level normalization: Prevents exponential amplitude drift across deep contractions

3. Wirtinger gradient clipping: Clips real components to \([-1.0, 1.0]\), bypassing JAX complex gradient instabilities

Result: Energy converges monotonically from 0.4718 to 0.4311 per site (vs. NaN crash without fixes).

1000-qubit bond dimension bottleneck: Reducing bond dimension to \(\chi = 64\) (TPU v5e-16 memory constraint) imposes \(S_\text{max} = 6.0\) bits entanglement ceiling. As training generates correlation beyond this ceiling, SVD truncation discards significant singular values, causing:

• V-bounce: Energy descends to ~0.464 at epoch 2, then rebounds to ~0.481 at epoch 4

• Singularity source: Passive site degeneracy (\(s_i \approx s_j \approx 0\)) causes \(\frac{1}{s_i^2 - s_j^2}\) → ∞ in JAX's SVD backward pass

Complete resolution (vqe_1024q_v5e16.py):

1. SVD jitter: \(10^{-9}\) complex noise before SVD breaks \(s_i = s_j\) degeneracy

2. Full Hamiltonian optimization: All qubits contribute to loss (no passive sites)

3. Momentum SGD (\(\mu = 0.9\)): Low-pass filters the V-bounce oscillation, enabling stable convergence to 10,000 epochs

Figure 9: Unstable MPS-VQE run (512 qubits, no stabilization). The energy curve shows a catastrophic spike followed by NaN divergence — caused by SVD derivative singularities when singular values approach zero degeneracy.

Figure 10: 1000-qubit MPS-VQE energy oscillation (V-bounce). The V-shaped rebound pattern appears when bond dimension χ=64 saturates the S_max=6 bit entanglement ceiling. Momentum SGD (μ=0.9) damps these oscillations, enabling monotonic convergence over 10,000 epochs.

5. Cross-Hardware Performance Benchmarks

5.1 Benchmark Methodology

Statistical rigor note: All benchmark data in Sections 5.3–5.7 were collected on 2026-05-30 using N=10 post-warmup timed runs (2 warmup runs discarded), with raw JSON logs archived in benchmarks/results/n10_benchmark_20260530_214024.json. We report mean ± σ computed from all 10 individual run measurements. GPU execution data at 25q/27q (Sections 5.4–5.5) retain †N=3 from original RTX 2050 hardware runs; the scaling sweep and gradient benchmarks (Sections 5.3, 5.6–5.8) are the new N=10 rigorous measurements. Independent reproducers may re-run benchmarks/n10_rigorous_benchmark.py (JAX 0.10.1+, CPU baseline) or benchmarks/cuda_vs_cpu_benchmark.py (GPU required).

5.2 Gate Application Speed (10-Qubit, Single CNOT)

Note: The JAX JIT first-call overhead is a one-time compilation cost amortized over all subsequent executions. The meaningful comparison for repeated inference or training is the "Subsequent calls" column.

5.3 Full Gradient Step — Measured vs. Reference Data

We measure gradient computation on a 15-qubit Hardware-Efficient Ansatz circuit with 3 entangling layers (120 trainable parameters). Results are from bench_D_gradient in n10_benchmark_20260530_214024.json (JAX 0.10.1, CPU backend, N=10 post-warmup runs).

5.3a Rigorous N=10 Measurements (15-Qubit HEA, 120 Params, CPU Backend)

Two independent N=10 benchmark runs were completed on the same CPU-only hardware. For jax.grad, the first timed run in both cases contains a residual JIT retracing event (106 ms in V5; 413 ms in V7) despite 2 declared warmup runs — this occurs because jax.value_and_grad retraces on the first differentiated call after warmup when gradient tape construction is deferred. We therefore report two statistics: the 10-run total mean (conservative, includes outlier) and the 9-run stable mean (runs 2–10, the physically meaningful post-JIT figure):

PSR 9-run mean = 10-run mean because PSR has no JIT retracing (it calls a pre-compiled function repeatedly).

V5 is the primary reference: the 9-run stable jax.grad mean (37.5 ms ± 1.8 ms) has negligible variance, confirming fully compiled steady-state execution. V7 was run under measurably higher OS load: the jax.grad 9-run stable mean of 107 ms ± 53 ms (49% CV) vs. V5's 37.5 ms ± 1.8 ms (5% CV) indicates a ~3× slowdown with high run-to-run jitter consistent with concurrent background processes competing for CPU time and L3 cache during the jax.value_and_grad backward pass¹. V7 PSR times (6,254 ms ± 2,016 ms, 32% CV) show the same OS-load signature. Both runs confirm jax.grad is >35× faster than PSR at 120 parameters; V5 is the primary reference for the clean steady-state figure.

¹ V7 was run on the same physical machine as V5 but at a different time of day with additional background tasks running. No system-level profiling was performed to isolate the exact source of the 3× slowdown.

V5 raw runs — jax.grad (ms): 106.11*, 35.87, 36.12, 35.67, 35.15, 36.43, 37.26, 40.34, 38.21, 39.77
V5 raw runs — PSR (ms): 1,759.98, 1,997.50, 1,775.52, 1,892.40, 1,805.04, 1,800.24, 1,692.28, 1,807.51, 1,881.80, 1,845.88
* Run 1 flagged as JIT retracing event; excluded from 9-run stable mean.
(Source: n10_benchmark_20260530_212827.json)

Primary speedup (V5, 9 stable runs, 120 params): jax.grad is 48.7× faster than PSR (37.5 ms vs 1,826 ms). Conservative 10-run figure: 41.4× (44.1 ms vs 1,826 ms).

5.3b Reference Data — Prior Literature / Preliminary Runs (50 Params, GPU, N=3)

[!NOTE] The rows below used 50 parameters on GPU hardware (RTX 2050, N=3 preliminary runs) and are not directly comparable to the 120-parameter CPU measurements above. They are included for orientation against published baselines. The PSR overhead scales linearly with parameter count: 50 params gives ~75× speedup; 120 params gives 48.7× (V5 primary, 9 stable runs) — both ratios are consistent with \(\text{speedup} \propto P\) where \(P\) is parameter count.

† Preliminary data; reproduced from the original GPU benchmark runs on RTX 2050. For cross-framework comparisons at matched parameter counts and hardware, re-run benchmarks/cuda_vs_cpu_benchmark.py with GPU access.

Key insight: The jax.grad advantage over PSR is algorithmic — it scales linearly with parameter count (\(P\)): at \(P=50\) (GPU, preliminary) we observe ~75×; at \(P=120\) (CPU, N=10 rigorous, 9 stable runs) we observe 48.7× (V5 primary) and 58.4× (V7). Both are consistent with \(\text{speedup} \approx P / C_{\text{rev}}\) where \(C_{\text{rev}}\) is the constant reverse-mode backward-pass cost. Furthermore, the 4× speedup over PennyLane's own JAX reverse-mode backend is the most direct apples-to-apples comparison (same autodiff paradigm, same hardware class). This 4× gap is driven by our simulator's monolithic XLA kernel fusion: our jax.lax.fori_loop compiles the entire circuit into a single optimized XLA graph, whereas PennyLane incurs Python-level dispatch overheads for tape construction and multi-kernel execution.

Figure A: Gradient benchmark (N=10, 15-qubit HEA, 120 parameters, CPU backend). Left: scatter of individual run times per method. Right: mean ± 1σ bar chart. jax.grad computes all 120 gradients in a single reverse-mode backward pass; PSR requires 240 forward evaluations (2 per parameter).

5.4 Qubit Scaling Sweep — N=10 Rigorous Measurements

We executed a scaling sweep from 4 to 20 qubits using a Hardware-Efficient Ansatz (HEA) with real RY/RZ rotations and nearest-neighbor CNOT entanglement. Circuits with \(n \leq 18\) used 3 entangling layers (L=3); circuits with \(n \geq 19\) used 2 layers (L=2) to stay within available RAM on the test machine†. All measurements are N=10 post-warmup timed runs from bench_C_scaling in n10_benchmark_20260530_214024.json. Circuit objects are constructed outside jax.jit to isolate pure execution time from compilation overhead.

† The benchmark script (n10_rigorous_benchmark.py V7, line 55) automatically reduces to L=2 for n≥19. Parameter count formula: \(P = n \times 2 \times (L+1)\), giving 152 for 19q/L=3 and 160 for 20q/L=3; the measured values 114 and 120 correctly reflect L=2 as used.

‡ n≥19 used L=2 layers (not L=3) to fit within available RAM. \(P = n \times 2 \times (L+1)\): at L=2, P=114 (19q) and P=120 (20q). At L=3 these would be 152 and 160 respectively.

The exponential scaling (\(O(2^n)\)) is clearly visible: from 10 qubits (0.36 ms) to 20 qubits (1,678 ms), execution time increases by approximately 4,670× — consistent with the theoretically expected \(2^{10} = 1,024×\) growth in state-vector size modulated by memory-bandwidth effects from cache eviction above ~16 qubits.

Figure B: Statevector simulation scaling (N=10 timed runs per data point). Log-scale y-axis. Blue line: measured jax_qsim execution times. Blue band: ±1σ across 10 runs. Orange dashed: theoretical O(2^n) reference. The inflection at 16–17 qubits corresponds to the state-vector (2 MB → 512 MB) exceeding L3 CPU cache, triggering a 10× slowdown attributable to DRAM bandwidth saturation rather than compute throughput.

Figure C: Individual run-by-run timing for three representative qubit counts. The 10-qubit circuit is cache-resident and shows sub-millisecond variance. The 20-qubit circuit exhibits 400 ms run-to-run variance driven by OS memory scheduling.

Figure D: Complete N=10 benchmark data summary. Source: n10_benchmark_20260530_214024.json. Green cells: measured execution times. Orange cells: gradient benchmark results.

5.5 JIT Compilation Overhead (Amortization Analysis — N=10 Measured)

XLA compilation times (measured, 2026-05-30 benchmark):

† 25-qubit data from original RTX 2050 hardware (requires 256 MB VRAM); 20-qubit is the largest tractable measurement on CPU-only JAX.

For training workflows with hundreds of gradient steps, the XLA compilation overhead is fully amortized after fewer than 20 calls at all tested qubit counts. The break-even point improves as circuit complexity grows — larger circuits have proportionally longer execution times relative to their (sub-linear) compile overhead.

Figure 14: GPU scaling benchmark (6-panel). Clockwise from top-left: execution time scaling, peak memory usage, gate throughput, VRAM utilization, compilation time, and speedup factor — all as functions of qubit count from 4 to 22 qubits on RTX 2050.

5.6 Maximum Qubit Threshold by Hardware

5.7 Negative Result / Scope Boundary: 37-Qubit RCS F_XEB Benchmark

The 37-qubit random circuit sampling uses TensorCircuit's tensor-network amplitude sampling (not full statevector). The underlying chaotic circuit operates on a 40-qubit 1D chain topology (20-layer RX/RZ + alternating CZ pattern) evaluated on TPU v6e-64, with TensorCircuit's contraction-path optimizer sampling individual output amplitudes:

Note: The F_XEB ≈ 0 result is physically expected and serves as a scope boundary for tensor-network methods. It indicates the sampled output distribution is approximately uniform — a null result expected for a deep chaotic circuit evaluated via tensor-network contraction with finite bond dimension. To achieve F_XEB ≈ 1.0 (indicating high fidelity), the simulator would need to retain exact full-statevector precision, which requires a prohibitive 8.79 TB of memory at 40 qubits, or exact tensor-network contraction with exponential bond dimension. This F_XEB value therefore characterizes the expected failure of the finite bond-dimension tensor-network approximation on deep chaotic circuits, not the underlying physical circuit quality.

6. Discussion

6.1 What This Work Demonstrates

Confirmed strengths:

1. Competitive GPU performance: Preliminary observations (see Appendix C) suggest our JAX simulator matches or outperforms PennyLane Lightning GPU and Qiskit-Aer GPU on the same hardware (RTX 2050) at 27 qubits.

2. Gradient computation advantage: The gradient advantage over PSR is architecturally principled and confirmed at 48.7× (N=10 rigorous, 120-param CPU circuit); the ~75× figure (2ms vs 150ms, 50-param GPU) is preliminary (N=3) and consistent with the linear \(\text{speedup} \propto P\) relationship. The 4× advantage over PennyLane's own JAX reverse-mode backend is also real and architecture-based.

3. TPU scalability: lax.fori_loop + PositionalSharding + jax.checkpoint together enable 33-qubit simulation on TPU v5e-16 that would be impossible with Python-loop-based approaches

4. Novel MPS numerical analysis: The SVD gradient singularity characterization and V-bounce identification are engineering contributions not documented in prior JAX-based MPS literature

6.2 Honest Limitations

Memory bandwidth is the binding constraint:
Preliminary tests at 25 qubits (256 MB state-vector) indicate the RTX 2050 (192 GB/s) is bandwidth-limited, not compute-limited. PennyLane Lightning CPU (using multi-core CPU DRAM at ~50–80 GB/s per-thread but higher effective bandwidth) wins at this scale on this specific hardware. However, a single run on a higher-bandwidth GPU (like an RTX 4090 or A100) would substantially flip this result in favor of the GPU.

Compilation overhead is real:
The first jax.jit call takes 1–18 seconds depending on qubit count. This is unsuitable for interactive circuit exploration but negligible for training workflows.

MPS accuracy is bond-dimension-limited:
At \(\chi = 64\), MPS can represent at most 6 bits of bipartite entanglement. Volume-law entanglement circuits (Grover, Shor, deep random circuits) exceed this ceiling rapidly. The 1024-qubit VQE result is physically meaningful only for near-product initial states with low entanglement.

No gate fusion pass:
Unlike cuQuantum, our simulator applies each gate as a separate memory operation. DAG-based gate fusion would theoretically reduce memory bandwidth requirements by 2–5×, which is identified as the primary optimization opportunity. Without gate fusion, our baseline simulator sits slightly below the absolute performance frontier of highly optimized C++ frameworks for raw gate application. However, we sacrifice this single-gate speed specifically to retain the end-to-end jax.grad differentiability that C++ simulators lack.

Statistical rigor — resolved:
Sections 5.3 (gradient timing) and 5.4 (scaling sweep 4–20 qubits) now report N=10 timed runs with mean ± σ from raw JSON logs (n10_benchmark_20260530_214024.json, timestamp 2026-05-30 21:40:24). The 25q/27q GPU data have been moved to Appendix C as preliminary observations since they retain †N=3 from the original RTX 2050 hardware sessions, and re-running 25q/27q XLA compilation on CPU-only hardware to achieve N=10 is impractical (requires 256 MB–1 GB VRAM).

7. Conclusion

We presented the JAX Quantum Research Suite, a differentiable quantum simulator built in pure JAX for statevector experiments and extending to TensorCircuit (JAX backend) for extreme-scale random circuit sampling. Key results:

• 29 qubits on consumer RTX 2050 GPU with full jax.grad differentiability

• 33 qubits via distributed sharding on Cloud TPU v5e-16 (Shor's algorithm demonstration)

• 36 qubits Grover's algorithm on TPU v6e-64 via 2 TB state-vector sharding

• 37-qubit RCS via TensorCircuit tensor-network amplitude sampling on TPU v6e-64

• 48.7× faster gradient computation vs PennyLane parameter-shift (N=10 rigorous, 120-param circuit); ~75× at 50 params (N=3 preliminary GPU); 4× vs PennyLane JAX reverse-mode

• Novel MPS numerical analysis (SVD jitter, V-bounce, Wirtinger clipping) enabling stable 1024-qubit VQE

• O(1) XLA graph size via lax.fori_loop enabling deep circuit simulation on TPU

The pure-JAX design makes any circuit composable with jax.grad, jax.vmap, and jax.pmap without code modification, providing a productive research tool for the NISQ algorithm development cycle. Because of this research and framework, researchers everywhere can now use our code to perform massive distributed quantum simulations directly on TPUs without having to write any complex C++ or CUDA distributed systems code.

Code: https://github.com/AshiteshSingh/Tpu-Accelerated-Quantum-JAX (MIT License)
Supported by: Google TPU Research Cloud (TRC) Program

Acknowledgements

The authors are deeply grateful to the Google TPU Research Cloud (TRC) program and the TRC team for providing access to Cloud TPU v6e-64chip and TPU v5e-16 hardware resources. This support enabled the large-scale distributed simulations (33-qubit Shor's, 36-qubit Grover's, 1024-qubit MPS-VQE) presented in this work.

References

[1] Preskill, J. (2018). Quantum Computing in the NISQ Era and Beyond. Quantum, 2, 79.

[2] Mitarai, K., Negoro, M., Kitagawa, M., Fujii, K. (2018). Quantum circuit learning. Physical Review A, 98(3), 032309.

[3] Bradbury, J., Frostig, R., et al. (2018). JAX: composable transformations of Python+NumPy programs. http://github.com/google/jax

[4] Duffield, S., Matos, G., Johannsen, M. (2023). qujax: Simulating quantum circuits with JAX. Journal of Open Source Software, 8(89), 5504. https://doi.org/10.21105/joss.05504

[5] Zhang, S.-X. et al. (2023). TensorCircuit: a Quantum Software Framework for the NISQ Era. Quantum, 7, 912. arXiv:2205.10091.

[6] Ittah, D., Asadi, A., Sanner, S., et al. (2024). Catalyst: A Python JIT compiler for auto-differentiable hybrid quantum programs. Journal of Open Source Software, 9(96), 6720. https://doi.org/10.21105/joss.06720

[7] Gray, J. (2018). quimb: A python library for quantum information and many-body calculations. Journal of Open Source Software, 3(29), 819. https://doi.org/10.21105/joss.00819

[8] Bergholm, V., Izaac, J., Schuld, M., et al. (2022). PennyLane: Automatic differentiation of hybrid quantum-classical computations. arXiv:1811.04968v4.

[9] McClean, J.R., Boixo, S., Smelyanskiy, V.N., Babbush, R., Neven, H. (2018). Barren plateaus in quantum neural network training landscapes. Nature Communications, 9, 4812.

[10] Jamadagni, A. et al. (2024). Benchmarking Quantum Computer Simulation Software Packages: State Vector Simulators. SciPost Physics Core.

[11] Shor, P.W. (1994). Algorithms for Quantum Computation: Discrete Logarithms and Factoring. Proceedings of FOCS 1994.

[12] Grover, L.K. (1996). A fast quantum mechanical algorithm for database search. Proceedings of STOC 1996.

[13] Farhi, E., Goldstone, J., Gutmann, S. (2014). A Quantum Approximate Optimization Algorithm. arXiv:1411.4028.

[14] Peruzzo, A. et al. (2014). A variational eigenvalue solver on a photonic quantum processor. Nature Communications, 5, 4213.

[15] NVIDIA Corporation (2023). cuQuantum SDK: cuStateVec State-Vector Library. https://developer.nvidia.com/cuquantum-sdk

[16] Vidal, G. (2003). Efficient Classical Simulation of Slightly Entangled Quantum Computations. Physical Review Letters, 91(14), 147902.

Appendix A: Circuit API Reference

from jax_qsim.circuit import Circuit
import jax, jax.numpy as jnp

c = Circuit(num_qubits=4)
c.h(0).cnot(0, 1).ry(2, param_index=0).rz(3, param_index=1)
c.noise_depolarizing(0, p=0.01)   # density_matrix mode only

params = jnp.array([0.5, 1.2])
state = c.run(params, state_type='statevector')   # JIT-compiled

# Gradient computation: all parameters in one backward pass
grad_fn = jax.jit(jax.grad(lambda p: compute_expectation(c.run(p, 'statevector'))))
grads = grad_fn(params)

# vmap over a batch of 64 parameter vectors simultaneously
batch_params = jnp.ones((64, 2))
batch_fn = jax.vmap(lambda p: c.run(p, 'statevector'))
batch_states = batch_fn(batch_params)   # all 64 circuits in one GPU kernel

Appendix B: TPU Cluster Deployment

# Provision TPU v5e-16 cluster
gcloud compute tpus tpu-vm create tpu-16chip-worker \
  --zone=us-central1-a \
  --accelerator-type=v5litepod-16 \
  --version=v2-alpha-tpuv5-lite

# Interactive control script
./tpu/run_tpu.sh
# Option 1: TERMINATE — kill zombie libtpu.so processes
# Option 2: SYNC & RUN — git sync all workers and run full suite  
# Option 3: DOWNLOAD — archive results and plots for download
# Option 4: CLEANUP — clear output directories

Appendix C: Preliminary Observations (25 & 27 Qubit GPU Benchmarks)

[!NOTE] The data in this appendix is derived from preliminary hardware sessions on an NVIDIA RTX 2050 (†N=3 runs) and is included as preliminary observations rather than rigorous N=10 benchmark claims.

C.1 25-Qubit State-Vector Benchmark (N=10 timed runs)

Full N=10 CPU baseline measurements completed (timestamp 20260530_212827):

V5 raw runs — 25q jax_qsim CPU (s): 19.23, 20.17, 22.01, 21.01, 25.41, 22.99, 19.02, 19.00, 19.12, 19.66
(Source: n10_benchmark_20260530_212827.json)

At 25 qubits, the CPU-only JAX baseline (20.76 s) is slower than both the GPU version (15.6 s) and PennyLane Lightning CPU (9.9 s). The CPU JAX slowness is expected: XLA statevector operates serially on a single cpu:0 device whereas PennyLane Lightning uses optimised multi-threaded C++ kernels. The GPU version achieves 15.6 s by parallelising across CUDA cores, approaching PennyLane Lightning throughput. On a GPU with 1+ TB/s bandwidth (e.g., RTX 4090, A100) the JAX simulator would outperform Lightning at 25 qubits.

Figure C.1: 25-qubit statevector benchmark (RTX 2050). At 25 qubits, PennyLane Lightning CPU (multi-threaded C++) achieves the fastest wall-clock time. The JAX GPU (RTX 2050 CUDA) is faster than JAX CPU but slower than Lightning at this state-vector size due to the 192 GB/s HBM bandwidth ceiling.

C.2 27-Qubit Cross-Framework Comparison (N=10 CPU Baseline + †N=3 GPU)

Full N=10 CPU baseline measurements for 27-qubit HEA (1 layer, 108 params):

V5 raw runs — 27q jax_qsim CPU (s): 160.45, 99.23, 109.34, 83.92, 92.57, 83.70, 92.90, 83.94, 96.37, 94.99
(Source: n10_benchmark_20260530_212827.json)

The CPU execution time (99.74 s) confirms the GPU is essential at 27 qubits: the RTX 2050 achieves 4.61 s — a 21.6× speedup over CPU JAX and a 1.3× speedup over PennyLane Lightning GPU. The first 27q CPU run (160.45 s) is an outlier likely caused by OS memory paging when the 1 GB state-vector first exceeds L3 cache; the 9 subsequent runs converge to 84–109 s.

At 27 qubits, the JAX JIT GPU achieves 1.3× speedup over PennyLane Lightning GPU and 1.5× over Qiskit-Aer GPU on identical hardware.

Figure C.2: 27-qubit GPU comparison. jax_qsim achieves the fastest execution across GPU frameworks at 27 qubits (1 GB state-vector). GPU is 21.6× faster than CPU-only JAX at this scale. All GPU measurements on NVIDIA RTX 2050.