Add support for qml.Snapshot

[joeycarter]

PennyLane added the qml.Snapshot operation in Release 0.22.0, which saves the internal state of devices at arbitrary points of execution. For example:

import pennylane as qml

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

@qml.qnode(dev)
def circuit():
    wires = list(range(NUM_QUBITS))

    qml.Snapshot("Initial state")

    for wire in wires:
        qml.Hadamard(wires=wire)

    qml.Snapshot("After applying Hadamard gates")

    return qml.probs()


results = qml.snapshots(circuit)()

for k, result in results.items():
    print(f"{k:<30}: {result}")

Output:

Initial state                 : [1.+0.j 0.+0.j 0.+0.j 0.+0.j]
After applying Hadamard gates : [0.5+0.j 0.5+0.j 0.5+0.j 0.5+0.j]
execution_results             : [0.25 0.25 0.25 0.25]

Adding Snapshot support to Catalyst would possibly be a nice feature to have for some users.

Note that currently if we try wrapping circuit() with a @qjit decorator, calling qml.snapshots(circuit)() will raise a TransformError exception since we cannot apply a PennyLane transform to a qjit-compiled function. This also requires changing the device to "lightning.qubit" in the example above.