Capture measurements with `jax.make_jaxpr`.

[albi3ro]

Context:

Description of the Change:

Benefits:

We can capture measurements and keep track of their resulting shapes in an extensible manner.

Possible Drawbacks:

Measurements do need a lot more hand-holding then observables due to the "duel mode" inputs of either wires or observables, and the need to fully specifying the resulting shape when we perform an actual measurment.

With the current framework, it will theoretically be extensible to add new measurements, but the process is slightly more error prone.

We also now provide duplicate information from MeasurementProcess.numeric_type, MeasurementProcess.shape, and MeasurementProcess._abstract_eval. But the instance-based property and method didn't play well with the need to track that information during an abstract evaluation.

Related GitHub Issues:

[sc-61200]

[codecov[bot]]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 99.67%. Comparing base (f5b805b) to head (08719a2). Report is 263 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5564      +/-   ##
==========================================
- Coverage   99.67%   99.67%   -0.01%     
==========================================
  Files         416      416              
  Lines       38686    38549     -137     
==========================================
- Hits        38562    38424     -138     
- Misses        124      125       +1     

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

[albi3ro]

One point that feels confusing is the qml.capture.measure during circuits. It correctly marks a quantum->classical transition, but it tells us that qml.capture.measure will not necessarily mark the endpoint of a circuit, which I thought it would 🤔

@dwierichs

We are eliminating the distinction between the "quantum part" and the "classical part", so terminal measurements don't necessarily need to be the last thing in the plxpr.

For example, here we retrieve samples, but then we postprocess them classically.

def f():
    mp = qml.sample(wires=(0,1))
    samples = qml.capture.measure(mp, shots=50)
    return jax.numpy.sum(*samples,axis=0)/50

jax.make_jaxpr(f)()

{ lambda ; . let
    a:AbstractMeasurement(n_wires=2) = sample_wires 0 1
    b:i32[50,2] = measure[num_device_wires=0 shots=Shots(total=50)] a
    c:i32[2] = reduce_sum[axes=(0,)] b
    d:f32[2] = convert_element_type[new_dtype=float32 weak_type=False] c
    e:f32[2] = div d 50.0
  in (e,) }

While this type of stuff won't necessarily be captured from the user in the current QNode, it dramatically increases what can do internally with transformations and compilation.

Maybe this wont' be the way we go long term, but I think it will be rather useful to be able to include classical postprocessing inside the same program.

[dwierichs]

Ah, yes, I should have been more clear with what I mean by "end of a circuit". I meant that we can not infer that a quantum circuit has been "finished", in the sense of deallocating/erasing the qubit register state, whether it is during a captured program or at its end. I was thinking about quantum gradients that need to dynamically find the scope of a quantum circuit and apply, say, the parameter-shift rule to the quantum parts of a hybrid program when differentiating it.

[josh146]

Note that this is something that we discussed with Romain and Mikhail last quarter during the discovery --- that is, to allow developers and user to easily write transforms when we have the new hybrid program capture, we need a way to (temporarily) 'separate' the hybrid program into classical and quantum 'blocks', so that transforms only apply to the quantum components.

This will be a huge quality of life improvement; in Catalyst currently, any quantum transformation in MLIR has to manually do this separation, apply the transform to the quantum part, and then glue the hybrid program back together.

@trbromley and I are thinking about having a discovery for transforms + PLXPR in Q3

[dime10]

in Catalyst currently, any quantum transformation in MLIR has to manually do this separation, apply the transform to the quantum part, and then glue the hybrid program back together

I don't think this is accurate for all transforms. The quantum and classical instructions live in the same scope, but you are free to traverse only operations from the quantum dialect, modify them, delete them, insert new ones, etc. This makes it particularly easy to write peephole optimizations on the quantum circuit. However for certain transformations, like differentiation, where you need to treat the classical scope and quantum scope fundamentally differently, and also transformations where you restructure or copy regions of quantum instructions, because you will also need to be aware of the classical values they depend on and where those come from.

One point that feels confusing is the qml.capture.measure during circuits. It correctly marks a quantum->classical transition, but it tells us that qml.capture.measure will not necessarily mark the endpoint of a circuit, which I thought it would 🤔

