Repository: PennyLaneAI/pennylane

Issue:
[BUG] Shot vectors with jitting don't work for parametrized `QNode`s
### Expected behavior

Shot vectors work with jitting all the time.

### Actual behavior

If a qnode accepts parameters, jitting it with a shot vector raises an error. This happens when shots are set during device creation as well as when shots are specified on the qnode. The errors do not occur if the qnode doesn't have any parameters, even if a parametrized operation inside a qnode uses values from an outer context as its parameters.

### Additional information

_No response_

### Source code

```shell
>>> dev = qml.device("default.qubit")
>>> @qml.qnode(dev, interface="jax")
... def f(p):
...     qml.RX(p, 0)
...     qml.CNOT([0, 1])
...     return qml.expval(qml.PauliZ(0))
... 
>>> f_jitted = jax.jit(f, static_argnames=["shots"])
```


### Tracebacks

```shell
>>> f_jitted(jax.numpy.array(np.pi / 2), shots=(10, 10))
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/mudit.pandey/repos/pennylane/pennylane/qnode.py", line 1010, in __call__
    res = qml.execute(
  File "/Users/mudit.pandey/repos/pennylane/pennylane/interfaces/execution.py", line 736, in execute
    results = ml_boundary_execute(
  File "/Users/mudit.pandey/repos/pennylane/pennylane/interfaces/jax_jit.py", line 278, in execute
    return _execute_bwd(
  File "/Users/mudit.pandey/repos/pennylane/pennylane/interfaces/jax_jit.py", line 363, in _execute_bwd
    return execute_wrapper(params)
  File "/Users/mudit.pandey/repos/pennylane/pennylane/interfaces/jax_jit.py", line 324, in execute_wrapper
    return jax.pure_callback(wrapper, shape_dtype_structs, inner_params, vectorized=True)
  File "/Users/mudit.pandey/.pyenv/versions/pennylane/lib/python3.10/site-packages/jax/_src/callback.py", line 334, in pure_callback_api
    return pure_callback(
  File "/Users/mudit.pandey/.pyenv/versions/pennylane/lib/python3.10/site-packages/jax/_src/callback.py", line 261, in pure_callback
    tree_util.tree_map(_check_shape_dtype, result_shape_dtypes)
  File "/Users/mudit.pandey/.pyenv/versions/pennylane/lib/python3.10/site-packages/jax/_src/callback.py", line 216, in _check_shape_dtype
    raise ValueError(
ValueError: Cannot return 64-bit values when `jax_enable_x64` is disabled
>>> jax.config.update("jax_enable_x64", True)
>>> f_jitted(jax.numpy.array(np.pi / 2), shots=(10, 10))
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/mudit.pandey/repos/pennylane/pennylane/qnode.py", line 1010, in __call__
    res = qml.execute(
  File "/Users/mudit.pandey/repos/pennylane/pennylane/interfaces/execution.py", line 736, in execute
    results = ml_boundary_execute(
  File "/Users/mudit.pandey/repos/pennylane/pennylane/interfaces/jax_jit.py", line 278, in execute
    return _execute_bwd(
  File "/Users/mudit.pandey/repos/pennylane/pennylane/interfaces/jax_jit.py", line 363, in _execute_bwd
    return execute_wrapper(params)
  File "/Users/mudit.pandey/repos/pennylane/pennylane/interfaces/jax_jit.py", line 324, in execute_wrapper
    return jax.pure_callback(wrapper, shape_dtype_structs, inner_params, vectorized=True)
  File "/Users/mudit.pandey/.pyenv/versions/pennylane/lib/python3.10/site-packages/jax/_src/callback.py", line 334, in pure_callback_api
    return pure_callback(
  File "/Users/mudit.pandey/.pyenv/versions/pennylane/lib/python3.10/site-packages/jax/_src/callback.py", line 262, in pure_callback
    result_avals = tree_util.tree_map(
  File "/Users/mudit.pandey/.pyenv/versions/pennylane/lib/python3.10/site-packages/jax/_src/callback.py", line 263, in <lambda>
    lambda x: core.ShapedArray(x.shape, x.dtype), result_shape_dtypes)
TypeError: Shapes must be 1D sequences of concrete values of integer type, got ((), ()).
```


### System information

