fixed an issue that arose due to custom_control/custom_inverse using …

[positr0nium]

Fix the custom_inversion / custom_control calling convention at its source

Description

custom_inversion and custom_control register a second Jaspr that stands in for the
decorated function when it is inverted or controlled. The two Jasprs are traced by
different mechanisms and the mechanisms disagree on how closed-over values are passed:

• The forward version goes through qache, i.e. jax.jit, so Jax performs its own
  closure conversion. Values the function captures become leading invars of the callee
  and extra leading operands of the wrapping jit equation.
• The registered variant is traced with make_jaspr, i.e. jax.make_jaxpr, which does
  not closure convert. The same values stay behind in constvars/consts, so the
  variant exposes fewer invars than the jit equation supplies, and it holds live
  tracers in consts.

Until now this was patched up at every use: invert_eqn and control_eqn folded the
constvars back into invars and rewrapped the result with Jaspr(normalized). That rewrap
produces a fresh object and silently drops everything the decorators had registered on
the Jaspr. PR #881 worked around one consequence by copying inv_jaspr across the
rewrap; ctrl_jaspr was still dropped by the same two lines, and on the control side the
rewrap downgraded a ControlledJaspr to a plain Jaspr, which no amount of field
copying can repair.

This PR fixes the mismatch where it originates. The variant is brought into the jit
calling convention once, where it is traced, and the transformation passes go back to a
plain .inverse() / .control(1). Since nothing is rewrapped at use time, nothing can be
lost.

Related Issues

Closes #
Related to #881 (this replaces the workaround introduced there)

Type of Change

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

Mostly a refactoring that removes the workaround, but it does fix two live defects on top
of what #881 addressed: ctrl_jaspr was dropped by the same rewrap, and a
ControlledJaspr was downgraded to a plain Jaspr, losing its base_jaspr/ctrl_state,
its inverse() override and its efficient nested control().

Breaking Change?

• Yes
• No

If yes, describe the impact and migration path:

Not breaking for working code. Code that was already broken now fails differently: a
custom_inversion/custom_control whose registered variant does not take the same
arguments as the function itself now raises at trace time, where the pair is registered,
instead of surfacing later as ValueError: Tried to evaluate jaxpr with insufficient arguments or as silently miswired operands.

What was changed?

• src/qrisp/jasp/jasp_expression/jaxpr_utils.py
  • fold_extra_constvars_into_invars moved here from inv_transform.py, with a
    docstring that describes the actual cause (the two tracing mechanisms) rather than
    "retracing reclassifies invars".
  • New closure_convert_jaspr(jaspr, insert_at=0): folds every constvar into an invar
    and rewraps without loss, carrying over inv_jaspr, ctrl_jaspr, permeability,
    isqfree and envs_flattened.
• src/qrisp/environments/custom_inversion_environment.py: the inverse is closure
  converted and its signature checked against the forward Jaspr before being cached as
  inv_jaspr.
• src/qrisp/environments/custom_control_environment.py: same for ctrl_jaspr, with
  insert_at=1 so the folded invars land behind the control qubit.
• src/qrisp/jasp/jasp_expression/inv_transform.py: invert_eqn is back to
  params["jaxpr"] = eqn.params["jaxpr"].inverse(). The fold in invert_jaspr, which
  normalizes the derived inverse after its retrace, stays.
• src/qrisp/jasp/jasp_expression/control_transform.py: control_eqn is back to
  new_params["jaxpr"] = eqn.params["jaxpr"].control(1). This is what lets a
  ControlledJaspr stay one.
• src/qrisp/jasp/jasp_expression/centerclass.py: check_aval_equivalence is now
  length-aware (it used zip, which truncates, so it could not detect an arity mismatch)
  and is now actually used. It had no callers before.
• src/qrisp/jasp/interpreter_tools/abstract_interpreter.py: reinterpret discarded the
  consts of its own retrace while keeping the input's, so any hoisting produced a
  ClosedJaxpr whose consts no longer lined up with its constvars. The two const lists
  are now concatenated in the same order as the constvars.
• Type hints added to the transformation passes: all module-level functions in
  inv_transform.py and control_transform.py, the ControlledJaspr methods, the two
  helpers in jaxpr_utils.py and check_aval_equivalence.
• tests/jax_tests/test_custom_inverse.py: two new tests (below). The three tests from
  Preserve the custom inverse when folding constvars back into invars #881 are kept unchanged.
• Changelog: the Preserve the custom inverse when folding constvars back into invars #881 entry was rewritten rather than a second one added, since it
  described a mechanism that no longer exists.

How was it tested?

No linked issue, so no Test-IDs. Run locally against the pinned jax==0.7.1.

