[WIP] Allow catalyst.cond to take in branch functions with arguments

[paul0403]

Context:
There is a restriction on catalyst.cond, that the conditional branch functions can never have arguments. This feature should be allowed, especially as we move towards FTQC.

Description of the Change:
The strategy used in #1232 to allow cond to take in pennylane gates is extended to apply to all callables with arguments.

Benefits:
catalyst.cond can take in branch functions with arguments, e.g

@qml.qnode(qml.device("lightning.qubit", wires=2))
def ref_func():
    qml.PauliX(wires=1) # |01>
    m0 = qml.measure(0)  # will measure 0

    def true_fn(wire):
        qml.PauliX(wires=wire)

    def false_fn(wire): # will come here
        qml.RX(1.23, wires=wire+1)

    qml.cond(m0 == 1, true_fn, false_fn)(0)

    return qml.probs()


print("ref: ", ref_func())


@qjit
@qml.qnode(qml.device("lightning.qubit", wires=2))
def func():
    qml.PauliX(wires=1) # |01>
    m0 = catalyst.measure(0)  # will measure 0

    @catalyst.cond(m0 == 1)
    def conditional(wire):
        qml.PauliX(wires=wire)

    @conditional.otherwise
    def false_fn(wire): # will come here
        qml.RX(1.23, wires=wire+1)

    conditional(0)

    return qml.probs()

print("cat: ", func())

ref:  [0.33288114 0.66711886 0.         0.        ]
cat:  [0.33288114 0.66711886 0.         0.        ]

[mlxd]

Thanks @paul0403 for the quick turnaround.

Just some stuff that was brought up: let's hold off on merging this for now, as we have a few things to resolve elsewhere in the ecosystem first, such as:

• Does it require modification/changes/updates to be compatible with the ongoing PLxPR conversions?
• Does it break the qml.cond dispatch we have in PennyLane?
• There exists a lot of documentation and tutorials requiring updates that currently mention not to allow this.

We can keep the PR, but maybe only in a WIP state.

[mlxd]

As commented above

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 95.96%. Comparing base (f1ac6b7) to head (b28d813).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1531      +/-   ##
==========================================
- Coverage   95.99%   95.96%   -0.03%     
==========================================
  Files          80       80              
  Lines        8411     8411              
  Branches      780      780              
==========================================
- Hits         8074     8072       -2     
- Misses        281      282       +1     
- Partials       56       57       +1     

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

[josh146]

Hey @paul0403, yep I have the same questions as Lee, but perhaps one additional one:

• Does this affect @rauletorresc's PLxPR conversion work?
• Does it break the qml.cond dispatch we have in PennyLane? Will we need to also update that?
• Does this match the semantics of qml.cond?

We need to also make sure we also update the documentation in this PR, as there are a few places (autograph guide, sharp bits, cond docstring) etc. that does not specify this behaviour or says this is not possible.

[paul0403]

Hi @josh146 , this PR is in essence a hot fix for @mlxd to experiment FTQC on Catalyst. It is not tied to any product feature for now (hence no shortcut story). I can look into those if we decide that this is worth spending time polishing.

But just a quick thought, I don't think this breaks anything in Raul's cond capture because this PR is after his capture PR #1468 went in and none of the tests broke.

[rauletorresc]

Hey @paul0403, yep I have the same questions as Lee, but perhaps one additional one:

• Does this affect @rauletorresc's PLxPR conversion work?
• Does it break the qml.cond dispatch we have in PennyLane? Will we need to also update that?
• Does this match the semantics of qml.cond?

We need to also make sure we also update the documentation in this PR, as there are a few places (autograph guide, sharp bits, cond docstring) etc. that does not specify this behaviour or says this is not possible.

I have added a test for PLxPR capture of conditional branches taking arguments: it seems to work out of the box :) So, it does not affect the PLxPR work so far.

[rauletorresc]

LGTM

[paul0403]

I have added a test for PLxPR capture of conditional branches taking arguments: it seems to work out of the box :) So, it does not affect the PLxPR work so far.

Thanks!

Does this new test fail without this PR and give the error message "catalyst.cond branch functions cannot have arguments"? Just for my peace of mind

[rauletorresc]

I have added a test for PLxPR capture of conditional branches taking arguments: it seems to work out of the box :) So, it does not affect the PLxPR work so far.

Thanks!

Does this new test fail without this PR and give the error message "catalyst.cond branch functions cannot have arguments"? Just for my peace of mind

It does fail: TypeError: Conditional 'True' function is not allowed to have any arguments