[Return-types #3] QNode integration, enable and disable functions

[rmoyard]

Context:

Third PR of the new return type project.

Description of the Change:

We add the QNode integration of the new return types.

The changes are made separately from the main pipeline. It can be activated by qml.enable_return()
, deactivated by qml.disable_return() and we can return mode by qml.active_return.

Example
Forward pass:

dev = qml.device("default.qubit", wires=2)

def circuit(x):
        qml.Hadamard(wires=[0])
        qml.CRX(x, wires=[0, 1])
        return (qml.probs(wires=[0]), qml.vn_entropy(wires=[0]), qml.probs(wires=0), qml.expval(wires=1))

qnode = qml.QNode(circuit, dev)
res = qnode(0.5)

>>> res
(tensor([0.5, 0.5], requires_grad=True), tensor(0.08014815, requires_grad=True), tensor([0.5, 0.5], requires_grad=True), tensor(0.93879128, requires_grad=True))

Backpropagation:

import jax

dev = qml.device("default.qubit", wires=2)
qml.enable_return()

@qml.qnode(dev, interface="jax")
def circuit(a):
    qml.RX(a[0], wires=0)
    qml.CNOT(wires=(0, 1))
    qml.RY(a[1], wires=1)
    qml.RZ(a[2], wires=1)
    return qml.expval(qml.PauliZ(wires=0)), qml.probs(wires=[0, 1]), qml.vn_entropy(wires=1)

x = jax.numpy.array([0.1, 0.2, 0.3])

res = jax.jacobian(circuit)(x)

>>> res

(DeviceArray([-9.9833414e-02, -7.4505806e-09, -3.9932679e-10], dtype=float32), 
DeviceArray([[-4.9419206e-02, -9.9086545e-02,  3.4938008e-09],
             [-4.9750542e-04,  9.9086538e-02,  1.2768372e-10],
             [ 4.9750548e-04,  2.4812977e-04,  4.8371929e-13],
             [ 4.9419202e-02, -2.4812980e-04,  2.6696912e-11]],            dtype=float32), 
DeviceArray([ 2.9899091e-01, -4.4703484e-08,  9.5104014e-10], dtype=float32))

[antalszava]

@rmoyard great push with this! 💪 🔥

No concrete blocker, just curious about some things.

[antalszava]

Looks good! 💪 🥇 Massive push here, great to have all these tests.

I'm wondering about the note on skipping some sampling tests.

The PR became quite big eventually, it could help to separate some conceptual parts (e.g., have backrop be tested in a separate PR) in other such cases.

[antalszava]

[sc-24246]