Repository: PennyLaneAI/pennylane (python)

Issue:
[BUG] Jax jacobian of batch when a circuit is not trainable
### Expected behavior

This is an extension of this test: https://github.com/PennyLaneAI/pennylane/blob/45b89b67d28bbe3f8a3f5ab49b0a80a07a92f6a9/tests/interfaces/test_jax.py#L108

where I checked jacobian and `jvp` instead of just a single forward execution.


Fails as well with parameter shift.

### Actual behavior

See traceback below

### Additional information

_No response_

### Source code

```shell
dev = qml.device("default.qubit.legacy", wires=2)

def cost(params):
    tape1 = qml.tape.QuantumScript(
        [qml.RY(params[0], 0), qml.RX(params[1], 0)],
        [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))],
    )

    tape2 = qml.tape.QuantumScript(
        [qml.RY(np.array(0.5), 0)],
        [qml.expval(qml.PauliZ(0))],
    )
    return qml.execute(
        [tape1, tape2],
        dev,
        gradient_fn="device",
        gradient_kwargs={
            "method": "adjoint_jacobian",
            "use_device_state": True,
        },
    )

a = jax.numpy.array([0.1, 0.2])
res = cost(a)

jax.jacobian(cost)(a)
```


### Tracebacks

```shell
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[6], line 1
----> 1 jax.jacobian(cost)(a)

File ~/Prog/pl311/lib/python3.11/site-packages/jax/_src/api.py:938, in jacrev.<locals>.jacfun(*args, **kwargs)
    936 tree_map(partial(_check_input_dtype_jacrev, holomorphic, allow_int), dyn_args)
    937 if not has_aux:
--> 938   y, pullback = _vjp(f_partial, *dyn_args)
    939 else:
    940   y, pullback, aux = _vjp(f_partial, *dyn_args, has_aux=True)

File ~/Prog/pl311/lib/python3.11/site-packages/jax/_src/api.py:2238, in _vjp(fun, has_aux, reduce_axes, *primals)
   2236 if not has_aux:
   2237   flat_fun, out_tree = flatten_fun_nokwargs(fun, in_tree)
-> 2238   out_primal, out_vjp = ad.vjp(
   2239       flat_fun, primals_flat, reduce_axes=reduce_axes)
   2240   out_tree = out_tree()
   2241 else:

File ~/Prog/pl311/lib/python3.11/site-packages/jax/_src/interpreters/ad.py:140, in vjp(traceable, primals, has_aux, reduce_axes)
    138 def vjp(traceable, primals, has_aux=False, reduce_axes=()):
    139   if not has_aux:
--> 140     out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
    141   else:
    142     out_primals, pvals, jaxpr, consts, aux = linearize(traceable, *primals, has_aux=True)

File ~/Prog/pl311/lib/python3.11/site-packages/jax/_src/interpreters/ad.py:129, in linearize(traceable, *primals, **kwargs)
    127 _, in_tree = tree_flatten(((primals, primals), {}))
    128 jvpfun_flat, out_tree = flatten_fun(jvpfun, in_tree)
--> 129 jaxpr, out_pvals, consts = pe.trace_to_jaxpr_nounits(jvpfun_flat, in_pvals)
    130 out_primals_pvals, out_tangents_pvals = tree_unflatten(out_tree(), out_pvals)
    131 assert all(out_primal_pval.is_known() for out_primal_pval in out_primals_pvals)

File ~/Prog/pl311/lib/python3.11/site-packages/jax/_src/profiler.py:340, in annotate_function.<locals>.wrapper(*args, **kwargs)
    337 @wraps(func)
    338 def wrapper(*args, **kwargs):
    339   with TraceAnnotation(name, **decorator_kwargs):
--> 340     return func(*args, **kwargs)
    341   return wrapper

File ~/Prog/pl311/lib/python3.11/site-packages/jax/_src/interpreters/partial_eval.py:777, in trace_to_jaxpr_nounits(fun, pvals, instantiate)
    775 with core.new_main(JaxprTrace, name_stack=current_name_stack) as main:
    776   fun = trace_to_subjaxpr_nounits(fun, main, instantiate)
--> 777   jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    778   assert not env
    779   del main, fun, env

File ~/Prog/pl311/lib/python3.11/site-packages/jax/_src/linear_util.py:191, in WrappedFun.call_wrapped(self, *args, **kwargs)
    188 gen = gen_static_args = out_store = None
    190 try:
--> 191   ans = self.f(*args, **dict(self.params, **kwargs))
    192 except:
    193   # Some transformations yield from inside context managers, so we have to
    194   # interrupt them before reraising the exception. Otherwise they will only
    195   # get garbage-collected at some later time, running their cleanup tasks
    196   # only after this exception is handled, which can corrupt the global
    197   # state.
    198   while stack:

Cell In[4], line 18, in cost(params)
      9 tape2 = qml.tape.QuantumScript(
     10     [qml.RY(np.array(0.5), 0)],
     11     [qml.expval(qml.PauliZ(0))],
     12 )
     14 tape3 = qml.tape.QuantumScript(
     15     [qml.RY(params[0], 0), qml.RX(params[1], 0)],
     16     [qml.expval(qml.PauliZ(0))],
     17 )
---> 18 return qml.execute(
     19     [tape1, tape2],
     20     dev,
     21     gradient_fn="device",
     22     gradient_kwargs={
     23         "method": "adjoint_jacobian",
     24         "use_device_state": True,
     25     },
     26 )

File ~/Prog/pennylane/pennylane/interfaces/execution.py:786, in execute(tapes, device, gradient_fn, interface, transform_program, config, grad_on_execution, gradient_kwargs, cache, cachesize, max_diff, override_shots, expand_fn, max_expansion, device_batch_transform, device_vjp)
    784     results = ml_boundary_execute(tapes, execute_fn, jpc)
    785 else:
--> 786     results = ml_boundary_execute(
    787         tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_diff=max_diff
    788     )
    790 return post_processing(results)

File ~/Prog/pennylane/pennylane/interfaces/jax.py:132, in execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n, max_diff)
    128 parameters = tuple(list(t.get_parameters()) for t in tapes)
    130 if gradient_fn is None:
    131     # PennyLane forward execution
--> 132     return _execute_fwd(
    133         parameters,
    134         tapes,
    135         execute_fn,
    136         gradient_kwargs,
    137         _n=_n,
    138     )
    140 # PennyLane backward execution
    141 return _execute_bwd(
    142     parameters,
    143     tapes,
   (...)
    149     max_diff=max_diff,
    150 )

File ~/Prog/pennylane/pennylane/interfaces/jax.py:253, in _execute_fwd(params, tapes, execute_fn, gradient_kwargs, _n)
    250     jvps = _compute_jvps(jacs, tangents[0], tapes)
    251     return (res, jacs), (jvps, jacs)
--> 253 res, _jacs = execute_wrapper(params)
    254 return res

    [... skipping hidden 5 frame]

File ~/Prog/pl311/lib/python3.11/site-packages/jax/_src/custom_derivatives.py:345, in _flatten_jvp(primal_name, jvp_name, in_tree, maybe_out_type, *args)
    340     msg = ("Custom JVP rule must produce primal and tangent outputs with "
    341            "equal shapes and dtypes, but got:\n{}")
    342     disagreements = (
    343         f"  primal {av1.str_short()} for tangent {av2.str_short()}"
    344         for av1, av2 in zip(primal_avals_out, tangent_avals_out) if av1 != av2)
--> 345     raise TypeError(msg.format('\n'.join(disagreements)))
    346 yield primals_out + tangents_out, (out_tree, primal_avals)

TypeError: Custom JVP rule must produce primal and tangent outputs with equal shapes and dtypes, but got:
  primal float32[] for tangent float32[1,0]
```


