[BUG] StatePrep and its decomposition differ in derivative

[dwierichs]

Expected behavior

The derivative is (approximately) independent of decomposition and diff method

Actual behavior

The derivative between StatePrep and its decomposition differs.

Additional information

No response

Source code

dev = qml.device("default.qubit")

def circuit(x):
    qml.StatePrep(x, wires=[0, 1])
    return qml.probs(wires=[0, 1])

node_ps = qml.QNode(circuit, dev ,diff_method="finite-diff", h=0.05)
node_bp = qml.QNode(circuit, dev ,diff_method="backprop")

x = jnp.array([0.5, 0.5, 0.5, 0.5])
print(jax.jacobian(node_ps)(x))
[[ 0.49979169  0.          0.         -0.49979169]
 [ 0.          0.49979169 -0.49979169  0.        ]
 [ 0.         -0.49979169  0.49979169  0.        ]
 [-0.49979169  0.          0.          0.49979169]]
print(jax.jacobian(node_bp)(x))
[[1. 0. 0. 0.]
 [0. 1. 0. 0.]
 [0. 0. 1. 0.]
 [0. 0. 0. 1.]]

Tracebacks

No response

System information

pl dev

Existing GitHub issues

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

[trbromley]

@dwierichs, node_ps is actually finite-difference in the example. I get an error if I try parameter shift:

node_ps = qml.QNode(circuit, dev ,diff_method="parameter-shift")

Running your code with finite difference, my Jacobian is also different to yours:

[[ 0.7496871  -0.24989586 -0.24989586 -0.24989538]
 [-0.2498959   0.7496871  -0.24989533 -0.24989581]
 [-0.24989586 -0.24989538  0.749687   -0.24989578]
 [-0.24989533 -0.24989581 -0.24989581  0.74968696]]

So yes, this does look like a bug!

[dwierichs]

Indeed, the derivative changed for me as well. We fixed Mottonen in the meantime, which is used by StatePrep when decomposing it.

When assigning GlobalPhase.grad_method = None, the example also works with parameter-shift. (c.f. #5620 )