[Return-types #5] Finite-difference update

[rmoyard]

Context:

Update the finite-diff gradient transform to reflect the new return types changes. It concerns the return for multiple measurements.

Description of the Change:

Instead of returning a single ragged array containing the results for multiple measurements, we consider that multiple measurements returns a tuple of measurements (m0, m1, ...). Therefore the derivative will have the following form: (m0:(p0, p1, ...), m1:(p0, p1, ...), ...) where each parameter tuple corresponds to a specific measurement.

This PR updates finite_diff gradient transform.

Example

dev = qml.device("default.qubit", wires=2)
x = 0.543
y = -0.654

with qml.tape.QuantumTape() as tape:
     qml.RX(x, wires=[0])
     qml.RY(y, wires=[1])
     qml.CNOT(wires=[0, 1])
     qml.expval(qml.PauliZ(0))
     qml.probs(wires=[0, 1])

tapes, fn = finite_diff_new(tape, approx_order=2, strategy="center")
res = fn(dev.batch_execute_new(tapes))

((array(-0.5167068), array(5.55111512e-10)), (array([-0.23169865, -0.02665475, 0.02665475, 0.23169865]), array([ 0.28230648, -0.28230648, -0.02187647, 0.02187647])))

where we obtain the same results as with Jax and backpropagation:

import jax

x = jax.numpy.array(0.543)
y = jax.numpy.array(-0.654)

qml.enable_return()

@qml.qnode(device=dev, interface="jax")
def circuit(x, y):
    qml.RX(x, wires=[0])
    qml.RY(y, wires=[1])
    qml.CNOT(wires=[0, 1])
    return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])

print(jax.jacobian(circuit, argnums=[0, 1])(x, y))

((DeviceArray(-0.5167068, dtype=float32, weak_type=True), DeviceArray(-2.9802322e-08, dtype=float32, weak_type=True)), (DeviceArray([-0.23169865, -0.02665475, 0.02665475, 0.23169865], dtype=float32, weak_type=True), DeviceArray([ 0.28230646, -0.2823065 , -0.02187647, 0.02187647], dtype=float32, weak_type=True)))

Benefits:

Clear and intuitive return system for multiple measurements.

TODO

• QNode integration

[sc-25812]

[github-actions]

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md with:

• A one-to-two sentence description of the change. You may include a small working example for new features.
• A link back to this PR.
• Your name (or GitHub username) in the contributors section.

[codecov]

Codecov Report

Merging #2966 (1f470a9) into master (1aa7df3) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff            @@
##           master    #2966    +/-   ##
========================================
  Coverage   99.67%   99.67%            
========================================
  Files         273      273            
  Lines       23349    23450   +101     
========================================
+ Hits        23273    23374   +101     
  Misses         76       76            

Impacted Files	Coverage Δ
pennylane/_qubit_device.py	99.65% <100.00%> (+<0.01%)	⬆️
pennylane/gradients/finite_difference.py	100.00% <100.00%> (ø)
pennylane/transforms/batch_transform.py	100.00% <100.00%> (ø)

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[AlbertMitjans]

Should be use qml.enable_return() in tests?

[antalszava]

Some logic may be specific only to qubits. Maybe worth adding a TODO there.

[antalszava]

Overall looks great! 💯 Approving with the condition that the CV case is checked and if found to be relevant it's resolved either via a TODO or via a test case.

[eddddddy]

Great work @rmoyard 💯 ! I still have a couple of suggestions on how we can clean up the code a bit. In particular, it looks like we can reuse a lot of the code from the new parameter-shift transform here.

I'm giving my approval now but this is conditioned on these comments being addressed.

[antalszava]

👏