### System information

```shell
PL master
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

Pull request description:
fix jax jacobians when one circuit not trainable
Fixes #4834 [sc-50137]

Long story short, our jvp tools, `qml.gradients.jvp` and `qml.interfaces.jacobian_products._compute_jvps` failed to return proper shapes and values when there is no trainable parameters. This fixes that.

Note to reviewers:

To properly test this fix, I needed integration tests for when the device provides a derivative with a shot vector.  So I added tests for device-provided derivatives with a shotvector.  Those test changes are a little painful, but they seem to work. Similar tests have already been added to autograd and torch in previous PR's.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 8cc6d5f68..ad52cb2cd 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -98,6 +98,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Jax can now differeniate a batch of circuits where one tape does not have trainable parameters.
+  [(#4837)](https://github.com/PennyLaneAI/pennylane/pull/4837)
+
 * Fixes a bug where the adjoint method differentiation would fail if
   an operation with `grad_method=None` that has a parameter is present.
   [(#4820)](https://github.com/PennyLaneAI/pennylane/pull/4820)
diff --git a/pennylane/gradients/jvp.py b/pennylane/gradients/jvp.py
index 3d7c9a4dc..f2ff2b6b8 100644
--- a/pennylane/gradients/jvp.py
+++ b/pennylane/gradients/jvp.py
@@ -294,7 +294,11 @@ def jvp(tape, tangent, gradient_fn, gradient_kwargs=None):
     if len(tape.trainable_params) == 0:
         # The tape has no trainable parameters; the JVP
         # is simply none.
-        return [], lambda _, num=None: None
+        def zero_vjp(_):
+            res = tuple(np.zeros(mp.shape(None, tape.shots)) for mp in tape.measurements)
+            return res[0] if len(tape.measurements) == 1 else res
+
+        return tuple(), zero_vjp
 
     multi_m = len(tape.measurements) > 1
 
diff --git a/pennylane/interfaces/jacobian_products.py b/pennylane/interfaces/jacobian_products.py
index 6cf385c22..65bf49b9f 100644
--- a/pennylane/interfaces/jacobian_products.py
+++ b/pennylane/interfaces/jacobian_products.py
@@ -21,6 +21,7 @@ import logging
 from typing import Tuple, Callable, Optional, Union
 
 from cachetools import LRUCache
+import numpy as np
 
 import pennylane as qml
 from pennylane.tape import QuantumScript
@@ -54,7 +55,15 @@ def _compute_jvps(jacs, tangents, tapes):
     jvps = []
     for jac, dx, t in zip(jacs, tangents, tapes):
         multi = len(t.measurements) > 1
-        if t.shots.has_partitioned_shots:
+        if len(t.trainable_params) == 0:
+            empty_shots = qml.measurements.Shots(None)
+            zeros_jvp = tuple(np.zeros(mp.shape(None, empty_shots)) for mp in t.measurements)
+            zeros_jvp = zeros_jvp[0] if len(t.measurements) == 1 else zeros_jvp
+            if t.shots.has_partitioned_shots:
+                jvps.append(tuple(zeros_jvp for _ in range(t.shots.num_copies)))
+            else:
+                jvps.append(zeros_jvp)
+        elif t.shots.has_partitioned_shots:
             jvps.append(tuple(f[multi](dx, j) for j in jac))
         else:
             jvps.append(f[multi](dx, jac))
diff --git a/pennylane/interfaces/jax.py b/pennylane/interfaces/jax.py
index 39ce3836e..2b04f487f 100644
--- a/pennylane/interfaces/jax.py
+++ b/pennylane/interfaces/jax.py
@@ -215,7 +215,6 @@ def _execute_bwd(
             new_tapes = set_parameters_on_copy_and_unwrap(tapes, primals[0], unwrap=False)
             jacs = gradient_fn(new_tapes, **gradient_kwargs)
             jvps = _compute_jvps(jacs, tangents[0], new_tapes)
-
         return res, jvps
 
     return execute_wrapper(params)


Test patch:
diff --git a/tests/gradients/core/test_jvp.py b/tests/gradients/core/test_jvp.py
index ac9730275..2aa09c217 100644
--- a/tests/gradients/core/test_jvp.py
+++ b/tests/gradients/core/test_jvp.py
@@ -350,8 +350,8 @@ class TestJVP:
         tangent = np.array([1.0])
         tapes, fn = qml.gradients.jvp(tape, tangent, param_shift)
 
-        assert tapes == []
-        assert fn(tapes) is None
+        assert tapes == tuple()
+        assert qml.math.allclose(fn(tapes), np.array(0.0))
 
     def test_zero_tangent_single_measurement_single_param(self, batch_dim):
         """A zero tangent vector will return no tapes and a zero matrix"""
@@ -831,7 +831,7 @@ class TestBatchJVP:
         # to the JVP, so only 2*2=4 quantum evals
         res = fn(dev.execute(v_tapes))
 
-        assert res[0] is None
+        assert qml.math.allclose(res[0], np.array(0.0))
         assert res[1] is not None
 
     def test_all_tapes_no_trainable_parameters(self):
@@ -859,7 +859,8 @@ class TestBatchJVP:
         v_tapes, fn = qml.gradients.batch_jvp(tapes, tangents, param_shift)
 
         assert v_tapes == []
-        assert fn([]) == (None, None)
+        assert qml.math.allclose(fn([])[0], np.array(0.0))
+        assert qml.math.allclose(fn([])[1], np.array(0.0))
 
     def test_zero_tangent(self):
         """A zero dy vector will return no tapes and a zero matrix"""
diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
index 5d92ea5b9..15f0913b8 100644
--- a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
@@ -15,11 +15,13 @@
 import pytest
 import numpy as np
 
+from param_shift_dev import ParamShiftDerivativesDevice
 
 import pennylane as qml
 from pennylane.devices import DefaultQubit
 from pennylane.gradients import param_shift
 from pennylane.interfaces import execute
+from pennylane.measurements import Shots
 
 jax = pytest.importorskip("jax")
 jnp = pytest.importorskip("jax.numpy")
@@ -122,16 +124,17 @@ class TestCaching:
 # add tests for lightning 2 when possible
 # set rng for device when possible
 test_matrix = [
-    ({"gradient_fn": param_shift}, 100000, DefaultQubit(seed=42)),
-    ({"gradient_fn": param_shift}, None, DefaultQubit()),
-    ({"gradient_fn": "backprop"}, None, DefaultQubit()),
-    ({"gradient_fn": "adjoint"}, None, DefaultQubit()),
+    ({"gradient_fn": param_shift}, Shots(100000), DefaultQubit(seed=42)),
+    ({"gradient_fn": param_shift}, Shots(None), DefaultQubit()),
+    ({"gradient_fn": "backprop"}, Shots(None), DefaultQubit()),
+    ({"gradient_fn": "adjoint"}, Shots(None), DefaultQubit()),
+    ({"gradient_fn": "device"}, Shots((10000, 10000)), ParamShiftDerivativesDevice(seed=54353453)),
 ]
 
 
 def atol_for_shots(shots):
     """Return higher tolerance if finite shots."""
-    return 1e-2 if shots else 1e-6
+    return 2e-2 if shots else 1e-6
 
 
 @pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix)
@@ -163,8 +166,9 @@ class TestJaxExecuteIntegration:
         assert device.tracker.totals["executions"] == 2  # different wires so different hashes
 
         assert len(res) == 2
-        assert res[0].shape == ()
-        assert res[1].shape == ()
+        if not shots.has_partitioned_shots:
+            assert res[0].shape == ()
+            assert res[1].shape == ()
 
         assert qml.math.allclose(res[0], jnp.cos(a) * jnp.cos(b), atol=atol_for_shots(shots))
         assert qml.math.allclose(res[1], jnp.cos(a) * jnp.cos(b), atol=atol_for_shots(shots))
@@ -178,7 +182,8 @@ class TestJaxExecuteIntegration:
             return execute([tape], device, **execute_kwargs)[0]
 
         res = jax.jacobian(cost)(a)
-        assert res.shape == ()  # pylint: disable=no-member
+        if not shots.has_partitioned_shots:
+            assert res.shape == ()  # pylint: disable=no-member
 
         # compare to standard tape jacobian
         tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))])
@@ -199,20 +204,32 @@ class TestJaxExecuteIntegration:
             ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])]
             m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))]
             tape = qml.tape.QuantumScript(ops, m, shots=shots)
-            return jnp.hstack(execute([tape], device, **execute_kwargs)[0])
+            return execute([tape], device, **execute_kwargs)[0]
 
         res = cost(a, b)
         expected = [jnp.cos(a), -jnp.cos(a) * jnp.sin(b)]
-        assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
+        if shots.has_partitioned_shots:
+            assert np.allclose(res[0], expected, atol=atol_for_shots(shots), rtol=0)
+            assert np.allclose(res[1], expected, atol=2 * atol_for_shots(shots), rtol=0)
+        else:
+            assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
 
-        res = jax.jacobian(cost, argnums=[0, 1])(a, b)
-        assert isinstance(res, tuple) and len(res) == 2
-        assert res[0].shape == (2,)
-        assert res[1].shape == (2,)
+        g = jax.jacobian(cost, argnums=[0, 1])(a, b)
+        assert isinstance(g, tuple) and len(g) == 2
 
         expected = ([-jnp.sin(a), jnp.sin(a) * jnp.sin(b)], [0, -jnp.cos(a) * jnp.cos(b)])
-        for _r, _e in zip(res, expected):
-            assert np.allclose(_r, _e, atol=atol_for_shots(shots))
+
+        if shots.has_partitioned_shots:
+            for i in (0, 1):
+                assert np.allclose(g[i][0][0], expected[0][0], atol=atol_for_shots(shots), rtol=0)
+                assert np.allclose(g[i][1][0], expected[0][1], atol=atol_for_shots(shots), rtol=0)
+                assert np.allclose(g[i][0][1], expected[1][0], atol=atol_for_shots(shots), rtol=0)
+                assert np.allclose(g[i][1][1], expected[1][1], atol=atol_for_shots(shots), rtol=0)
+        else:
+            assert np.allclose(g[0][0], expected[0][0], atol=atol_for_shots(shots), rtol=0)
+            assert np.allclose(g[1][0], expected[0][1], atol=atol_for_shots(shots), rtol=0)
+            assert np.allclose(g[0][1], expected[1][0], atol=atol_for_shots(shots), rtol=0)
+            assert np.allclose(g[1][1], expected[1][1], atol=atol_for_shots(shots), rtol=0)
 
     def test_tape_no_parameters(self, execute_kwargs, shots, device):
         """Test that a tape with no parameters is correctly
@@ -241,7 +258,12 @@ class TestJaxExecuteIntegration:
             tape4 = qml.tape.QuantumScript(
                 [qml.RY(jnp.array(0.5), 0)], [qml.probs(wires=[0, 1])], shots=shots
             )
-            return sum(jnp.hstack(execute([tape1, tape2, tape3, tape4], device, **execute_kwargs)))
+            res = execute([tape1, tape2, tape3, tape4], device, **execute_kwargs)
+            res = jax.tree_util.tree_leaves(res)
+            out = sum(jnp.hstack(res))
+            if shots.has_partitioned_shots:
+                out = out / shots.num_copies
+            return out
 
         params = jnp.array([0.1, 0.2])
         x, y = params
@@ -250,20 +272,14 @@ class TestJaxExecuteIntegration:
         expected = 2 + jnp.cos(0.5) + jnp.cos(x) * jnp.cos(y)
         assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
 
-        # TODO: jax does not allow computing tapes with different gradient
-        # shapes like [None, None, *, None] here
-
-        # grad = jax.grad(cost)(params)
-        # expected = [-jnp.cos(y) * jnp.sin(x), -jnp.cos(x) * jnp.sin(y)]
-        # assert np.allclose(grad, expected, atol=atol_for_shots(shots), rtol=0)
+        grad = jax.grad(cost)(params)
+        expected = [-jnp.cos(y) * jnp.sin(x), -jnp.cos(x) * jnp.sin(y)]
+        assert np.allclose(grad, expected, atol=atol_for_shots(shots), rtol=0)
 
+    # pylint: disable=too-many-statements
     def test_tapes_with_different_return_size(self, execute_kwargs, shots, device):
         """Test that tapes wit different can be executed and differentiated."""
 
-        # TODO: Will probably fail if we update jax to fix this test
-        # if execute_kwargs["gradient_fn"] == "backprop":
-        #     pytest.xfail("backprop is not compatible with something about this situation.")
-
         def cost(params):
             tape1 = qml.tape.QuantumScript(
                 [qml.RY(params[0], 0), qml.RX(params[1], 0)],
@@ -280,35 +296,63 @@ class TestJaxExecuteIntegration:
                 [qml.expval(qml.PauliZ(0))],
                 shots=shots,
             )
-            return jnp.hstack(execute([tape1, tape2, tape3], device, **execute_kwargs))
+            res = qml.execute([tape1, tape2, tape3], device, **execute_kwargs)
+            leaves = jax.tree_util.tree_leaves(res)
+            return jnp.hstack(leaves)
 
         params = jnp.array([0.1, 0.2])
         x, y = params
 
         res = cost(params)
         assert isinstance(res, jax.Array)
-        assert res.shape == (4,)
-
-        assert np.allclose(res[0], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots))
-        assert np.allclose(res[1], 1, atol=atol_for_shots(shots))
-        assert np.allclose(res[2], jnp.cos(0.5), atol=atol_for_shots(shots))
-        assert np.allclose(res[3], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots))
-
-        # TODO: jax does not allow computing tapes with different gradient shapes
-
-        # jac = jax.jacobian(cost)(params)
-        # assert isinstance(jac, jnp.ndarray)
-        # assert jac.shape == (4, 2)  # pylint: disable=no-member
-
-        # assert np.allclose(jac[1:3], 0, atol=atol_for_shots(shots))
+        assert res.shape == (4 * shots.num_copies,) if shots.has_partitioned_shots else (4,)
+
+        if shots.has_partitioned_shots:
+            assert np.allclose(res[0], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots))
+            assert np.allclose(res[2], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots))
+            assert np.allclose(res[1], 1, atol=atol_for_shots(shots))
+            assert np.allclose(res[3], 1, atol=atol_for_shots(shots))
+            assert np.allclose(res[4], jnp.cos(0.5), atol=atol_for_shots(shots))
+            assert np.allclose(res[5], jnp.cos(0.5), atol=atol_for_shots(shots))
+            assert np.allclose(res[6], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots))
+            assert np.allclose(res[7], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots))
+        else:
+            assert np.allclose(res[0], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots))
+            assert np.allclose(res[1], 1, atol=atol_for_shots(shots))
+            assert np.allclose(res[2], jnp.cos(0.5), atol=atol_for_shots(shots))
+            assert np.allclose(res[3], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots))
 
-        # d1 = -jnp.sin(x) * jnp.cos(y)
-        # assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots))
-        # assert np.allclose(jac[3, 0], d1, atol=atol_for_shots(shots))
+        jac = jax.jacobian(cost)(params)
+        assert isinstance(jac, jnp.ndarray)
+        assert (
+            jac.shape == (8, 2) if shots.has_partitioned_shots else (4, 2)
+        )  # pylint: disable=no-member
 
-        # d2 = -jnp.cos(x) * jnp.sin(y)
-        # assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots))
-        # assert np.allclose(jac[3, 1], d2, atol=atol_for_shots(shots))
+        if shots.has_partitioned_shots:
+            assert np.allclose(jac[1], 0, atol=atol_for_shots(shots))
+            assert np.allclose(jac[3:5], 0, atol=atol_for_shots(shots))
+        else:
+            assert np.allclose(jac[1:3], 0, atol=atol_for_shots(shots))
+
+        d1 = -jnp.sin(x) * jnp.cos(y)
+        if shots.has_partitioned_shots:
+            assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots))
+            assert np.allclose(jac[2, 0], d1, atol=atol_for_shots(shots))
+            assert np.allclose(jac[6, 0], d1, atol=atol_for_shots(shots))
+            assert np.allclose(jac[7, 0], d1, atol=atol_for_shots(shots))
+        else:
+            assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots))
+            assert np.allclose(jac[3, 0], d1, atol=atol_for_shots(shots))
+
+        d2 = -jnp.cos(x) * jnp.sin(y)
+        if shots.has_partitioned_shots:
+            assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots))
+            assert np.allclose(jac[2, 1], d2, atol=atol_for_shots(shots))
+            assert np.allclose(jac[6, 1], d2, atol=atol_for_shots(shots))
+            assert np.allclose(jac[7, 1], d2, atol=atol_for_shots(shots))
+        else:
+            assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots))
+            assert np.allclose(jac[3, 1], d2, atol=atol_for_shots(shots))
 
     def test_reusing_quantum_tape(self, execute_kwargs, shots, device):
         """Test re-using a quantum tape by passing new parameters"""
@@ -339,7 +383,11 @@ class TestJaxExecuteIntegration:
         # values of the parameters for subsequent calls
         res2 = cost(2 * a, b)
         expected = [jnp.cos(2 * a), -jnp.cos(2 * a) * jnp.sin(b)]
-        assert np.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0)
+        if shots.has_partitioned_shots:
+            assert np.allclose(res2[:2], expected, atol=atol_for_shots(shots), rtol=0)
+            assert np.allclose(res2[2:], expected, atol=atol_for_shots(shots), rtol=0)
+        else:
+            assert np.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0)
 
         jac_fn = jax.jacobian(lambda a, b: cost(2 * a, b), argnums=[0, 1])
         jac = jac_fn(a, b)
@@ -348,8 +396,15 @@ class TestJaxExecuteIntegration:
             [0, -jnp.cos(2 * a) * jnp.cos(b)],
         )
         assert isinstance(jac, tuple) and len(jac) == 2
-        for _j, _e in zip(jac, expected):
-            assert np.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0)
+        if shots.has_partitioned_shots:
+            for offset in (0, 2):
+                assert np.allclose(jac[0][0 + offset], expected[0][0], atol=atol_for_shots(shots))
+                assert np.allclose(jac[0][1 + offset], expected[0][1], atol=atol_for_shots(shots))
+                assert np.allclose(jac[1][0 + offset], expected[1][0], atol=atol_for_shots(shots))
+                assert np.allclose(jac[1][1 + offset], expected[1][1], atol=atol_for_shots(shots))
+        else:
+            for _j, _e in zip(jac, expected):
+                assert np.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0)
 
     def test_classical_processing(self, execute_kwargs, shots, device):
         """Test classical processing within the quantum tape"""
@@ -371,8 +426,9 @@ class TestJaxExecuteIntegration:
 
         # Only two arguments are trainable
         assert isinstance(res, tuple) and len(res) == 2
-        assert res[0].shape == ()
-        assert res[1].shape == ()
+        if not shots.has_partitioned_shots:
+            assert res[0].shape == ()
+            assert res[1].shape == ()
 
         # I tried getting analytic results for this circuit but I kept being wrong and am giving up
 
@@ -392,7 +448,8 @@ class TestJaxExecuteIntegration:
 
         jac_fn = jax.jacobian(cost)
         jac = jac_fn(a, U)
-        assert isinstance(jac, jnp.ndarray)
+        if not shots.has_partitioned_shots:
+            assert isinstance(jac, jnp.ndarray)
         assert np.allclose(jac, jnp.sin(a), atol=atol_for_shots(shots), rtol=0)
 
     def test_differentiable_expand(self, execute_kwargs, device, shots):
@@ -482,13 +539,17 @@ class TestJaxExecuteIntegration:
                 ],
             ]
         )
-        assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
+        if shots.has_partitioned_shots:
+            assert np.allclose(res[:, 0:2].flatten(), expected, atol=atol_for_shots(shots))
+            assert np.allclose(res[:, 2:].flatten(), expected, atol=atol_for_shots(shots))
+        else:
+            assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
 
         jac_fn = jax.jacobian(cost, argnums=[0, 1])
         res = jac_fn(x, y)
         assert isinstance(res, tuple) and len(res) == 2
-        assert res[0].shape == (4,)
-        assert res[1].shape == (4,)
+        assert res[0].shape == (2, 4) if shots.has_partitioned_shots else (4,)
+        assert res[1].shape == (2, 4) if shots.has_partitioned_shots else (4,)
 
         expected = (
             jnp.array(
@@ -508,8 +569,14 @@ class TestJaxExecuteIntegration:
             ),
         )
 
-        assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0)
-        assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0)
+        if shots.has_partitioned_shots:
+            assert np.allclose(res[0][:, 0:2].flatten(), expected[0], atol=atol_for_shots(shots))
+            assert np.allclose(res[0][:, 2:].flatten(), expected[0], atol=atol_for_shots(shots))
+            assert np.allclose(res[1][:, :2].flatten(), expected[1], atol=atol_for_shots(shots))
+            assert np.allclose(res[1][:, 2:].flatten(), expected[1], atol=atol_for_shots(shots))
+        else:
+            assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0)
+            assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0)
 
     def test_ragged_differentiation(self, execute_kwargs, device, shots):
         """Tests correct output shape and evaluation for a tape
@@ -521,7 +588,8 @@ class TestJaxExecuteIntegration:
             ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))]
             m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)]
             tape = qml.tape.QuantumScript(ops, m, shots=shots)
-            return jnp.hstack(execute([tape], device, **execute_kwargs)[0])
+            res = qml.execute([tape], device, **execute_kwargs)[0]
+            return jnp.hstack(jax.tree_util.tree_leaves(res))
 
         x = jnp.array(0.543)
         y = jnp.array(-0.654)
@@ -530,20 +598,30 @@ class TestJaxExecuteIntegration:
         expected = jnp.array(
             [jnp.cos(x), (1 + jnp.cos(x) * jnp.cos(y)) / 2, (1 - jnp.cos(x) * jnp.cos(y)) / 2]
         )
-        assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
+        if shots.has_partitioned_shots:
+            assert np.allclose(res[:3], expected, atol=atol_for_shots(shots), rtol=0)
+            assert np.allclose(res[3:], expected, atol=atol_for_shots(shots), rtol=0)
+        else:
+            assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
 
         jac_fn = jax.jacobian(cost, argnums=[0, 1])
         res = jac_fn(x, y)
         assert isinstance(res, tuple) and len(res) == 2
-        assert res[0].shape == (3,)
-        assert res[1].shape == (3,)
+        assert res[0].shape == (3 * shots.num_copies,) if shots.has_partitioned_shots else (3,)
+        assert res[1].shape == (3 * shots.num_copies,) if shots.has_partitioned_shots else (3,)
 
         expected = (
             jnp.array([-jnp.sin(x), -jnp.sin(x) * jnp.cos(y) / 2, jnp.sin(x) * jnp.cos(y) / 2]),
             jnp.array([0, -jnp.cos(x) * jnp.sin(y) / 2, jnp.cos(x) * jnp.sin(y) / 2]),
         )
-        assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0)
-        assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0)
+        if shots.has_partitioned_shots:
+            assert np.allclose(res[0][:3], expected[0], atol=atol_for_shots(shots), rtol=0)
+            assert np.allclose(res[0][3:], expected[0], atol=atol_for_shots(shots), rtol=0)
+            assert np.allclose(res[1][:3], expected[1], atol=atol_for_shots(shots), rtol=0)
+            assert np.allclose(res[1][3:], expected[1], atol=atol_for_shots(shots), rtol=0)
+        else:
+            assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0)
+            assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0)
 
 
 class TestHigherOrderDerivatives:
@@ -652,7 +730,10 @@ class TestHamiltonianWorkflows:
                 qml.expval(H2)
 
             tape = qml.tape.QuantumScript.from_queue(q, shots=shots)
-            return jnp.hstack(execute([tape], device, **execute_kwargs)[0])
+            res = execute([tape], device, **execute_kwargs)[0]
+            if shots.has_partitioned_shots:
+                return jnp.hstack(res[0] + res[1])
+            return jnp.hstack(res)
 
         return _cost_fn
 
@@ -698,11 +779,19 @@ class TestHamiltonianWorkflows:
 
         res = cost_fn(weights, coeffs1, coeffs2)
         expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
-        assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
+        if shots.has_partitioned_shots:
+            assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0)
+            assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0)
+        else:
+            assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
 
         res = jax.jacobian(cost_fn)(weights, coeffs1, coeffs2)
         expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2]
-        assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
+        if shots.has_partitioned_shots:
+            assert np.allclose(res[:2, :], expected, atol=atol_for_shots(shots), rtol=0)
+            assert np.allclose(res[2:, :], expected, atol=atol_for_shots(shots), rtol=0)
+        else:
+            assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
 
     def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots, use_new_op_math):
         """Test hamiltonian with trainable parameters."""
@@ -717,8 +806,17 @@ class TestHamiltonianWorkflows:
 
         res = cost_fn(weights, coeffs1, coeffs2)
         expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
-        assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
+        if shots.has_partitioned_shots:
+            assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0)
+            assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0)
+        else:
+            assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
 
         res = jnp.hstack(jax.jacobian(cost_fn, argnums=[0, 1, 2])(weights, coeffs1, coeffs2))
         expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)
-        assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
+        if shots.has_partitioned_shots:
+            pytest.xfail(
+                "multiple hamiltonians with shot vectors does not seem to be differentiable."
+            )
+        else:
+            assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0)
diff --git a/tests/interfaces/test_jax.py b/tests/interfaces/test_jax.py
index 8d774e698..6374b2772 100644
--- a/tests/interfaces/test_jax.py
+++ b/tests/interfaces/test_jax.py
@@ -148,6 +148,13 @@ class TestJaxExecuteUnitTests:
         assert dev.num_executions == 3
         spy.assert_called()
 
+        g = jax.jacobian(cost)(a)
+        expected_g = (-np.sin(x) * np.cos(y), -np.cos(x) * np.sin(y))
+        assert qml.math.allclose(g[0][0], expected_g)
+        assert qml.math.allclose(g[0][1], np.zeros(2))
+        assert qml.math.allclose(g[1], np.zeros(2))
+        assert qml.math.allclose(g[2], expected_g)
+
     def test_no_grad_on_execution(self, mocker):
         """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway"""
         dev = qml.device("default.qubit.legacy", wires=1)
@@ -592,7 +599,7 @@ class TestJaxExecuteIntegration:
         res = cost(a, U, device=dev)
         assert np.allclose(res, -np.cos(a), atol=tol, rtol=0)
 
-        jac_fn = jax.grad(cost, argnums=(0))
+        jac_fn = jax.grad(cost, argnums=0)
         res = jac_fn(a, U, device=dev)
         assert np.allclose(res, np.sin(a), atol=tol, rtol=0)
 
@@ -629,7 +636,7 @@ class TestJaxExecuteIntegration:
         )
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-        jac_fn = jax.grad(cost_fn, argnums=(1))
+        jac_fn = jax.grad(cost_fn, argnums=1)
         res = jac_fn(a, p, device=dev)
         expected = jax.numpy.array(
             [