Jasp refactoring (Part 2)

[PietropaoloFrisoni]

Description

I worked with Claude to remove a lot of boilerplate, dead code, code repetitions and perform code improvements (mostly based on pylance, pylint and ruff, as usual) across the Jasp infrastructure . We also catched and fixed some bugs during the process. Review should be hopefully easy as this is basically a massive internal refactoring with no consequences.

Removes duplicated and dead code across Qrisp's jasp (JAX-tracing) subsystem: the same control-flow, equation-copying, and caching boilerplate had been independently copy-pasted across the Catalyst, classical-boolean-simulation, profiling, and post-processing interpreter backends. Everything beyond the dedup itself is static-analysis cleanup (pylint/pyright/ruff) surfaced while touching these files, plus the handful of real bugs that surfaced along the way.

Related Issues

Closes #
Related to #

Type of Change

• Feature (new functionality)
• Change Request (modification of existing functionality)
• Bug Fix
• Refactoring (no behavior change) — primary focus of this PR
• Performance improvement
• Documentation
• CI / Build

Breaking Change?

• Yes
• No

stimulate()'s signature narrowed from stimulate(func: Callable | None = None) to stimulate(func: Callable). In practice this is not expected to affect anyone: the None default was already dead (no handling existed for it, so @stimulate() with empty parens would have failed downstream anyway), and no call site anywhere in the codebase used the zero-arg form. Migration path: use @stimulate (no parentheses), which is how it was already always used.

What was changed?

Centralization (the main point of this PR):

• Unified the "replay while/scan under an active trace" logic (previously duplicated across the Catalyst, classical-boolean-simulation, profiling, and post-processing interpreters) into evaluate_while_loop_under_trace/evaluate_scan_under_trace in control_flow_interpretation.py; same for cond into evaluate_cond_under_trace where the two consumers shared the same strategy.
• Unified flatten_signature/unflatten_signature (Catalyst vs. classical-boolean-simulation interpreters) and the copy_jaxpr_eqn/_copy_eqn equation-copying helper (previously reimplemented 5 times across environment_flattening.py, composite_gate_interpreter.py, inv_transform.py, injection_transform.py, and conjugation_environment.py).
• New shared helpers replacing other scattered duplicates: insert_call_outvalues (call-equation output wrapping, 4 sites), build_metric_profiler/get_cached_jaspr (profiling-metric boilerplate), rebuild_jaxpr (new jasp_expression/jaxpr_utils.py, 3 sites), _wrap_branch_for_tracing/_get_last_cond_eqn/_finalize_branch_eqn (q_cond/q_switch), tracing_scope (exception-safe tracing context manager, qache/quantum_kernel), _bump_gate_count (BufferedQuantumState), _decompose_and_exec (composite-gate decomposition), _register_vars/_slice_field (VarTracker), get_op_counts (gate-count breakdown).
• Removed dead code: Jaspr.inline() (zero callers anywhere, now delegates to embedd(inline=True)), a redundant simulator-validation branch in simulate_jaspr (identical error already raised one line later by BufferedQuantumState.__init__).

Bug fixes found while centralizing:

