[BUG] Zero-valued JVPs of gradient-free tapes with shot vectors have the wrong shape

[dwierichs]

Expected behavior

Found by @astralcai

Gradient-free tapes with shot vectors can be differentiated, giving a zero with the right output shape.

Actual behavior

The output shape is off.

Additional information

No response

Source code

import jax
import jax.numpy as jnp
import pennylane as qml

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

def cost(theta):
    tape1 = qml.tape.QuantumScript([qml.RX(0.5, wires=0)], [qml.expval(qml.Z(0))], shots=[20, 20])
    tape2 = qml.tape.QuantumScript([qml.RY(theta, wires=0)], [qml.expval(qml.Z(0))], shots=[20, 20])
    res1, res2 = qml.execute([tape1, tape2], dev, gradient_fn=qml.gradients.param_shift)
    return res1, res2

params = jnp.array(0.5)
jax.jacobian(cost)(params)

Tracebacks

TypeError: Custom JVP rule _execute_and_compute_jvp for function _execute_wrapper must produce primal and tangent outputs with equal container (pytree) structures, but got PyTreeDef(((*, *), (*, *))) and PyTreeDef((*, (*, *))) respectively.

System information

pl dev

Existing GitHub issues

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