Allow standard qml.cond usage of qml.cond(pred, qml.some_gate)(*args, **kwargs) in qjit

[paul0403]

Context:
The standard qml.cond usage for single gates fails with qjit:

@qjit
@qml.qnode(dev)
def func(x, y):
    qml.cond(x == 42, qml.Hadamard)(wires=0)
    qml.cond(y == 37, qml.Hadamard)(wires=0)
    return qml.probs()

print(func(42, 37), func(42, 2))

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

While users can work around this by using catalyst.cond syntax, it is a bit incovenient.

Moreover, the fewer code changes required to qjit a pennylane circuit, the better. Since qml.cond just aliases to catalyst.cond within qjit, we should make sure qjit does not reject any valid plain pennylane use patterns.

Description of the Change:
The error arises because the conditional branch functions are not allowed to have arguments. This is because the CondCallable class's __call__() method does not take any arguments.

While we do not want to change the overarching design of CondCallable, we can make a new class that inherits from it, which specifically deals with the case where the conditional branch function is a pennylane gate (like qml.Hadamard, ...).

This new class, named CondCallableSingleGateHandler, has a __call__() method that can take in arbitrary args and kwargs. The class creates a wrapper conditional function taking in no arguments and sends the wrapper to the usual CondCallable, but the body of the wrapper calls the actual conditional branch function that is passed in. This allows us to perform the conditional branch function with arguments.

Benefits:
The standard usage qml.cond(predicate, qml.some_gate)(*args, **kwargs) works within qjit.

Related GitHub Issues: closes #449

[sc-55079]

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.97%. Comparing base (71b43c7) to head (e339f97).

Additional details and impacted files

@@            Coverage Diff             @@
##           v0.9.0-rc    #1232   +/-   ##
==========================================
  Coverage      97.96%   97.97%           
==========================================
  Files             77       77           
  Lines          11278    11318   +40     
  Branches         976      981    +5     
==========================================
+ Hits           11049    11089   +40     
  Misses           180      180           
  Partials          49       49           

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

[joeycarter]

Looks good @paul0403, thanks! 🚀

I only have a few minor suggestions. I was also wondering if you could comment on the fact that the else_if and otherwise methods in CondCallable also restrict their inputs to be functions without arguments. For example, would this be allowed/supported?

qml.cond(x > 42, qml.RY, qml.RX)(0.7, wires=0)

[rmoyard]

Thanks it looks good, just one question 👍

[paul0403]

Looks good @paul0403, thanks! 🚀

I only have a few minor suggestions. I was also wondering if you could comment on the fact that the else_if and otherwise methods in CondCallable also restrict their inputs to be functions without arguments. For example, would this be allowed/supported?

qml.cond(x > 42, qml.RY, qml.RX)(0.7, wires=0)

Good catch! I'll add the other conditional branches.

[github-actions]

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md on your branch with:

• A one-to-two sentence description of the change. You may include a small working example for new features.
• A link back to this PR.
• Your name (or GitHub username) in the contributors section.

[willjmax]

Thanks for this! This gets us one step closer to making PennyLane's iterative_qpe operation QJIT compatible. The remaining issue is that catalyst.cond does not accept MeasurementValue type for the input conditional. I opened an issue addressing this #1237.

[paul0403]

Looks good @paul0403, thanks! 🚀

I only have a few minor suggestions. I was also wondering if you could comment on the fact that the else_if and otherwise methods in CondCallable also restrict their inputs to be functions without arguments. For example, would this be allowed/supported?

qml.cond(x > 42, qml.RY, qml.RX)(0.7, wires=0)

Done, now all three branches (true, false, sequence of else_ifs) are all supported