• evaluate_scan_under_trace: a while/scan with exactly one carried value crashed (carry/pytree-structure mismatch) across all four consumer backends — the duplicated copies all shared this bug.
• extract_post_processing: a KeyError when a jaxpr-level constant (e.g. a scan's xs array) was referenced, from consts never being bound into the evaluation context.
• AbstractQubitArray._getitem: an off-by-one slice-bound error (dead code, no reachable caller, but now shares correct logic with the live DynamicQubitArray.__getitem__ path).

Static-analysis cleanup: type hints, pylint findings, and pyright errors resolved across primitives/, jaspification.py, boolean_simulation.py, BufferedQuantumState, and the control-flow interpreter files — no intended behavior change. Two latent type-hint bugs (Jaspr.debug_info wrongly Optional, Jaspr.count_ops/.depth/.num_qubits declaring meas_behavior: str when str | Callable is actually accepted) were corrected as part of this.

New tests: coverage added for previously-untested lines across test_basic_primitives.py, test_boolean_simulation.py, test_buffered_quantum_state.py (new), test_catalyst_interface.py, test_count_ops.py, test_jaspify.py, test_post_processing_extraction.py, test_stim_simulation.py (new) — including regression tests for the num-carry-one scan bug across all four affected backends.

How was it tested?

No issue-linked Test-IDs apply here (this PR isn't tied to a specific bug-tracker issue), so this table describes verification method rather than Test-IDs:

Verification	Status
New/updated pytest tests (8 files, listed above)	✅
Static analysis: pyright, pylint, ruff (format + lint) on every touched file	✅
Fresh-process import checks after every change to eagerly-loaded modules	✅
Manual functional verification of the centralization batch not yet covered by pytest (see "further tests to add" in reviewer notes)	✅ (manual only)

Screenshots / Output (if applicable)

N/A — internal refactor, no user-visible output changes.

Checklist

• My code follows the project's coding standards
• I have performed a self-review
• I have added/updated tests (referencing issue Test-IDs) — partial: see reviewer notes, the most recent centralization batch has manual-only verification
• All tests pass locally and in CI — please confirm in CI; two test failures found post-push (exact-match error-message assertions) have been fixed
• I have updated the documentation — not needed; stimulate's doc page is Sphinx autofunction-generated from the signature
• I have added a changelog entry — this repo has no CHANGELOG file to update
• Breaking changes are documented with migration path

Reviewer Notes

• The stimulate() signature change (see "Breaking Change?" above) is the only public-API-visible change in this PR; everything else is internal interpreter/tracing-tool plumbing.
• Worth a close look: evaluate_scan_under_trace's num-carry-one fix, since it changes behavior for a previously-crashing case across four backends at once (Catalyst, classical simulation, profiling, post-processing) — regression tests were added for all four.
• Not yet covered by automated tests (manual verification only): the q_cond/q_switch state-stripped branch path, tracing_scope's exception path, get_cached_jaspr's cache-hit path. Happy to add these before merge if preferred.
• A follow-up PR could pick up ~11 more duplication-removal opportunities identified during the same audit (not included here to keep this PR reviewable).

[Copilot]

Pull request overview

This PR continues the JASP refactor by centralizing repeated tracing/interpreter patterns, strengthening type hints, and adding regression tests around Stim dispatch, profiling/post-processing, and JAX control-flow edge cases (notably scan with num_carry == 1).

Changes:

• Introduces shared utilities for rebuilding Jaxprs and copying equations, and reuses them across multiple transform/interpreter modules.
• Refactors control-flow handling to preserve cond/while/scan structure under trace, and reduces duplicated interpreter logic.
• Adds targeted regression tests for Stim gate dispatch, scan(num_carry==1) across multiple backends, and error cases (returning QuantumVariables).

Reviewed changes

Copilot reviewed 47 out of 47 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
tests/jax_tests/test_stim_simulation.py	Adds Stim dispatch coverage tests and a stimulate return-value error test.
tests/jax_tests/test_post_processing_extraction.py	Adds regression test for scan post-processing extraction with scalar carry.
tests/jax_tests/test_jaspify.py	Adds regression test that jaspify rejects returning QuantumVariables.
tests/jax_tests/test_count_ops.py	Adds regression test for scan profiling with scalar carry.
tests/jax_tests/test_catalyst_interface.py	Adds regression test for scan with scalar carry under qjit (if Catalyst installed).
tests/jax_tests/test_buffered_quantum_state.py	New direct unit tests for BufferedQuantumState constructor/copy behavior.
tests/jax_tests/test_boolean_simulation.py	Adds regression test for scan with scalar carry under boolean simulation.
tests/jax_tests/test_basic_primitives.py	Adds direct tests for abstract type __eq__/__hash__/__repr__.
src/qrisp/jasp/tracing_logic/tracing_quantum_session.py	Adds tracing_scope helper contextmanager to centralize exception-safe tracing scopes.
src/qrisp/jasp/tracing_logic/quantum_kernel.py	Refactors to use tracing_scope for exception safety.
src/qrisp/jasp/tracing_logic/qaching.py	Refactors to use tracing_scope for exception safety.
src/qrisp/jasp/tracing_logic/dynamic_qubit_array.py	Reuses shared slice normalization and qubit caching helpers.
src/qrisp/jasp/program_control/prefix_control.py	Refactors cond/switch tracing logic into helpers for reuse.
src/qrisp/jasp/primitives/quantum_primitive.py	Documents intended Primitive subclassing/registration behavior.
src/qrisp/jasp/primitives/parity_primitive.py	Improves dtype checks and refactors array parity logic readability.
src/qrisp/jasp/primitives/operation_primitive.py	Refactors imports and improves error message formatting/clarity.
src/qrisp/jasp/primitives/measurement_primitive.py	Refactors imports, adds docstrings, and simplifies branching.
src/qrisp/jasp/primitives/kernel_primitives.py	Renames abstract-eval functions for clarity and ignores unused args.
src/qrisp/jasp/primitives/abstract_qubit.py	Adds docstring and clarifies deferred import rationale.
src/qrisp/jasp/primitives/abstract_quantum_state.py	Refactors imports, adds docstrings, and clarifies impl behavior.
src/qrisp/jasp/primitives/abstract_quantum_register.py	Centralizes slice normalization and qubit caching; adds docstrings.
src/qrisp/jasp/jasp_expression/jaxpr_utils.py	New rebuild_jaxpr helper to centralize Jaxpr reconstruction.
src/qrisp/jasp/jasp_expression/inv_transform.py	Switches to shared eqn copy + jaxpr rebuild utilities and reinterpretation.
src/qrisp/jasp/jasp_expression/injection_transform.py	Removes local eqn-copy helper in favor of shared implementation.
src/qrisp/jasp/jasp_expression/environment_collection.py	Uses shared eqn copy + rebuild_jaxpr; refactors VarTracker internals.
src/qrisp/jasp/jasp_expression/control_transform.py	Uses rebuild_jaxpr instead of a local Jaxpr copy helper.
src/qrisp/jasp/jasp_expression/centerclass.py	Simplifies inline path; widens meas_behavior types; tweaks debug_info typing.
src/qrisp/jasp/jasp_expression/init.py	Exposes new jaxpr utilities via package import.
src/qrisp/jasp/interpreter_tools/interpreters/utilities.py	Adds get_op_counts and reuses it for operation discovery.
src/qrisp/jasp/interpreter_tools/interpreters/qc_extraction_interpreter.py	Uses new insert_call_outvalues for call-like primitives.
src/qrisp/jasp/interpreter_tools/interpreters/profiling_interpreter.py	Refactors cond/while/scan to shared under-trace evaluators; adds profiler builder + cached Jaspr helper.
src/qrisp/jasp/interpreter_tools/interpreters/post_processing_interpreter.py	Refactors control flow evaluation and fixes constvar binding in context.
src/qrisp/jasp/interpreter_tools/interpreters/pjit_flattening.py	Uses insert_call_outvalues for correct single-outvar handling.
src/qrisp/jasp/interpreter_tools/interpreters/num_qubits_metric.py	Reuses shared metric-profiler builder.
src/qrisp/jasp/interpreter_tools/interpreters/environment_flattening.py	Uses shared eqn-copy helper.
src/qrisp/jasp/interpreter_tools/interpreters/depth_metric.py	Reuses shared metric-profiler builder.
src/qrisp/jasp/interpreter_tools/interpreters/count_ops_metric.py	Reuses shared metric-profiler builder and shared op counting.
src/qrisp/jasp/interpreter_tools/interpreters/control_flow_interpretation.py	Adds under-trace cond/while/scan evaluation and signature flatten/unflatten helpers.
src/qrisp/jasp/interpreter_tools/interpreters/composite_gate_interpreter.py	Reuses shared eqn-copy helper and consolidates decomposition/exec logic.
src/qrisp/jasp/interpreter_tools/interpreters/cl_func_interpreter.py	Reuses shared scan-under-trace + signature (un)flattening helpers; improves typing.
src/qrisp/jasp/interpreter_tools/interpreters/catalyst_interpreter.py	Reuses shared scan-under-trace + signature (un)flattening helpers.
src/qrisp/jasp/interpreter_tools/abstract_interpreter.py	Strengthens typing and introduces shared copy_jaxpr_eqn/insert_call_outvalues.
src/qrisp/jasp/evaluation_tools/profiler.py	Reuses shared cached-Jaspr helper for decorator caching.
src/qrisp/jasp/evaluation_tools/jaspification.py	Refactors simulate_jaspr jit handling into helper; improves typing and correctness commentary.
src/qrisp/jasp/evaluation_tools/buffered_quantum_state.py	Refactors stim dispatch via a method-name table; improves typing and copy semantics.
src/qrisp/jasp/evaluation_tools/boolean_simulation.py	Refactors documentation/comments and improves typing/narrowing in decorator.
src/qrisp/environments/conjugation_environment.py	Removes local eqn-copy helper in favor of shared implementation.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[renezander90]

Deviates from the NumPy style used throughout the Qrisp documentation.

[PietropaoloFrisoni]

Thanks, I should tell my friend Claude to use the same convention as well (this are not user-facing functions, but it is worth using the same convention everywhere)

[renezander90]

Deviates from the NumPy style used throughout the Qrisp documentation.

[renezander90]

Deviates from the NumPy style used throughout the Qrisp documentation.

[renezander90]

LGTM 🎊