```shell
Name: PennyLane
Version: 0.33.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /Users/mudit.pandey/repos/pennylane
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing_extensions
Required-by: PennyLane-Lightning

Platform info:           macOS-13.6-arm64-arm-64bit
Python version:          3.10.12
Numpy version:           1.23.5
Scipy version:           1.11.1
Installed devices:
- default.gaussian (PennyLane-0.33.0.dev0)
- default.mixed (PennyLane-0.33.0.dev0)
- default.qubit (PennyLane-0.33.0.dev0)
- default.qubit.autograd (PennyLane-0.33.0.dev0)
- default.qubit.jax (PennyLane-0.33.0.dev0)
- default.qubit.legacy (PennyLane-0.33.0.dev0)
- default.qubit.tf (PennyLane-0.33.0.dev0)
- default.qubit.torch (PennyLane-0.33.0.dev0)
- default.qutrit (PennyLane-0.33.0.dev0)
- null.qubit (PennyLane-0.33.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.32.0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.


Discussion:
Additional information:

specifying a shot vector on the device also doesn't work.
Fixed just this case locally by changing `pennylane/interfaces/jax_jit.py:_create_shape_dtype_struct`

```
def _create_shape_dtype_struct(tape: "qml.tape.QuantumScript", device: "qml.Device"):
    """Auxiliary function for creating the shape and dtype object structure
    given a tape."""

    shape = tape.shape(device)
    if len(shape) == 1:
        tape_dtype = _numeric_type_to_dtype(tape.numeric_type)
        return jax.ShapeDtypeStruct(tuple(shape), tape_dtype)

    if tape.shots.has_partitioned_shots:
        tape_dtype = _numeric_type_to_dtype(tape.numeric_type)
        return tuple(jax.ShapeDtypeStruct(tuple(s), tape_dtype) for s in shape)

    tape_dtype = tuple(_numeric_type_to_dtype(elem) for elem in tape.numeric_type)
    return tuple(jax.ShapeDtypeStruct(tuple(s), d) for s, d in zip(shape, tape_dtype))
