Repository: PennyLaneAI/pennylane (python)

Issue:
[BUG] Shape mismatch occurs in gradient calculation when using an indexing operation with batched inputs and autograd 
### Expected behavior

When having batched inputs for a QNode (in the example script below: input of shape `(5, 1)`), I want to calculate the probabilities of the outcomes of the circuit and then use the first probability of outcome "0" for each input.
That is, I obtain a output of shape `(5, 2)` from the circuit, where the first dimension (5) represents the batch dimension, and the second dimension (2) represents the respective probabilities for either "0", or "1".

The probability for all "0"-es, can be obtained by indexing, i.e. when the outcomes of the QNode are `prediction` (shape `(5, 2)`), the probabilities for "0"-es would be `prediction[:,0]` (shape `(5,)`).

The "forward" pass works as expected. One would also expect to obtain gradients, wrt. the parameters, when e.g. calling `qml.AdamOptimizer().step_and_cost`.

### Actual behavior

Computing the gradients fails for batchsizes > 1 with the Traceback below.

### Additional information

In the non-working example script, either setting the `BATCH_SIZE = 1`, or using the torch interface with a torch optimiser, e.g. as in the following code sample, does provide the expected behaviour. Therefore, I assume the bug might be related to either the gradient calculation with autograd, or the `qml.AdamOptimizer` used in the example.

``` 
## Working example with torch
import pennylane as qml
import torch

BATCH_SIZE = 5

def _circuit(w, x):
    qml.RX(w[0], wires=0)
    qml.RZ(w[1], wires=0)
    qml.RX(x[:, 0], wires=0)
    return qml.probs(0)

circuit = qml.QNode(
    _circuit,
    qml.device("default.qubit", shots=1024, wires=1),
)

x = torch.rand(BATCH_SIZE, 1)
y = torch.rand(BATCH_SIZE)
w = torch.rand(2, requires_grad=True)

opt = torch.optim.Adam([w], lr=0.1)

def mse(prediction, target):
    return torch.mean((prediction - target) ** 2)

def cost(params, target, **kwargs):
    prediction = circuit(w=params, **kwargs)
    print(f"Circuit prediction: {prediction}")
    if len(prediction.shape) == 1:
        prediction = prediction[0]
    else:
        prediction = prediction[:,0]
    print(f"Processed circuit prediction: {prediction}")
    return mse(prediction, target)

print(f"Weights before optimising: {w}")
cost_val = cost(w, target=y, x=x)

opt.step()

print(f"Weights after optimising: {w}")
print(f"Cost: {cost_val}")
```

#### Output
```
Weights before optimising: tensor([0.2912, 0.7003], requires_grad=True)
Circuit prediction: tensor([[0.8311, 0.1689],
        [0.9014, 0.0986],
        [0.9570, 0.0430],
        [0.9004, 0.0996],
        [0.8955, 0.1045]], dtype=torch.float64, grad_fn=<ExecuteTapesBackward>)
Processed circuit prediction: tensor([0.8311, 0.9014, 0.9570, 0.9004, 0.8955], dtype=torch.float64,
       grad_fn=<SelectBackward0>)
Weights after optimising: tensor([0.2912, 0.7003], requires_grad=True)
Cost: 0.23868222326400854
```

### Source code

```shell
## Non-working example with autograd
import pennylane as qml
import pennylane.numpy as pnp
import numpy as np

BATCH_SIZE = 5

def _circuit(w, x):
    qml.RX(w[0], wires=0)
    qml.RZ(w[1], wires=0)
    qml.RX(x[:, 0], wires=0)
    return qml.probs(0)

circuit = qml.QNode(
    _circuit,
    qml.device("default.qubit", shots=1024, wires=1),
)

x = np.random.rand(BATCH_SIZE, 1)
y = np.random.rand(BATCH_SIZE)
w = pnp.random.rand(2, requires_grad=True)

opt = qml.AdamOptimizer(0.1)

def mse(prediction, target):
    return pnp.mean((prediction - target) ** 2)

def cost(params, target, **kwargs):
    prediction = circuit(w=params, **kwargs)
    print(f"Circuit prediction: {prediction}")
    if len(prediction.shape) == 1:
        prediction = prediction[0]
    else:
        prediction = prediction[:,0]
    print(f"Processed circuit prediction: {prediction}")
    return mse(prediction, target)

print(f"Weights before optimising: {w}")
w, cost_val = opt.step_and_cost(
    cost,
    w,
    target=y,
    x=x,
)

print(f"Weights after optimising: {w}")
print(f"Cost: {cost_val}")
```


### Tracebacks

