Update QNode for PyTree support

[andrijapau]

Context:

Catalyst currently supports the handling of PyTrees as measurement outputs,

@qml.qnode(dev)
def circuit():
    return {"Exp Value": qml.expval(qml.Z(0)), "Probability": qml.probs()}

Pennylane currently doesn't have support for this but can be added using the existing qml.pytrees.

Description of the Change:

Makes use of the existing qml.pytrees.flatten and qml.pytrees.unflatten functions and updates how the results and qfunc_output is handled/validated.

Example

@qml.qnode(qml.device("default.qubit", wires=3, shots=500))
def circuit():
    qml.Hadamard(0)
    qml.measure(0)
    qml.CNOT([0,1])
    return (
             [
                  qml.probs(wires=1),
                  {"a": qml.probs(wires=0)}, 
                  qml.expval(qml.Z(0))
             ] , 
             {
                  "b": qml.probs(wires=0), 
                  "c": qml.expval(qml.X(1))
             }
     )

>>> circuit()
([array([0.494, 0.506]), {'a': array([0.494, 0.506])}, -0.012], {'b': array([0.494, 0.506]), 'c': -0.096})

Benefits: QNode now has PyTree support! 😄

Possible Drawbacks: None

[sc-66732]

[albi3ro]

Perhaps here we can switch this to measurement_processes = qml.pytrees.flatten(self._qfunc_output)[0].

[albi3ro]

Potentially we could switch to storing self._pytree_structure instead of self._qfunc_output. Basically we just need to store information somehow to remember how to postprocess the results.

[andrijapau]

Yes, that could be an improvement. However, the typing of qfunc_output would also need to be stored as well since it is used to type convert results. 😓

[albi3ro]

Note that while the main example was a dictinary, we should be able to handle any pytree of measurements. So another example might be:

return ([mp1, mp2], [mp3, mp4], {"a": mp5, "b": [mp6, mp7]})

as a convoluted example.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.46%. Comparing base (828c1ec) to head (3f4d8b9).

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #6378   +/-   ##
=======================================
  Coverage   99.46%   99.46%           
=======================================
  Files         454      454           
  Lines       42649    42657    +8     
=======================================
+ Hits        42420    42428    +8     
  Misses        229      229           

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

---

🚨 Try these New Features:

• Flaky Tests Detection - Detect and resolve failed and flaky tests

[albi3ro]

Also just making a comment here. If we are having obscure edge cases that are difficult to fit into nice logic here and are causing problems, we can consider making a breaking change, if that means more consistent, logical behavior in the QNode.

[albi3ro]

❤️