```

But we will also need to include logical branches for multiple measurements and shot vectors.

Resolving patch:
diff --git a/pennylane/interfaces/jax.py b/pennylane/interfaces/jax.py
--- a/pennylane/interfaces/jax.py
+++ b/pennylane/interfaces/jax.py
@@ -25,6 +25,8 @@
 import pennylane as qml
 from pennylane.transforms import convert_to_numpy_parameters
 
+from .jacobian_products import _compute_jvps
+
 dtype = jnp.float64
 
 logger = logging.getLogger(__name__)
@@ -212,8 +214,7 @@ def execute_wrapper_jvp(primals, tangents):
             # Backward: Gradient function is a device method.
             new_tapes = set_parameters_on_copy_and_unwrap(tapes, primals[0], unwrap=False)
             jacs = gradient_fn(new_tapes, **gradient_kwargs)
-            multi_measurements = [len(tape.measurements) > 1 for tape in new_tapes]
-            jvps = _compute_jvps(jacs, tangents[0], multi_measurements)
+            jvps = _compute_jvps(jacs, tangents[0], new_tapes)
 
         return res, jvps
 
@@ -245,26 +246,14 @@ def execute_wrapper(params):
     def execute_wrapper_jvp(primals, tangents):
         """Primals[0] are parameters as Jax tracers and tangents[0] is a list of tangent vectors as Jax tracers."""
         res, jacs = execute_wrapper(primals[0])
-        multi_measurements = [len(tape.measurements) > 1 for tape in tapes]
 
-        jvps = _compute_jvps(jacs, tangents[0], multi_measurements)
+        jvps = _compute_jvps(jacs, tangents[0], tapes)
         return (res, jacs), (jvps, jacs)
 
     res, _jacs = execute_wrapper(params)
     return res
 
 
-def _compute_jvps(jacs, tangents, multi_measurements):
-    """Compute the jvps of multiple tapes, directly for a Jacobian and tangents."""
-    jvps = []
-    for i, multi in enumerate(multi_measurements):
-        compute_func = (
-            qml.gradients.compute_jvp_multi if multi else qml.gradients.compute_jvp_single
-        )
-        jvps.append(compute_func(tangents[i], jacs[i]))
-    return tuple(jvps)
-
-
 def _is_count_result(r):
     """Checks if ``r`` is a single count (or broadcasted count) result"""
     return isinstance(r, dict) or isinstance(r, list) and all(isinstance(i, dict) for i in r)
diff --git a/pennylane/interfaces/jax_jit.py b/pennylane/interfaces/jax_jit.py
--- a/pennylane/interfaces/jax_jit.py
+++ b/pennylane/interfaces/jax_jit.py
@@ -22,9 +22,10 @@
 import jax.numpy as jnp
 
 import pennylane as qml
-from pennylane.interfaces.jax import _compute_jvps
 from pennylane.transforms import convert_to_numpy_parameters
 
+from .jacobian_products import _compute_jvps
+
 dtype = jnp.float64
 Zero = jax.custom_derivatives.SymbolicZero
 
@@ -62,9 +63,15 @@ def _create_shape_dtype_struct(tape: "qml.tape.QuantumScript", device: "qml.Devi
     shape = tape.shape(device)
     if len(tape.measurements) == 1:
         tape_dtype = _numeric_type_to_dtype(tape.numeric_type)
+        if tape.shots.has_partitioned_shots:
+            return tuple(jax.ShapeDtypeStruct(s, tape_dtype) for s in shape)
         return jax.ShapeDtypeStruct(tuple(shape), tape_dtype)
 
     tape_dtype = tuple(_numeric_type_to_dtype(elem) for elem in tape.numeric_type)
+    if tape.shots.has_partitioned_shots:
+        return tuple(
+            tuple(jax.ShapeDtypeStruct(tuple(s), d) for s, d in zip(si, tape_dtype)) for si in shape
+        )
     return tuple(jax.ShapeDtypeStruct(tuple(s), d) for s, d in zip(shape, tape_dtype))
 
 
@@ -197,7 +204,6 @@ def wrapper(inner_params):
             jacs = execute_fn(jvp_tapes)[0]
             jacs = res_processing_fn(jacs)
             all_jacs.append(jacs)
-
         return tuple(all_jacs)
 
     expected_shapes = tuple(_jac_shape_dtype_struct(t, device) for t in tapes)
@@ -354,8 +360,7 @@ def execute_wrapper_jvp(primals, tangents):
         tangents_trainable = tuple(
             tuple(t for t in tangent if not isinstance(t, Zero)) for tangent in tangents[0]
         )
-        multi_measurements = [len(tape.measurements) > 1 for tape in tapes]
-        jvps = _compute_jvps(jacobians, tangents_trainable, multi_measurements)
+        jvps = _compute_jvps(jacobians, tangents_trainable, tapes)
 
         results = execute_wrapper(params)
         return results, jvps
@@ -425,9 +430,8 @@ def execute_wrapper_jvp(primals, tangents):
         updated_jacs = updated_jacs[0] if len(tapes) == 1 else tuple(updated_jacs)
 
         # Get the jvps
-        multi_measurements = [len(tape.measurements) > 1 for tape in tapes]
         tangents = [[t for t in tangent if not isinstance(t, Zero)] for tangent in tangents[0]]
-        jvps = _compute_jvps(jacs_, tangents, multi_measurements)
+        jvps = _compute_jvps(jacs_, tangents, tapes)
 
         return (res, updated_jacs), (jvps, updated_jacs)
 
diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -811,7 +811,7 @@ def shape(self, device):
         if len(shapes) == 1:
             return shapes[0]
 
-        if len(shots.shot_vector) > 1:
+        if shots.num_copies > 1:
             # put the shot vector axis before the measurement axis
             shapes = tuple(zip(*shapes))
 


Test patch:
diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
--- a/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
@@ -775,6 +775,48 @@ def cost_fn(a, b):
         # if we set the shots to None, backprop can now be used
         assert spy.call_args[1]["gradient_fn"] == "backprop"
 
+    @pytest.mark.parametrize("shots", [(10000, 10000), (10000, 10005)])
+    def test_shot_vectors_single_measurements(self, interface, shots):
+        """Test jax-jit can work with shot vectors."""
+
+        dev = qml.device("default.qubit", shots=shots, seed=4747)
+
+        @jax.jit
+        @qml.qnode(dev, interface=interface, diff_method="parameter-shift")
+        def circuit(x):
+            qml.RX(x, wires=0)
+            return qml.var(qml.PauliZ(0))
+
+        res = circuit(0.5)
+        expected = 1 - np.cos(0.5) ** 2
+        assert qml.math.allclose(res[0], expected, atol=1e-2)
+        assert qml.math.allclose(res[1], expected, atol=3e-2)
+
+        g = jax.jacobian(circuit)(0.5)
+
+        expected_g = 2 * np.cos(0.5) * np.sin(0.5)
+        assert qml.math.allclose(g[0], expected_g, atol=2e-2)
+        assert qml.math.allclose(g[1], expected_g, atol=2e-2)
+
+    @pytest.mark.parametrize("shots", [(10000, 10000), (10000, 10005)])
+    def test_shot_vectors_multiple_measurements(self, interface, shots):
+        """Test jax-jit can work with shot vectors."""
+
+        dev = qml.device("default.qubit", shots=shots, seed=987548)
+
+        @jax.jit
+        @qml.qnode(dev, interface=interface, diff_method="parameter-shift")
+        def circuit(x):
+            qml.RX(x, wires=0)
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=0)
+
+        res = circuit(0.5)
+        assert qml.math.allclose(res[0][0], np.cos(0.5), atol=5e-3)
+        assert qml.math.allclose(res[1][0], np.cos(0.5), atol=5e-3)
+        expected_probs = np.array([np.cos(0.25) ** 2, np.sin(0.25) ** 2])
+        assert qml.math.allclose(res[0][1], expected_probs, atol=5e-3)
+        assert qml.math.allclose(res[1][1], expected_probs, atol=5e-3)
+
 
 @pytest.mark.parametrize(
     "interface,dev,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method