[BUG] Some templates not properly differentiated with parameter shift on the legacy device

[astralcai]

Expected behavior

ControlledSequence, Reflection, Qubitization, and AmplitudeAmplification can be differentiated with both default.qubit and default.qubit.legacy.

Actual behavior

Consider the following circuit:

def circuit(x):
    qml.PauliX(2)
    qml.ControlledSequence(qml.RX(x, wires=3), control=[0, 1, 2])
    return qml.probs(wires=range(4))

With default.qubit, it produces the expected results when taking the gradient:

>>> dev = qml.device("default.qubit")
>>> qnode = qml.QNode(circuit, dev, interface="autograd", diff_method="parameter-shift")
>>> x = qml.numpy.array(0.5, requires_grad=True)
>>> qml.jacobian(qnode)(x)
array([ 0.        ,  0.        , -0.23971277,  0.23971277,  0.        ,
        0.        ,  0.        ,  0.        ,  0.        ,  0.        ,
        0.        ,  0.        ,  0.        ,  0.        ,  0.        ,
        0.        ])

But with default.qubit.legacy,

>>> dev = qml.device("default.qubit.legacy", wires=4)
>>> qnode = qml.QNode(circuit, dev, interface="autograd", diff_method="parameter-shift")
>>> qml.jacobian(qnode)(x)
TypeError: 'NoneType' object is not subscriptable

This error can be fixed quite easily by adding

# check the grad_recipe validity
if self.grad_recipe is None:
    # Make sure grad_recipe is an iterable of correct length instead of None
    self.grad_recipe = [None] * self.num_params

to the end of ControlledSequence.__init__(), but when shots are added:

>>> dev = qml.device("default.qubit.legacy", wires=4, shots=50000)
>>> qnode = qml.QNode(circuit, dev, interface="autograd", diff_method="parameter-shift")
>>> qml.jacobian(qnode)(x)
array([     0.,      0., -26200.,  26200.,      0.,      0.,      0.,
            0.,      0.,      0.,      0.,      0.,      0.,      0.,
            0.,      0.])

We get non-sensible results. This issue also appears for Reflection, Qubitization, and AmplitudeAmplification

Additional information

The issue is due to default.qubit has a preprocess that happens before the ml boundary, decomposing the templates to lower level gates, but with default.qubit.legacy, expansion happens in inner_execute, which means that the ml boundary is dealing with the template operation itself, and they themselves are not properly differentiable.

This bug is blocking #5791

Source code

No response

Tracebacks

No response

System information

Name: PennyLane
Version: 0.37.0.dev0
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author:
Author-email:
License: Apache License 2.0
Location: /Users/astral.cai/Workspace/pennylane/venv/lib/python3.9/site-packages
Editable project location: /Users/astral.cai/Workspace/pennylane
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane_Lightning

Platform info:           macOS-14.5-arm64-arm-64bit
Python version:          3.9.19
Numpy version:           1.26.4
Scipy version:           1.11.4
Installed devices:
- default.clifford (PennyLane-0.37.0.dev0)
- default.gaussian (PennyLane-0.37.0.dev0)
- default.mixed (PennyLane-0.37.0.dev0)
- default.qubit (PennyLane-0.37.0.dev0)
- default.qubit.autograd (PennyLane-0.37.0.dev0)
- default.qubit.jax (PennyLane-0.37.0.dev0)
- default.qubit.legacy (PennyLane-0.37.0.dev0)
- default.qubit.tf (PennyLane-0.37.0.dev0)
- default.qubit.torch (PennyLane-0.37.0.dev0)
- default.qutrit (PennyLane-0.37.0.dev0)
- default.qutrit.mixed (PennyLane-0.37.0.dev0)
- null.qubit (PennyLane-0.37.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.36.0)

Existing GitHub issues

• I have searched existing GitHub issues to make sure the issue does not already exist.

[dwierichs]

but when shots are added [...] We get non-sensible results.

This is a result of finite differences. Some operation in the decomposition will have grad_method="F" ControlledSequence with the fixed grad_recipe is not caught by _param_shift_stopping_condition and thus not decomposed. It then falls back to finite diff instead of parameter shift. The result is a numerically unstable derivative that can be resolved with float precision (shots=None) but not with shots!=None precision (for all shot counts that fit on a laptop).