Fix jacobian shape in VJP for measurements with shape and batched inputs

doc/releases/changelog-dev.md

@@ -70,6 +70,9 @@
 * `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)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):

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

tests/gradients/core/test_vjp.py

@@ -676,3 +676,19 @@ def test_reduction_extend(self):
         # 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 if jacobian gets calculated when inputs are batched, multiple
+        trainable parameters are used and the measurement has a shape (probs)"""
+
+        @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)