Backend sampling

[positr0nium]

Description

Introduces @backend_sampler — a Jasp decorator that routes sample() and expectation_value() calls through a real quantum backend instead of the Jaspify simulator. The quantum circuit is extracted once, executed on the backend for all shots, and the classical post-processing is replayed via the Jaspr's own while-loop — JIT-compiled through JAX. Control-flow primitives (fori_loop, while, cond, scan, jit) are handled via a custom equation evaluator that propagates through all nesting levels. Example:

from qrisp import QuantumFloat, h, measure
from qrisp.jasp import sample, backend_sampler
from qrisp.interface import StimBackend

backend = StimBackend()

@backend_sampler(backend=backend)
def main(k):
    def kernel(k):
        qf = QuantumFloat(4)
        h(qf[0])
        return measure(qf)
    return sample(kernel, shots=100)(k)

result = main(1)
# result is a JAX array of shape (100,) with backend results

Related Issues

Closes #
Related to the sampling_upgrades branch

Type of Change

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

Breaking Change?

• Yes
• No

What was changed?

• New decorator: @backend_sampler(backend=...) — routes sample() / expectation_value() calls through a real backend. Uses make_jaxpr (not make_jaspr) for the outer trace — no quantum tracing context, so direct quantum operations outside sampling kernels fail cleanly at trace time.

• New interpreter module: backend_sampling_interpreter.py — contains the Jaspr-level logic:

  • _extract_to_qc_args() — mini-interpreter that runs the eval Jaspr until the first sampling_body_func call, capturing all arguments automatically (handles JAX-implicitly-prepended closure variables).
  • _make_backend_sampling_fn() — factory returning a backend_sampling_fn that extracts the circuit via Jaspr.to_qc(), runs it once on the backend, and replays the post-processing loop.
  • _body_loop_evaluator() — intercepts sampling_body_func in the eval Jaspr, replacing it with Jaspr.extract_post_processing() using pre-computed measurement bits.

• Control-flow propagation — the outer equation evaluator handles jit/pjit, while, cond, and scan primitives by recursively calling eval_jaxpr with the custom evaluator, ensuring sample()/expectation_value() calls are intercepted regardless of nesting depth inside control flow.

• Real-time feedback detection — catches the "Tried to convert real-time feedback into QuantumCircuit" exception from to_qc() and re-raises as RuntimeError with a clear message pointing to jaspify.

• _backend_shots_marker — identity @jax.jit marker inserted at the top of sampling_eval_function and expectation_value_eval_function so the backend sampler can reliably locate the shot count in the traced Jaxpr without fragile position-based extraction.

• Supporting changes:

  • sampling.py / ev.py: added _backend_shots_marker call; renamed state_prep → sampling_kernel in expectation_value for consistency.
  • measurement_primitive.py: fixed measure_implementation to index into the qubit array.
  • tracing_quantum_session.py: allow qubit reuse from templates.
  • qaoa_problem.py: accept QuantumVariableTemplate in qarg parameter.
  • qc_extraction_interpreter.py: support QuantumVariableTemplate and parity handle fixes.

• Documentation:

  • Enriched backend_sampler docstring with usage examples, parameter docs, Raises section, and architecture notes.
  • Added BackendSampling.rst to Sphinx docs and updated Jasp/index.rst toctree.

• Tests (57 total):

  • 51 tests covering sampling patterns, expectation values, corner cases, error handling.
  • 6 control-flow propagation tests: fori_loop, cond, while_loop, scan, nested jit, switch.

How was it tested?

All 57 tests pass locally:

pytest tests/jax_tests/test_backend_sampling.py \
       tests/jax_tests/test_backend_sampling_ev.py \
       tests/jax_tests/test_to_qc_in_tracing_context.py -q
# 57 passed in ~25s

Checklist

• My code follows the project's coding standards
• I have performed a self-review
• I have added/updated tests
• All tests pass locally
• I have updated the documentation
• I have added a changelog entry
• Breaking changes are documented with migration path

Reviewer Notes

• The split between backend_sampling.py (decorator + outer evaluator) and backend_sampling_interpreter.py (Jaspr interpreters) keeps the architecture modular.
• The make_jaxpr vs make_jaspr choice is deliberate: make_jaxpr avoids opening a quantum tracing context for the outer orchestration function, so any accidental direct quantum operations fail at trace time with a clear error.
• The control-flow handlers (jit/while/cond/scan) follow the same pattern as the post-processing interpreter — propagating the custom evaluator downward via eval_jaxpr(eqn_evaluator=...).