[Frontend] PennyLane Measurements PyTrees

[rmoyard]

Context:
PennyLane 0.33.0 will introduce measurements as pytrees PennyLaneAI/pennylane#4607
Most measurements have therefore no leaves, this breaks a Catalyst assumptions for capturing the program.

Description of the Change:

• Unflatten the return to get it
• Flatten the return again but this time with is_leaf true for measurement processes.

Benefits:

Catalyst is up to date with PennyLane master

Possible Drawbacks:

Potential slow down
Benchmark:

import pennylane as qml
from jax import numpy as jnp

from catalyst import qjit

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

import timeit

def my_function_v1():

    @qjit
    @qml.qnode(device=dev)
    def circuit(x: float, y: float):
        qml.RX(x, wires=0)
        qml.RY(y, wires=1)
        qml.CNOT(wires=[0, 1])
        return [tuple([qml.expval(qml.PauliZ(wires=0))]), jnp.sin(y)], {"expval": qml.expval(qml.PauliZ(wires=1))}, tuple([qml.probs(wires=[0,1]), qml.expval(qml.PauliZ(wires=2))])

def benchmark_function(func):
    setup_code = f"from __main__ import {func.__name__}"
    stmt = f"{func.__name__}()"
    execution_time = timeit.timeit(stmt, setup_code, number=100)
    return execution_time

if __name__ == "__main__":
    time_v1 = benchmark_function(my_function_v1)
    print(f"Version execution time: {time_v1:.6f} seconds")

PL master:
For 100 runs: Version execution time: 6.361672 seconds
PL 0.32.0:
For 100 runs: Version execution time: 6.301733 seconds
Diff: around 0.06s for hundred runs

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (faf5934) 99.56% compared to head (690798e) 99.56%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #315   +/-   ##
=======================================
  Coverage   99.56%   99.56%           
=======================================
  Files          41       41           
  Lines        7345     7353    +8     
  Branches      428      430    +2     
=======================================
+ Hits         7313     7321    +8     
  Misses         16       16           
  Partials       16       16           

Files	Coverage Δ
frontend/catalyst/jax_tracer.py	99.56% <100.00%> (+<0.01%)	⬆️
frontend/catalyst/pennylane_extensions.py	99.63% <100.00%> (+<0.01%)	⬆️

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

[josh146]

Awesome work @rmoyard!

[rmoyard]

[sc-46568]

[maliasadi]

🎉

[rmoyard]

Need to wait for this one amazon-braket/amazon-braket-pennylane-plugin-python#210 and update the braket plugin version.

[dime10]

@rmoyard looks like it has been merged