```shell
Weights before optimising: [0.52479668 0.21144672]
Circuit prediction: Autograd ArrayBox with value [[0.78417969 0.21582031]
 [0.61816406 0.38183594]
 [0.74902344 0.25097656]
 [0.79589844 0.20410156]
 [0.6328125  0.3671875 ]]
Processed circuit prediction: Autograd ArrayBox with value [0.78417969 0.61816406 0.74902344 0.79589844 0.6328125 ]
Traceback (most recent call last):
  File "<path>/.venv/lib/python3.11/site-packages/pennylane/gradients/vjp.py", line 156, in compute_vjp_single
    res = jac @ dy_row
          ~~~~^~~~~~~~
ValueError: matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 10 is different from 2)

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<path>/test.py", line 39, in <module>
    w, cost_val = opt.step_and_cost(
                  ^^^^^^^^^^^^^^^^^^
  File "<path>/.venv/lib/python3.11/site-packages/pennylane/optimize/gradient_descent.py", line 64, in step_and_cost
    g, forward = self.compute_grad(objective_fn, args, kwargs, grad_fn=grad_fn)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<path>/.venv/lib/python3.11/site-packages/pennylane/optimize/gradient_descent.py", line 122, in compute_grad
    grad = g(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "<path>/.venv/lib/python3.11/site-packages/pennylane/_grad.py", line 166, in __call__
    grad_value, ans = grad_fn(*args, **kwargs)  # pylint: disable=not-callable
                      ^^^^^^^^^^^^^^^^^^^^^^^^
  File "<path>/.venv/lib/python3.11/site-packages/autograd/wrap_util.py", line 20, in nary_f
    return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<path>/.venv/lib/python3.11/site-packages/pennylane/_grad.py", line 192, in _grad_with_forward
    grad_value = vjp(vspace(ans).ones())
                 ^^^^^^^^^^^^^^^^^^^^^^^
  File "<path>/.venv/lib/python3.11/site-packages/autograd/core.py", line 14, in vjp
    def vjp(g): return backward_pass(g, end_node)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<path>/.venv/lib/python3.11/site-packages/autograd/core.py", line 21, in backward_pass
    ingrads = node.vjp(outgrad[0])
              ^^^^^^^^^^^^^^^^^^^^
  File "<path>/.venv/lib/python3.11/site-packages/autograd/core.py", line 67, in <lambda>
    return lambda g: (vjp(g),)
                      ^^^^^^
  File "<path>/.venv/lib/python3.11/site-packages/pennylane/workflow/interfaces/autograd.py", line 199, in grad_fn
    vjps = jpc.compute_vjp(tapes, dy)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<path>/.venv/lib/python3.11/site-packages/pennylane/workflow/jacobian_products.py", line 297, in compute_vjp
    return _compute_vjps(jacs, dy, tapes)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<path>/.venv/lib/python3.11/site-packages/pennylane/workflow/jacobian_products.py", line 46, in _compute_vjps
    vjps.append(f[multi](dy, jac))
                ^^^^^^^^^^^^^^^^^
  File "<path>/.venv/lib/python3.11/site-packages/pennylane/gradients/vjp.py", line 158, in compute_vjp_single
    res = qml.math.tensordot(jac, dy_row, [[1], [0]])
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<path>/.venv/lib/python3.11/site-packages/pennylane/math/multi_dispatch.py", line 152, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "<path>/.venv/lib/python3.11/site-packages/pennylane/math/multi_dispatch.py", line 403, in tensordot
    return np.tensordot(tensor1, tensor2, axes=axes, like=like)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<path>/.venv/lib/python3.11/site-packages/autoray/autoray.py", line 81, in do
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "<path>/.venv/lib/python3.11/site-packages/numpy/core/numeric.py", line 1099, in tensordot
    raise ValueError("shape-mismatch for sum")
ValueError: shape-mismatch for sum
```


### System information

```shell
Name: PennyLane
Version: 0.37.0
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: <path>/.venv/lib/python3.11/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane_Lightning

Platform info:           Linux-6.1.90-1-MANJARO-x86_64-with-glibc2.39
Python version:          3.11.5
Numpy version:           1.26.4
Scipy version:           1.14.0
Installed devices:
- default.clifford (PennyLane-0.37.0)
- default.gaussian (PennyLane-0.37.0)
- default.mixed (PennyLane-0.37.0)
- default.qubit (PennyLane-0.37.0)
- default.qubit.autograd (PennyLane-0.37.0)
- default.qubit.jax (PennyLane-0.37.0)
- default.qubit.legacy (PennyLane-0.37.0)
- default.qubit.tf (PennyLane-0.37.0)
- default.qubit.torch (PennyLane-0.37.0)
- default.qutrit (PennyLane-0.37.0)
- default.qutrit.mixed (PennyLane-0.37.0)
- default.tensor (PennyLane-0.37.0)
- null.qubit (PennyLane-0.37.0)
- lightning.qubit (PennyLane_Lightning-0.37.0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

Pull request description:
Fix jacobian shape in VJP for measurements with shape and batched inputs
**Context:**
This PR should fix a bug related to computing the gradient of a circuit which supports:
1. batching in non-trainable data
2. Measurement with a shape (i.e. probs)
3. More than one trainable parameter in the circuit

Localised example that fails, without the fix:
```
@qml.qnode(qml.device('default.qubit'), diff_method="parameter-shift")
def circuit(x, data):
    qml.RX(x[0], 0)
    qml.RX(x[1], 0)
    qml.RY(data, 0)
    return qml.probs(wires=0)

