Users can capture conditionals using qml.cond

[PietropaoloFrisoni]

Context: The purpose of this PR is to allow qml.cond to be captured into plxpr, although with some restrictions on its usage (see below for details).

Description of the Change: Now, qml.cond can be captured into plxpr. The function's behavior resembles the one provided in Catalyst, with the executed branch determined at runtime.

Benefits: If qml.capture.enabled() is True, qml.cond can be used.

Possible Drawbacks: Here there's a list of necessary restrictions at this stage.

• The arguments provided to the branches must be the same for all branches.

• The arguments provided to branches must be JAX-compatible. For example, this example would not work:

qml.capture.enable()

dev = qml.device("default.qubit", wires=3)

@qml.qnode(dev)
def circuit(pred, wires):
    def true_fn(wires):
        qml.RY(0.1, wires=wires)
    qml.cond(pred > 0, true_fn)(wires)
    return qml.expval(qml.PauliZ(wires=0))
circuit(1, qml.wires.Wires([0]))

because wires are not a valid JAX type, and the arguments provided to the branch are traced.

• Tests on nested conditionals have not been included as they belong to a different (future) epic, but in principle they should work

• If a branch returns one or more variables, every other branch must return the same abstract values.

For example, the following example would not work:

qml.capture.enable()

dev = qml.device("default.qubit", wires=3)

@qml.qnode(dev)
def circuit(pred, arg):
    def true_fn(arg):
        return qml.RY(arg, wires=0)
    def false_fn(arg):
        return qml.RZ(arg, wires=0), 5
    qml.cond(pred > 0, true_fn, false_fn)(arg)
    return qml.expval(qml.PauliZ(wires=0))

because true_fn and false_fn do not return the same abstract values. This is similar to Catalyst, where values produced inside the scope of a conditional can be returned to the outside context, but the return type signature of each branch must be identical.

• The function cannot branch on mid-circuit measurements, because these cannot be captured into plxpr yet.

Related GitHub Issues: None.

Related Shortcut Stories
[sc-66774]
[sc-69642]

[albi3ro]

One suggestion to try.

What would you think of combining jaxpr_true, jaxpr_false, and all the jaxpr_elifs? Same with combining the condition with the elifs_conditions.

We could essentially have:

jaxpr_branches = (jaxpr_true, *jaxpr_elifs, jaxpr_false)
conditions = jax.numpy.array([condition, *elifs_conditions, True])

And then the implementation could just be something like:

for condition, branch in zip(conditions, branches):
    if condition:
        return jax.core.eval_jaxpr(branch, consts, *args)

[mudit2812]

@PietropaoloFrisoni if you think this is in a reasonably good state, I would like to use this as a base for my PR so that I can capture qml.cond with predicates that use mid-circuit measurements.

[PietropaoloFrisoni]

@mudit2812 Sure, no problem for me. I'll tell you if I need to make relevant changes, but hopefully, this is not the case

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.65%. Comparing base (f9adf90) to head (db95c9b).
Report is 296 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5999      +/-   ##
==========================================
- Coverage   99.66%   99.65%   -0.01%     
==========================================
  Files         430      430              
  Lines       41544    41341     -203     
==========================================
- Hits        41404    41200     -204     
- Misses        140      141       +1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[dime10]

Nice work!

[mudit2812]

I created a branch locally that uses the changes from this PR as well as from #6015. Seems like after creating a primitive for qml.measure which has a scalar boolean abstract eval, using mid-circuit measurements with qml.cond just works out of the box for capture. It cannot be executed for the same reason why any other classical processing on MCMs can't be executed, but I guess I don't need to do any further work other than adding tests 😄

[albi3ro]

Excellent work 🎉