Test	Status
tests/jax_tests/test_custom_inverse.py (6 tests, incl. the 3 from #881)	Pass
tests/jax_tests/test_control_flow_interpretation.py (incl. the qswitch invert/control regression test)	Pass
tests/jax_tests/test_jasp_qswitch.py, prepare and control modules (20 tests)	Pass
tests/jax_tests full suite	880 pass, 2 skipped, 31 fail
ruff format --diff on the changed files	Pass
ruff check on the changed files	Pass (no new findings)

The 31 failures are all ModuleNotFoundError for optional dependencies that were not
installed locally (catalyst, stim, xdsl), in test_mlir*, test_stim_simulation,
test_catalyst_interface, test_qubit_array_fusion, test_buffered_quantum_state and
test_to_qc_in_tracing_context. None are related to this change. ruff check reports the
same per-file counts as before the change for the test file (pre-existing F405
star-import noise) and strictly fewer for every source file.

Two new tests:

• test_registered_inverse_matches_forward_signature walks the Jaspr of a prepare with
  run-time amplitudes and asserts that every registered inverse carries no
  constvars/consts and takes exactly the arguments of the function it inverts. This pins
  the invariant directly rather than a downstream symptom.
• test_double_inversion_of_prepare_under_control_round_trips applies
  invert(control(invert(invert(prepare)))), so the outer inversion acts on an already
  controlled equation and the inverse callee is actually invoked.

Both were confirmed not to be vacuous: with the normalization temporarily disabled they
fail with ValueError: Tried to evaluate jaxpr with insufficient arguments: expected 7 (including 3 consts), got 10 - the three leaked constvars. Note that the three tests from
#881 still pass in that state, because two adjacent inversions cancel through the
back-pointer without ever invoking the inverse callee.

Screenshots / Output (if applicable)

n/a

Checklist

• My code follows the project's coding standards
• I have performed a self-review
• I have added/updated tests (referencing issue Test-IDs)
• All tests pass locally and in CI
• I have updated the documentation
• I have added a changelog entry
• Breaking changes are documented with migration path

Notes on the unchecked items:

• Tests pass locally apart from the missing-optional-dependency failures listed above. CI
  has not run yet.
• No user-facing documentation needed changing. The note in prepare about
  incompatibility with invert/control was already removed earlier; the changes here
  are to internals and their docstrings.
• Not a breaking change, so no migration path.

Reviewer Notes

• The load-bearing assumption. Folding drops the values of the folded constvars and
  relies on the caller supplying equivalent values positionally. That holds because Jax's
  own closure conversion put the same values in the same leading positions on the forward
  path. It was previously unchecked anywhere; check_aval_equivalence now compares the
  two signatures where both Jasprs are in scope. It compares aval types, so it catches
  an arity mismatch or a Qubit lined up against an array, not two same-typed arrays in
  swapped order. Worth a second opinion on whether a stronger comparison is wanted.
• The new exception in both decorators is the only place this PR can newly raise.
• The reinterpret change is the one part with no test coverage, because the path
  that hoists consts there does not currently fire. It is included because the remaining
  fold in invert_jaspr sits directly behind it and would otherwise slice a consts list
  of the wrong length.
• Why Jaspr(normalized) had to go rather than be made non-lossy. Copying the
  attributes across would still not preserve ControlledJaspr, since the rewrap changes
  the class. Normalizing at creation avoids the rewrap entirely.

Generated with Claude Code

[Copilot]

🟡 Changes recommended

Closure operands can still be silently miswired, and the new signature check accepts incompatible array abstract values.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Moves custom inverse/control closure conversion to registration time while preserving Jaspr metadata.

Changes:

• Adds closure-conversion and signature-validation helpers.
• Simplifies inversion/control transformations.
• Adds regression tests and corrects retraced constant handling.

File summaries

File	Description
tests/jax_tests/test_custom_inverse.py	Adds signature and nested transformation tests.
src/qrisp/jasp/jasp_expression/jaxpr_utils.py	Adds closure-conversion utilities.
src/qrisp/jasp/jasp_expression/inv_transform.py	Removes late inverse rewrapping.
src/qrisp/jasp/jasp_expression/control_transform.py	Removes late control rewrapping and adds typing.
src/qrisp/jasp/jasp_expression/centerclass.py	Makes signature comparison length-aware.
src/qrisp/jasp/interpreter_tools/abstract_interpreter.py	Preserves newly hoisted constants.
src/qrisp/environments/custom_inversion_environment.py	Normalizes and validates registered inverses.
src/qrisp/environments/custom_control_environment.py	Normalizes and validates registered controls.
documentation/source/general/changelog/changelog-dev.rst	Updates the bug-fix explanation.

Review details

• Files reviewed: 9/9 changed files
• Comments generated: 4
• Review effort level: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.