x = qml.numpy.array([0.5, 0.8], requires_grad=True)
data = qml.numpy.array([1.2, 2.3, 3.4], requires_grad=False)
circuit(x, data)
qml.jacobian(circuit)(x, data)
```
Reshaping the jacobian, when computing the VJP does resolve the issue.

**Description of the Change:**
Jacobian is reshaped according to the shape of dy when computing the VJP for measurements with dimension.

**Benefits:**
One less bug :)

**Possible Drawbacks:**
None

**Related GitHub Issues:**
This PR fixes #5979


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index b403126a9..3af8b773b 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -259,10 +259,12 @@
 * `qml.AmplitudeEmbedding` has better support for features using low precision integer data types.
 [(#5969)](https://github.com/PennyLaneAI/pennylane/pull/5969)
 
+* Jacobian shape is fixed for measurements with dimension in `qml.gradients.vjp.compute_vjp_single`.
+[(5986)](https://github.com/PennyLaneAI/pennylane/pull/5986)
+
 * `qml.lie_closure` works with sums of Paulis.
   [(#6023)](https://github.com/PennyLaneAI/pennylane/pull/6023)
 
-
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
@@ -270,11 +272,12 @@ This release contains contributions from (in alphabetical order):
 Tarun Kumar Allamsetty,
 Guillermo Alonso,
 Utkarsh Azad,
+Gabriel Bottrill,
 Ahmed Darwish,
 Astral Cai,
 Yushao Chen,
-Gabriel Bottrill,
 Ahmed Darwish,
+Maja Franz,
 Lillian M. A. Frederiksen,
 Pietropaolo Frisoni,
 Emiliano Godinez,
diff --git a/pennylane/gradients/vjp.py b/pennylane/gradients/vjp.py
index 82e1fec59..650624b0c 100644
--- a/pennylane/gradients/vjp.py
+++ b/pennylane/gradients/vjp.py
@@ -151,7 +151,7 @@ def compute_vjp_single(dy, jac, num=None):
 
         # Single measurement with dimension e.g. probs
         else:
-            jac = qml.math.stack(jac)
+            jac = qml.math.reshape(qml.math.stack(jac), (-1, num))
             try:
                 res = jac @ dy_row
             except Exception:  # pylint: disable=broad-except


Test patch:
diff --git a/tests/gradients/core/test_vjp.py b/tests/gradients/core/test_vjp.py
index 5817563b0..ce972c46b 100644
--- a/tests/gradients/core/test_vjp.py
+++ b/tests/gradients/core/test_vjp.py
@@ -676,3 +676,24 @@ class TestBatchVJP:
         # list will correspond to a single input parameter of the combined
         # tapes.
         assert len(res) == sum(len(t.trainable_params) for t in tapes)
+
+    def test_batched_params_probs_jacobian(self):
+        """Test that the VJP gets calculated correctly when inputs are batched, multiple
+        trainable parameters are used and the measurement has a shape (probs)"""
+        data = np.array([1.2, 2.3, 3.4])
+        x0, x1 = 0.5, 0.8
+        ops = [qml.RX(x0, 0), qml.RX(x1, 0), qml.RY(data, 0)]
+        tape = qml.tape.QuantumScript(ops, [qml.probs(wires=0)], trainable_params=[0, 1])
+        dy = np.array([[0.6, -0.7], [0.2, -0.7], [-5.2, 0.6]])
+        v_tapes, fn = qml.gradients.batch_vjp([tape], [dy], qml.gradients.param_shift)
+
+        dev = qml.device("default.qubit")
+        vjp = fn(dev.execute(v_tapes))
+
+        # Analytically expected Jacobian and VJP
+        expected_jac = [-0.5 * np.cos(data) * np.sin(x0 + x1), 0.5 * np.cos(data) * np.sin(x0 + x1)]
+        expected_vjp = np.tensordot(expected_jac, dy, axes=[[0, 1], [1, 0]])
+        assert qml.math.shape(vjp) == (1, 2)  # num tapes, num trainable tape parameters
+        assert np.allclose(
+            vjp, expected_vjp
+        )  # Both parameters essentially feed into the same RX rotation