Repository: PennyLaneAI/pennylane

Issue:
[BUG] `qml.gradients.param_shift` fails for QNode with multiple array arguments (autograd)
### Expected behavior

Running `qml.gradients.param_shift` on a QNode should return the hybrid Jacobian of the QNode with respect to its trainable arguments.

### Actual behavior

An error is thrown. See the minimal working example and traceback below.

### Additional information

### Source of the Bug

I suspect it's due to shape of the classical Jacobian, or the way the classical and quantum Jacobians are contracted together.
The classical Jacobian is computed on the function of QNode inputs to tape parameters (i.e. gate arguments).
The quantum Jacobian is computed on the function of tape parameters to QNode outputs.

In the example given below, QNode arguments have a combined shape `(2, 1)`, tape parameters shape `(3)`, and QNode outputs shape `(4)`.
Computing the parameter shift Jacobian has the following execution path (from bottom to top):

**_autograd_** -> Computes a classical Jacobian with respect to each QNode argument, in this case `2` Jacobians with shape `(3,1)`.

**_qml.jacobian_** -> Attempts to stack the classical Jacobians into a single ndarray and transposes it, resulting in shape `(1,3,2)`. This shape seems off to me (reversed shape of one QNode argument, # gate arguments, # QNode arguments).
```py
return np.stack(jacobians).T
```
Where stacking is not possible (due to QNode arguments having different shapes), a tuple of the individual Jacobians would be returned (in this case it would have a shape of `(2,3,1)`).
```py
return tuple(jacobians)
```

**_qml.gradients.param_shift_** -> Computes the quantum Jacobian of shape `(4,3)`,

**_qml.gradients.gradient_transform_** -> Since the stacking was successful, attempts to compute the hybrid Jacobian via _`cjac.T x(-1,-1) qjac`_. This results in an error due to the mismatched dimensions `(1,3,2)[-1]` and `(4,3)[-1]`.
```py
jac = qml.math.squeeze(qml.math.tensordot(qml.math.T(cjac), qjac, [[-1], [-1]]))
return qml.math.T(jac)
```
This error does not occur in the case Jacobians are returned as tuples. Here, each `(3,1)` classical Jacobian would be contracted with the `(4,3)` quantum Jacobian via _`cjac x(0,-1) qjac`_.
```py
jacs = [
    qml.math.squeeze(qml.math.tensordot(c, qjac, [[0], [-1]]))
    for c in cjac
    if c is not None
]
return jacs
```

### Fix Discussion

One could fix the gradient_transform to account for this, but to me the shape returned by qml.jacobian seems off.
Something like (# gate arguments, QNode inputs shape) `(3,2,1)` would make sense, as the dimensions of the multiple QNode arguments would be stacked together in the inner-most dimensions. This matches what one would expect if the QNode arguments were passed as a single ndarray instead. It is also what the gradient_transform seems to expect.

Alternatively (# QNode arguments, # gate arguments, shape of one QNode argument) `(2,3,1)` might also make sense, as it matches the case were the Jacobians can't be stacked due to different QNode argument dimensions.

Note that the current behavior is (reversed shape of one QNode argument, # gate arguments, # QNode arguments) `(1,3,2)`. I'm not sure whether there is an argument in favor of this shape. Some of our tests seem to expect it.

### Source code

```shell
import pennylane as qml
from pennylane import numpy as np

dev = qml.device("default.qubit", wires = 2)

@qml.qnode(dev, max_diff=2)
def circuit(x, y):
    qml.RX(x[0], wires=0)
    qml.RY(y[0], wires=0)
    qml.RZ(x[0], wires=0)
    return qml.probs(wires=[0,1])

x = np.array([0.1], requires_grad=True)
y = np.array([0.2], requires_grad=True)

qml.gradients.param_shift(circuit)(x,y)
```


### Tracebacks

```shell
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_6871/4095522486.py in <module>
     15 z = np.array([0.1], requires_grad=True)
     16 
---> 17 qml.gradients.param_shift(circuit)(x,y)

~/penny/pennylane/gradients/gradient_transform.py in jacobian_wrapper(*args, **kwargs)
    146 
    147             # Classical processing of a single argument is present. Return qjac @ cjac.
--> 148             jac = qml.math.squeeze(qml.math.tensordot(qml.math.T(cjac), qjac, [[-1], [-1]]))
    149             return qml.math.T(jac)
    150 

~/penny/pennylane/math/multi_dispatch.py in tensordot(tensor1, tensor2, axes)
    280     """
    281     interface = _multi_dispatch([tensor1, tensor2])
--> 282     return np.tensordot(tensor1, tensor2, axes=axes, like=interface)
    283 
    284 

~/.local/lib/python3.8/site-packages/autoray/autoray.py in do(fn, like, *args, **kwargs)
     82         backend = infer_backend(like)
     83 
---> 84     return get_lib_fn(backend, fn)(*args, **kwargs)
     85 
     86 

~/penny/pennylane/numpy/wrapper.py in _wrapped(*args, **kwargs)
     73 
     74         # evaluate the original object
---> 75         res = obj(*args, **kwargs)
     76 
     77         if isinstance(res, _np.ndarray):

~/.local/lib/python3.8/site-packages/autograd/tracer.py in f_wrapped(*args, **kwargs)
     46             return new_box(ans, trace, node)
     47         else:
---> 48             return f_raw(*args, **kwargs)
     49     f_wrapped.fun = f_raw
     50     f_wrapped._is_autograd_primitive = True

<__array_function__ internals> in tensordot(*args, **kwargs)

~/.local/lib/python3.8/site-packages/numpy/core/numeric.py in tensordot(a, b, axes)
   1108                 axes_b[k] += ndb
   1109     if not equal:
-> 1110         raise ValueError("shape-mismatch for sum")
   1111 
   1112     # Move the axes to sum over to the end of "a"

ValueError: shape-mismatch for sum
```


### System information

```shell
-
```


### 

- [X] I have searched exisisting GitHub issues to make sure the issue does not already exist.


Discussion:
Thanks @dime10 for the extensive bug report and tracing this back so precisely.
I agree that the output shape of `qml.jacobian` is unintuitive. @josh146 will have more to say about the reasons why the current behavious is the way it is, but I'll add my five cents on the fix discussion:

**Behaviour of example in other interfaces**
- In JAX, the default output will be only the Jacobian with respect to the 0th argument, as always. When requesting the full classical Jacobian via `argnum=[0,1]`, we obtain a 2-tuple of arrays with total shape `(2, 3, 1)` if we do not discriminate between array and tuple dimensions.
- In TF, the behaviour is the same as for JAX
- In Torch, the classical Jacobian is `(2, 3, 1)` as well. (On a side note: `qml.math.shape` does not detect this shape but reports `(2,)` whereas it reports `(2, 3, 1)` for TF and JAX)

I think therefore there is a strong case to adopt the second suggestion in the Fix discussion above to match the behaviour in all other interfaces. In particular this leads to the (to me personally) intuitive slicing behaviour:
`classical_jacobian(circuit)(*qnode_args)[a:b] == classical_jacobian(circuit, argnum=list(range(a,b)))(*qnode_args)`

I agree with the discussion above.

> @josh146 will have more to say about the reasons why the current behavious is the way it is

There is no good reason as to why it is like this! If I had to guess, I would say that in the early days we were writing QNodes like this:

```python
@qml.qnode(dev)
def circuit(x, y, z):
    op(x)
    op(y)
    op(z)
```
That is, we were using a `circuit(*weights)` pattern. However, this does not scale!

Stepping back a bit, the most intuitive behaviour is such that the return type of `qml.jacobian` should be determined by the signature of the function, rather than the shape of the inputs.

e.g., as a user, I guess I expect

```python
qml.jacobian(circuit)(x, y)
```

(where both `x` and `y` requires grad) to match JAX's `argnums=[0, 1]` and return a tuple `(jac_x, jac_y)`, whereas I would expect

```python
qml.jacobian(circuit)(weights)
```

to return a tensor.

In other words, I think it is probably least surprising if the output of `qml.jacobian` matches the function signature.

Whereas, at the moment, `qml.jacobian(circuit)(x, y)` could return either a tuple or a tensor, it is impossible to say without knowing the shapes of `x`, `y`, and the output shape of circuit!

Resolving patch:
diff --git a/pennylane/gradients/gradient_transform.py b/pennylane/gradients/gradient_transform.py
--- a/pennylane/gradients/gradient_transform.py
+++ b/pennylane/gradients/gradient_transform.py
@@ -14,6 +14,8 @@
 """This module contains utilities for defining custom gradient transforms,
 including a decorator for specifying gradient expansions."""
 # pylint: disable=too-few-public-methods
+import warnings
+
 import pennylane as qml
 from pennylane.transforms.tape_expand import expand_invalid_trainable
 
@@ -108,6 +110,27 @@ def __init__(
         self.hybrid = hybrid
         super().__init__(transform_fn, expand_fn=expand_fn, differentiable=differentiable)
 
+    @staticmethod
+    def _jacobian_trainable_args(args, interface):
+        """Return the indices of QNode arguments for which a Jacobian was
+        computed by `qml.transforms.classical_jacobian` with argnum=None.
+        """
+        trainable_args = []
+
+        if interface == "autograd":
+            for idx, arg in enumerate(args):
+                # TODO: make default False once this change is done in qml.jacobian
+                if getattr(arg, "requires_grad", True):
+                    trainable_args.append(idx)
+
+        elif interface == "jax":
+            trainable_args = [0]
+
+        # Torch and Tensorflow interfaces are not considered since `classical_jacobian`
+        # always returns a tuple for them, thus not invoking this function.
+
+        return trainable_args
+
     def default_qnode_wrapper(self, qnode, targs, tkwargs):
         # Here, we overwrite the QNode execution wrapper in order
         # to take into account that classical processing may be present
@@ -144,8 +167,42 @@ def jacobian_wrapper(*args, **kwargs):
                 # is present inside the QNode.
                 return qjac
 
-            # Classical processing of a single argument is present. Return qjac @ cjac.
-            jac = qml.math.squeeze(qml.math.tensordot(qml.math.T(cjac), qjac, [[-1], [-1]]))
-            return qml.math.T(jac)
+            # Classical processing present of either:
+            #   a) a single argument or
+            #   b) multiple arguments of the same shape
+            # The shape of the classical jacobian returned by qml.jacobian (invoked by
+            # classical_jacobian with autograd) depends on the scenario:
+            #   a) (# gate args, qnode arg shape)
+            #   b) (reverse qnode arg shape, # gate args, # qnode args)
+            # It then needs to be contracted with the quantum jacobian of shape:
+            #      (qnode output shape, # gate args)
+            # The result should have shape:
+            #   a) (qnode ouput shape, qnode arg shape)
+            #   b) (reverse qnode arg shape, qnode output shape, # qnode args)
+            num_gate_args = qml.math.shape(qjac)[-1]
+            # Consider only QNode arguments regarded as trainable by the interfaces.
+            trainable_args_idx = self._jacobian_trainable_args(args, qnode.interface)
+            # Since all arguments have the same shape, obtain shape from the first trainable arg
+            qnode_arg_shape = qml.math.shape(args[trainable_args_idx[0]])
+            num_qnode_args = len(trainable_args_idx)
+
+            if qml.math.shape(cjac) == (num_gate_args, *qnode_arg_shape):
+                # single QNode argument
+                jac = qml.math.tensordot(qjac, cjac, [[-1], [0]])
+            elif qml.math.shape(cjac) == (*qnode_arg_shape[::-1], num_gate_args, num_qnode_args):
+                # multiple QNode arguments with stacking
+                cjac = qml.math.swapaxes(cjac, -1, -2)
+                jac = qml.math.tensordot(cjac, qjac, [[-1], [-1]])
+                jac = qml.math.moveaxis(jac, -1, len(qnode_arg_shape))
+            else:  # pragma: no cover
+                jac = ()
+                warnings.warn(
+                    "Unexpected classical Jacobian encoutered, could not compute the hybrid "
+                    "Jacobian of the QNode. You can still attempt to obtain the quantum Jacobian "
+                    "with the `hybrid=False` parameter.",
+                    UserWarning,
+                )
+
+            return qml.math.squeeze(jac)
 
         return jacobian_wrapper


Test patch:
diff --git a/tests/gradients/test_gradient_transform.py b/tests/gradients/test_gradient_transform.py
--- a/tests/gradients/test_gradient_transform.py
+++ b/tests/gradients/test_gradient_transform.py
@@ -135,6 +135,46 @@ def circuit(weights):
         expected = qml.jacobian(circuit)(w)
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+    def test_multiple_tensor_arguments(self, tol):
+        """Test that a gradient transform acts on QNodes
+        correctly when multiple tensor QNode arguments are present"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(x[0, 0], wires=0)
+            qml.RY(y[0, 0], wires=0)
+            qml.RZ(x[1, 0], wires=0)
+            return qml.probs(wires=[0, 1])
+
+        x = np.array([[0.1], [0.2]], requires_grad=True)
+        y = np.array([[0.2], [0.3]], requires_grad=True)
+
+        expected = qml.jacobian(circuit)(x, y)
+        res = qml.gradients.param_shift(circuit)(x, y)
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_first_non_trainable_argument(self, tol):
+        """Test that a gradient transform acts on QNodes
+        correctly when the first argument is non-trainable"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(y[0], wires=0)
+            qml.RY(y[0], wires=0)
+            qml.RZ(y[1], wires=0)
+            return qml.probs(wires=[0, 1])
+
+        x = np.array([0.1], requires_grad=False)
+        y = np.array([0.2, 0.3], requires_grad=True)
+
+        expected = qml.jacobian(circuit)(x, y)
+        res = qml.gradients.param_shift(circuit)(x, y)
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
     def test_classical_processing_arguments(self, mocker, tol):
         """Test that a gradient transform acts on QNodes
         correctly when the QNode arguments are classically processed"""