We are eliminating the distinction between the "quantum part" and the "classical part", so terminal measurements don't necessarily need to be the last thing in the plxpr. For example, here we retrieve samples, but then we postprocess them classically. I think it will be rather useful to be able to include classical postprocessing inside the same program.

Ah, yes, I should have been more clear with what I mean by "end of a circuit". I meant that we can not infer that a quantum circuit has been "finished", in the sense of deallocating/erasing the qubit register state, whether it is during a captured program or at its end. I was thinking about quantum gradients that need to dynamically find the scope of a quantum circuit and apply, say, the parameter-shift rule to the quantum parts of a hybrid program when differentiating it.

@albi3ro is there only one qml.capture.measure operation, or can there be multiple in different points of the program? I agree with David that it is somewhat important to have a well defined scope for where the qubits are "live", and to define the scope of what an instruction like qml.expval should act on. A situation like the following would leave a lot of ambiguity I think:

{
  # some gates
  qml.capture.measure(qml.expval)
  # some more gates
  qml.capture.measure(qml.probs)
}

What is the scope of the expectation value, some gates or some gates + some more gates? What about probs, is just some more gates or some gates + some more gates? Do the two measures imply two different circuits? Could you use the qml.expval value as argument to some more gates? Note that we did experiment in Catalyst with allowing measurement processes to appear anywhere in a quantum function, because on simulators those can just be computed on the current internal state without affecting anything. On hardware, since the measurement processes imply (repeated) execution of a complete quantum circuit, what that circuit is has to be well defined. I think this is one of the reasons other frameworks define measurement processes as acting on a "quantum kernel" or "quantum circuit", rather than being part of it.

I do agree that including post-processing in the program should definitely be a goal, although I'm not sure if it has to be in the same scope. The qnode primitive you are working on could delineate that scope for example, but it could also be something else.

Differentiation is definitely made easier when everything is separated into pre-processing, quantum, post-processing, but I don't think that structure can cover some of the more hybrid/dynamic constructs we are targeting anyways.

[josh146]

I don't think this is accurate for all transforms. The quantum and classical instructions live in the same scope, but you are free to traverse only operations from the quantum dialect, modify them, delete them, insert new ones, etc. This makes it particularly easy to write peephole optimizations on the quantum circuit. However for certain transformations, like differentiation, where you need to treat the classical scope and quantum scope fundamentally differently, and also transformations where you restructure or copy regions of quantum instructions, because you will also need to be aware of the classical values they depend on and where those come from.

@dime10 ah good to know! I was thinking of the error mitigation and gradient transforms when chatting to Romain --- these were examples where he brought up this separation last quarter.

[dwierichs]

Should the test file be renamed to tests/capture/test_measurements.py?

[albi3ro]

Looks good!

I didn't see anything about the "capture.measure" (not the mcm) operation in the code, has been shelved for now?

@dime10 Decided to break the capture.measure quantum-classical boundary into a later step to simplify the review process on this first component. Been shelved for now, and we can come back to some version of it when we need to work with mid circuit measurements better.

[dime10]

Looks good! I didn't see anything about the "capture.measure" (not the mcm) operation in the code, has been shelved for now?

@dime10 Decided to break the capture.measure quantum-classical boundary into a later step to simplify the review process on this first component. Been shelved for now, and we can come back to some version of it when we need to work with mid circuit measurements better.

Sounds good to me! But isn't this role filled by the qnode primitive now, that is the primitive both delineates the quantum scope as well turns abstract measurement return values into numeric data?

[albi3ro]

Looks good! I didn't see anything about the "capture.measure" (not the mcm) operation in the code, has been shelved for now?

@dime10 Decided to break the capture.measure quantum-classical boundary into a later step to simplify the review process on this first component. Been shelved for now, and we can come back to some version of it when we need to work with mid circuit measurements better.

Sounds good to me! But isn't this role filled by the qnode primitive now, that is the primitive both delineates the quantum scope as well turns abstract measurement return values into numeric data?

We'll need something like it anyway for mid circuit measurements. As for allowing it to be called on other measurements, we can probably wait to add that support till we actually need it and have a better idea of our constraints.