The manual Jacobian path (jacobian_manual.py:manual_jacobian, used for >5-qubit systems) builds gate-product chains in a sequential Python loop. For G gate segments this is O(G) sequential matmuls on the JAX device.
JAX's jax.lax.associative_scan can compute all prefix products in O(log G) parallel steps, which can significantly reduce wall time on GPU/multi-CPU for deep pulse sequences.
Proposed change:
Replace the explicit loop in manual_jacobian (lines 126–154 of jacobian_manual.py) with an associative_scan-based prefix-product computation, keeping the same output shape (G, d, d, K).
Acceptance criteria:
- Numerical output of the refactored function matches the current implementation to within JAX default tolerances.
- Existing tests in
test_geope.py covering Jacobians continue to pass.
- A benchmark (comment or docstring) documenting the speedup for a representative circuit depth.
The manual Jacobian path (
jacobian_manual.py:manual_jacobian, used for >5-qubit systems) builds gate-product chains in a sequential Python loop. ForGgate segments this is O(G) sequential matmuls on the JAX device.JAX's
jax.lax.associative_scancan compute all prefix products in O(log G) parallel steps, which can significantly reduce wall time on GPU/multi-CPU for deep pulse sequences.Proposed change:
Replace the explicit loop in
manual_jacobian(lines 126–154 ofjacobian_manual.py) with anassociative_scan-based prefix-product computation, keeping the same output shape(G, d, d, K).Acceptance criteria:
test_geope.pycovering Jacobians continue to pass.