[Capture] Add support for control flow using autograph

[andrijapau]

Context:

Support for autograph in our native program capture pipeline was introduced last release (v0.40). This release, we wish to utilize this feature to add support for control flow when integrated with the QNode.

Description of the Change:

Source Code

Apply run_autograph to the QNode.func if autograph is requested (default to True). Tests were added to ensure proper functionality with Python native control flow (if/else, for, while) as well as with our Pennylane native control flow (cond, while_loop, for_loop).

Test Suite

• Added tests for new feature under test_capture_qnode.py in the TestQNodeAutographIntegration class.
• Specifically set the autograph kwarg to False in tests/capture/autograph since the point of those tests are to test the autograph functions themselves, and QNode should not be interfering.

Benefits:

We can now do things like (note autograph is defaulted to True),

qml.capture.enable()

dev = qml.device("default.qubit", wires=[0, 1, 2])

@qml.qnode(dev, autograph=True)
def for_loop_circuit(n):
    for i in range(n):
        qml.H(i)
    return qml.probs()

@qml.qnode(dev, autograph=True)
def while_loop_circuit(n):
    i = 0
    while i < n:
        qml.H(i)
        i += 1
    return qml.probs()

@qml.qnode(dev, autograph=True)
def conditional_circuit(x):
    if x > 1:
        qml.H(0)
    else:
        qml.I(0)
    return qml.probs()

>>> print(for_loop_circuit(2))
[0.24999997 0.         0.24999997 0.         0.24999997 0.
 0.24999997 0.        ]
>>> print(while_loop_circuit(2))
[0.24999997 0.         0.24999997 0.         0.24999997 0.
 0.24999997 0.        ]
>>> print(conditional_circuit(2))
[0.24999997 0.         0.24999997 0.         0.24999997 0.
 0.24999997 0.        ]

However, if you use autograph=False with these circuits, you will get,

CaptureError: Autograph must be used when Python control flow is dependent on a dynamic variable (a function input). Please ensure autograph=True or use native control flow functions like for_loop, while_loop, etc.

Possible Drawbacks:

Found a bug where if you use qml.while_loop with a lambda condition function, autograph will error out.

[sc-82285]

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.59%. Comparing base (61dbc71) to head (b7ee987).
Report is 4 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #6837      +/-   ##
==========================================
- Coverage   99.59%   99.59%   -0.01%     
==========================================
  Files         477      477              
  Lines       45233    45248      +15     
==========================================
+ Hits        45050    45063      +13     
- Misses        183      185       +2     

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