Repository: PennyLaneAI/pennylane

Issue:
[BUG] Jitting with JAX produces incorrect results with parameter broadcasting
### Expected behavior

Consider the following QNode:
```python
dev = qml.device("default.mixed", wires=2)

@qml.qnode(dev, diff_method="backprop", interface="jax")
def circuit(x):
    qml.RX(x, wires=0)
    qml.CNOT(wires=[0, 1])
    return qml.expval(qml.PauliZ(1))
```
When the QNode is called with a parameter, the following should be output:
```pycon
>>> x = jnp.array([0.3, 0.6])
>>> jax.jit(circuit)(x)
[0.9553365 0.8253356]
```


### Actual behavior

Instead, the following is output:
```pycon
>>> x = jnp.array([0.3, 0.6])
>>> jax.jit(circuit)(x)
[0.95533645 0.95533645]
```
It seems that only the first parameter is used for all calculations and the rest are ignored.

### Additional information

If `diff_method="parameter-shift"` is set in the QNode, the issue disappears.

If `jax.jit` is removed in the above code, the issue disappears.

If `"default.qubit"` is used instead, the issue disappears.

The issue also appears if the QNode is decorated with `qml.transforms.batch_params`.


### Source code

_No response_

### Tracebacks

_No response_

### System information

```shell
Name: PennyLane
Version: 0.24.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: None
Author-email: None
License: Apache License 2.0
Location: c:\users\edward.jiang\documents\pennylane
Requires: numpy, scipy, networkx, retworkx, autograd, toml, appdirs, semantic-version, autoray, cachetools, pennylane-lightning
Required-by: PennyLane-Lightning

Platform info:           Windows-10-10.0.19042-SP0
Python version:          3.8.10
Numpy version:           1.22.3
Scipy version:           1.8.0
Installed devices:
- default.gaussian (PennyLane-0.24.0.dev0)
- default.mixed (PennyLane-0.24.0.dev0)
- default.mixed.autograd (PennyLane-0.24.0.dev0)
- default.qubit (PennyLane-0.24.0.dev0)
- default.qubit.autograd (PennyLane-0.24.0.dev0)
- default.qubit.jax (PennyLane-0.24.0.dev0)
- default.qubit.tf (PennyLane-0.24.0.dev0)
- default.qubit.torch (PennyLane-0.24.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.23.0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.
[BUG] Jitting with JAX produces incorrect results with parameter broadcasting
### Expected behavior

Consider the following QNode:
```python
dev = qml.device("default.mixed", wires=2)

@qml.qnode(dev, diff_method="backprop", interface="jax")
def circuit(x):
    qml.RX(x, wires=0)
    qml.CNOT(wires=[0, 1])
    return qml.expval(qml.PauliZ(1))
```
When the QNode is called with a parameter, the following should be output:
```pycon
>>> x = jnp.array([0.3, 0.6])
>>> jax.jit(circuit)(x)
[0.9553365 0.8253356]
```


### Actual behavior

Instead, the following is output:
```pycon
>>> x = jnp.array([0.3, 0.6])
>>> jax.jit(circuit)(x)
[0.95533645 0.95533645]
```
It seems that only the first parameter is used for all calculations and the rest are ignored.

### Additional information

If `diff_method="parameter-shift"` is set in the QNode, the issue disappears.

If `jax.jit` is removed in the above code, the issue disappears.

If `"default.qubit"` is used instead, the issue disappears.

The issue also appears if the QNode is decorated with `qml.transforms.batch_params`.


### Source code

_No response_

### Tracebacks

_No response_

### System information

```shell
Name: PennyLane
Version: 0.24.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: None
Author-email: None
License: Apache License 2.0
Location: c:\users\edward.jiang\documents\pennylane
Requires: numpy, scipy, networkx, retworkx, autograd, toml, appdirs, semantic-version, autoray, cachetools, pennylane-lightning
Required-by: PennyLane-Lightning

Platform info:           Windows-10-10.0.19042-SP0
Python version:          3.8.10
Numpy version:           1.22.3
Scipy version:           1.8.0
Installed devices:
- default.gaussian (PennyLane-0.24.0.dev0)
- default.mixed (PennyLane-0.24.0.dev0)
- default.mixed.autograd (PennyLane-0.24.0.dev0)
- default.qubit (PennyLane-0.24.0.dev0)
- default.qubit.autograd (PennyLane-0.24.0.dev0)
- default.qubit.jax (PennyLane-0.24.0.dev0)
- default.qubit.tf (PennyLane-0.24.0.dev0)
- default.qubit.torch (PennyLane-0.24.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.23.0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.


Discussion:
> If the `qml.transforms.batch_params` decorator is removed, the issue stays. This gives reason to believe that the issue is with the newly introduced parameter broadcasting.

@eddddddy given that `batch_params` does not seem to be the cause of this issue, might be worth editing the title + description + main example to reflect this 🙂 

Perhaps this is more related to the recent additions to JAX-JIT support than to parameter broadcasting? Not entirely sure 😕.
Hi @eddddddy,

> Perhaps this is more related to the recent additions to JAX-JIT support than to parameter broadcasting? Not entirely sure confused.

as we noted this during our chat, the latest changes were added to the JAX-JIT interface and those changes should not be relevant for `diff_method="backprop"` (unless the interface is mistakenly being applied).

The idea behind interfaces is that
1) they convert between NumPy and ML framework (i.e., JAX) objects to execute circuits on any device and
2) they provide custom gradient rules.

Interfaces allow us to feed in ML objects into a QNode even when the target device would not natively understand the ML objects.

When we have `diff_method="backprop"`, none of those are required because the entire computation is performed using the ML framework's objects (e.g., `jnp.array` objects). In such a case with the `default.qubit` device, under the hood, we would be swapping our device to `default.qubit.jax`.
Referencing [this test](https://github.com/PennyLaneAI/pennylane/pull/2754#discussion_r903101424) here so we remember to change it once this is fixed.
Thanks for digging this one out! :+1:  The reason seems to be, that `qml.transforms.broadcast_expand` is not JIT compatible.

**edit:** Actually, that's not quite true. The problem is that by default, `qml.execute`, which is called in `QNode.__call__`, makes use of `qml.interfaces.cache_execute`, which in turn caches the result and therefore does not produce multiple different, but only a single result, which then is retrieved from cache. I patched PL to print out the cache values and receive, when calling the example above:
```pycon
>>> out = circuit(x)
cached values: [Traced<ShapedArray(float64[1])>with<DynamicJaxprTrace(level=0/1)>]
```
As we can see, only one value is computed and stored in the cache. I think this is because all traced tapes after applying `broadcast_transform` have the same hash, and I'm not sure it's possible to change that.

**What next?**

1. With the introduction of parameter broadcasting to `DefaultQubit`, this problem will be gone, but only for `DefaultQubit` devices. In particular, `DefaultMixed` and basically all other devices will still have this problem.
2. This problem for example did not occur when the device's `batch_transform` produces multiple tapes because of a Hamiltonian decomposition. That's because the resulting terms differ in the measured observables and the tape hashes also differ (if they don't, the Hamiltonian would not be decomposed, I suppose)
3. I think we should deactivate caching when executing broadcasted tapes, but I am not sure how to do this best. One option would be to override (but not overwrite!) the execution kwarg `cache` in `QNode.__call__` if the QNode tape after construction has a `batch_size` that is not None, and if the QNode device does not support broadcasting (there is a device flag for this). @josh146 what do you think?

Also tagging @antalszava for visibility. Do you have an idea how to proceed best here? :)
I looked into this once more.
1. `default.qubit` does not longer produce the bug, because it now makes use of broadcasting. Replace it by `default.mixed` for reproducing/testing purposes. This also reduces the urgency/severity of the bug I think.
2. The problem is that `qml.tape.QuantumTape.hash` is not jit-compatible, because the parameters of the tape operations are converted to strings, making different calls with a tracer indistinguishable. 
3.  Taking `hash(tape)` into account in `tape.hash`, for example by simply including it in the `fingerprint` created in `QuantumTape.hash`, does resolve this bug, but it also prevents the correct recycling of cached results.
4. It seems that the best hotfix would be to internally deactivate caching in the particular use case of
   a. `dev.capabilities()["supports_broadcasting"] = False` , and
   b. using parameter-broadcasting, and
   c. `diff_method = "backprop"`, and
   d. `interface = "jax"`, and
   e. jitting, and
   f. `cache = True` (the default) in the QNode kwargs (so this would be overridden).
5. However, this hotfix might be unnecessary, because an immutable data structure replacing `QuantumTape` is on its way, and it is expected to not have this problem.
6. Side comment: `tape.hash` is 3 orders of magnitude slower than the default `hash(tape)`, with 200 µs vs 150 ns for a small tape. The latter does not work for our purposes, so we can not simply replace it, but it's good to keep this overhead in mind I think.
I can no longer reproduce this bug on master. @eddddddy, @josh146 , can you confirm?
@dwierichs if the minimal non-working example in the issue now works in `master`, I would say safe to close
Made the mistake of closing this, because it is fixed for `"default.qubit"`, although I had mentioned already that this will happen :facepalm: : 
>With the introduction of parameter broadcasting to DefaultQubit, this problem will be gone, but only for DefaultQubit devices. In particular, DefaultMixed and basically all other devices will still have this problem.

I reopen the issue and modify the reported example to use `"default.mixed"` instead.
> If the `qml.transforms.batch_params` decorator is removed, the issue stays. This gives reason to believe that the issue is with the newly introduced parameter broadcasting.

@eddddddy given that `batch_params` does not seem to be the cause of this issue, might be worth editing the title + description + main example to reflect this 🙂 

Perhaps this is more related to the recent additions to JAX-JIT support than to parameter broadcasting? Not entirely sure 😕.
Hi @eddddddy,

> Perhaps this is more related to the recent additions to JAX-JIT support than to parameter broadcasting? Not entirely sure confused.

as we noted this during our chat, the latest changes were added to the JAX-JIT interface and those changes should not be relevant for `diff_method="backprop"` (unless the interface is mistakenly being applied).

The idea behind interfaces is that
1) they convert between NumPy and ML framework (i.e., JAX) objects to execute circuits on any device and
2) they provide custom gradient rules.

Interfaces allow us to feed in ML objects into a QNode even when the target device would not natively understand the ML objects.

When we have `diff_method="backprop"`, none of those are required because the entire computation is performed using the ML framework's objects (e.g., `jnp.array` objects). In such a case with the `default.qubit` device, under the hood, we would be swapping our device to `default.qubit.jax`.
Referencing [this test](https://github.com/PennyLaneAI/pennylane/pull/2754#discussion_r903101424) here so we remember to change it once this is fixed.
Thanks for digging this one out! :+1:  The reason seems to be, that `qml.transforms.broadcast_expand` is not JIT compatible.

**edit:** Actually, that's not quite true. The problem is that by default, `qml.execute`, which is called in `QNode.__call__`, makes use of `qml.interfaces.cache_execute`, which in turn caches the result and therefore does not produce multiple different, but only a single result, which then is retrieved from cache. I patched PL to print out the cache values and receive, when calling the example above:
```pycon
>>> out = circuit(x)
cached values: [Traced<ShapedArray(float64[1])>with<DynamicJaxprTrace(level=0/1)>]
```
As we can see, only one value is computed and stored in the cache. I think this is because all traced tapes after applying `broadcast_transform` have the same hash, and I'm not sure it's possible to change that.

**What next?**

1. With the introduction of parameter broadcasting to `DefaultQubit`, this problem will be gone, but only for `DefaultQubit` devices. In particular, `DefaultMixed` and basically all other devices will still have this problem.
2. This problem for example did not occur when the device's `batch_transform` produces multiple tapes because of a Hamiltonian decomposition. That's because the resulting terms differ in the measured observables and the tape hashes also differ (if they don't, the Hamiltonian would not be decomposed, I suppose)
3. I think we should deactivate caching when executing broadcasted tapes, but I am not sure how to do this best. One option would be to override (but not overwrite!) the execution kwarg `cache` in `QNode.__call__` if the QNode tape after construction has a `batch_size` that is not None, and if the QNode device does not support broadcasting (there is a device flag for this). @josh146 what do you think?

Also tagging @antalszava for visibility. Do you have an idea how to proceed best here? :)
I looked into this once more.
1. `default.qubit` does not longer produce the bug, because it now makes use of broadcasting. Replace it by `default.mixed` for reproducing/testing purposes. This also reduces the urgency/severity of the bug I think.
2. The problem is that `qml.tape.QuantumTape.hash` is not jit-compatible, because the parameters of the tape operations are converted to strings, making different calls with a tracer indistinguishable. 
3.  Taking `hash(tape)` into account in `tape.hash`, for example by simply including it in the `fingerprint` created in `QuantumTape.hash`, does resolve this bug, but it also prevents the correct recycling of cached results.
4. It seems that the best hotfix would be to internally deactivate caching in the particular use case of
   a. `dev.capabilities()["supports_broadcasting"] = False` , and
   b. using parameter-broadcasting, and
   c. `diff_method = "backprop"`, and
   d. `interface = "jax"`, and
   e. jitting, and
   f. `cache = True` (the default) in the QNode kwargs (so this would be overridden).
5. However, this hotfix might be unnecessary, because an immutable data structure replacing `QuantumTape` is on its way, and it is expected to not have this problem.
6. Side comment: `tape.hash` is 3 orders of magnitude slower than the default `hash(tape)`, with 200 µs vs 150 ns for a small tape. The latter does not work for our purposes, so we can not simply replace it, but it's good to keep this overhead in mind I think.
I can no longer reproduce this bug on master. @eddddddy, @josh146 , can you confirm?
@dwierichs if the minimal non-working example in the issue now works in `master`, I would say safe to close
Made the mistake of closing this, because it is fixed for `"default.qubit"`, although I had mentioned already that this will happen :facepalm: : 
>With the introduction of parameter broadcasting to DefaultQubit, this problem will be gone, but only for DefaultQubit devices. In particular, DefaultMixed and basically all other devices will still have this problem.

I reopen the issue and modify the reported example to use `"default.mixed"` instead.

Resolving patch:
diff --git a/pennylane/gradients/metric_tensor.py b/pennylane/gradients/metric_tensor.py
--- a/pennylane/gradients/metric_tensor.py
+++ b/pennylane/gradients/metric_tensor.py
@@ -364,7 +364,6 @@ def circuit(weights):
             "as trainable."
         )
 
-    # pylint: disable=too-many-arguments
     if approx in {"diag", "block-diag"}:
         # Only require covariance matrix based transform
         diag_approx = approx == "diag"
diff --git a/pennylane/operation.py b/pennylane/operation.py
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -387,16 +387,22 @@ def classproperty(func):
 
 
 def _process_data(op):
+    def _mod_and_round(x, mod_val):
+        if mod_val is None:
+            return x
+        return qml.math.round(qml.math.real(x) % mod_val, 10)
+
     # Use qml.math.real to take the real part. We may get complex inputs for
     # example when differentiating holomorphic functions with JAX: a complex
     # valued QNode (one that returns qml.state) requires complex typed inputs.
     if op.name in ("RX", "RY", "RZ", "PhaseShift", "Rot"):
-        return str([qml.math.round(qml.math.real(d) % (2 * np.pi), 10) for d in op.data])
-
-    if op.name in ("CRX", "CRY", "CRZ", "CRot"):
-        return str([qml.math.round(qml.math.real(d) % (4 * np.pi), 10) for d in op.data])
+        mod_val = 2 * np.pi
+    elif op.name in ("CRX", "CRY", "CRZ", "CRot"):
+        mod_val = 4 * np.pi
+    else:
+        mod_val = None
 
-    return str(op.data)
+    return str([id(d) if qml.math.is_abstract(d) else _mod_and_round(d, mod_val) for d in op.data])
 
 
 class Operator(abc.ABC):


Test patch:
diff --git a/tests/test_operation.py b/tests/test_operation.py
--- a/tests/test_operation.py
+++ b/tests/test_operation.py
@@ -1004,7 +1004,9 @@ class DummyObserv(qml.operation.Observable):
         op = DummyObserv(1.0, wires=0, id="test")
         assert op.id == "test"
 
-    def test_wire_is_given_in_argument(self):
+    def test_raises_if_no_wire_is_given(self):
+        """Test that an error is raised if no wire is passed at initialization."""
+
         class DummyObservable(qml.operation.Observable):
             num_wires = 1
 
@@ -1193,11 +1195,8 @@ def test_divide_not_supported(self):
         """Test that the division of an operator with an unknown object is not supported."""
         obs = qml.PauliX(0)
 
-        class UnknownObject:
-            pass
-
         with pytest.raises(TypeError, match="unsupported operand type"):
-            _ = obs / UnknownObject()
+            _ = obs / object()
 
     def test_dunder_method_with_new_class(self):
         """Test that when calling any Operator dunder method with a non-supported class that
diff --git a/tests/transforms/test_broadcast_expand.py b/tests/transforms/test_broadcast_expand.py
--- a/tests/transforms/test_broadcast_expand.py
+++ b/tests/transforms/test_broadcast_expand.py
@@ -21,65 +21,50 @@
 from pennylane import numpy as pnp
 
 
-class RX_broadcasted(qml.RX):
-    """A version of qml.RX that detects batching."""
-
-    ndim_params = (0,)
-    compute_decomposition = staticmethod(lambda theta, wires=None: qml.RX(theta, wires=wires))
-
-
-class RZ_broadcasted(qml.RZ):
-    """A version of qml.RZ that detects batching."""
-
-    ndim_params = (0,)
-    compute_decomposition = staticmethod(lambda theta, wires=None: qml.RZ(theta, wires=wires))
-
-
 dev = qml.device("default.qubit", wires=2)
 """Defines the device used for all tests"""
 
 
-def make_tape(x, y, z, obs):
-    """Construct a tape with three parametrized, two unparametrized
-    operations and expvals of provided observables."""
-    with qml.queuing.AnnotatedQueue() as q:
-        qml.StatePrep(np.array([1, 0, 0, 0]), wires=[0, 1])
-        RX_broadcasted(x, wires=0)
-        qml.PauliY(0)
-        RX_broadcasted(y, wires=1)
-        RZ_broadcasted(z, wires=1)
-        qml.Hadamard(1)
-        for ob in obs:
-            qml.expval(ob)
-
-    tape = qml.tape.QuantumScript.from_queue(q)
-    return tape
-
-
-parameters_and_size = [
-    [(0.2, np.array([0.1, 0.8, 2.1]), -1.5), 3],
-    [(0.2, np.array([0.1]), np.array([-0.3])), 1],
-    [
-        (
-            0.2,
-            pnp.array([0.1, 0.3], requires_grad=True),
-            pnp.array([-0.3, 2.1], requires_grad=False),
-        ),
-        2,
-    ],
+def make_ops(x, y, z):
+    """Queue three parametrized and two unparametrized operations and return them."""
+    ops = [
+        qml.RX(x, wires=0),
+        qml.PauliY(0),
+        qml.RX(y, wires=1),
+        qml.RZ(z, wires=1),
+        qml.Hadamard(1),
+    ]
+    return ops
+
+
+parameters = [
+    (0.2, np.array([0.1, 0.8, 2.1]), -1.5),
+    (0.2, np.array([0.1]), np.array([-0.3])),
+    (
+        0.2,
+        pnp.array([0.1, 0.3], requires_grad=True),
+        pnp.array([-0.3, 2.1], requires_grad=False),
+    ),
 ]
+sizes = [3, 1, 2]
 
 coeffs0 = [0.3, -5.1]
 H0 = qml.Hamiltonian(qml.math.array(coeffs0), [qml.PauliZ(0), qml.PauliY(1)])
 
+
 # Here we exploit the product structure of our circuit
-exp_fn_Z0 = lambda x, y, z: -qml.math.cos(x) * qml.math.ones_like(y) * qml.math.ones_like(z)
-exp_fn_Y1 = lambda x, y, z: qml.math.sin(y) * qml.math.cos(z) * qml.math.ones_like(x)
+def exp_fn_Z0(x, y, z):
+    out = -qml.math.cos(x) * qml.math.ones_like(y) * qml.math.ones_like(z)
+    return out[0] if len(out) == 1 else out
+
+
+def exp_fn_Y1(x, y, z):
+    out = qml.math.sin(y) * qml.math.cos(z) * qml.math.ones_like(x)
+    return out[0] if len(out) == 1 else out
+
+
 exp_fn_Z0Y1 = lambda x, y, z: exp_fn_Z0(x, y, z) * exp_fn_Y1(x, y, z)
-exp_fn_Z0_and_Y1 = lambda x, y, z: qml.math.array(
-    [exp_fn_Z0(x, y, z), exp_fn_Y1(x, y, z)],
-    like=exp_fn_Z0(x, y, z) + exp_fn_Y1(x, y, z),
-)
+exp_fn_Z0_and_Y1 = lambda x, y, z: qml.math.stack([exp_fn_Z0(x, y, z), exp_fn_Y1(x, y, z)])
 exp_fn_H0 = lambda x, y, z: exp_fn_Z0(x, y, z) * coeffs0[0] + exp_fn_Y1(x, y, z) * coeffs0[1]
 
 observables_and_exp_fns = [
@@ -93,11 +78,13 @@ def make_tape(x, y, z, obs):
 class TestBroadcastExpand:
     """Tests for the broadcast_expand transform"""
 
-    @pytest.mark.parametrize("params, size", parameters_and_size)
+    @pytest.mark.parametrize("params, size", list(zip(parameters, sizes)))
     @pytest.mark.parametrize("obs, exp_fn", observables_and_exp_fns)
     def test_expansion(self, params, size, obs, exp_fn):
         """Test that the expansion works as expected."""
-        tape = make_tape(*params, obs)
+        ops = make_ops(*params)
+        expvals = [qml.expval(ob) for ob in obs]
+        tape = qml.tape.QuantumScript(ops, expvals)
         assert tape.batch_size == size
 
         tapes, fn = qml.transforms.broadcast_expand(tape)
@@ -107,33 +94,23 @@ def test_expansion(self, params, size, obs, exp_fn):
         result = fn(qml.execute(tapes, dev, None))
         expected = exp_fn(*params)
 
-        if len(tape.measurements) > 1 and size == 1:
-            expected = expected.T
-
         assert qml.math.allclose(result, expected)
 
-    @pytest.mark.parametrize("params, size", parameters_and_size)
+    @pytest.mark.parametrize("params", parameters)
     @pytest.mark.parametrize("obs, exp_fn", observables_and_exp_fns)
-    def test_expansion_qnode(self, params, size, obs, exp_fn):
+    def test_expansion_qnode(self, params, obs, exp_fn):
         """Test that the transform integrates correctly with the transform program"""
 
         @qml.transforms.broadcast_expand
         @qml.qnode(dev)
         def circuit(x, y, z, obs):
             qml.StatePrep(np.array([1, 0, 0, 0]), wires=[0, 1])
-            RX_broadcasted(x, wires=0)
-            qml.PauliY(0)
-            RX_broadcasted(y, wires=1)
-            RZ_broadcasted(z, wires=1)
-            qml.Hadamard(1)
+            _ = make_ops(x, y, z)
             return [qml.expval(ob) for ob in obs]
 
         result = circuit(*params, obs)
         expected = exp_fn(*params)
 
-        if len(obs) > 1 and size == 1:
-            expected = expected.T
-
         assert qml.math.allclose(result, expected)
 
     def test_state_prep(self):
@@ -174,54 +151,52 @@ def test_not_copied(self):
 
     @pytest.mark.autograd
     @pytest.mark.filterwarnings("ignore:Output seems independent of input")
-    @pytest.mark.parametrize("params, size", parameters_and_size)
+    @pytest.mark.parametrize("params", parameters)
     @pytest.mark.parametrize("obs, exp_fn", observables_and_exp_fns)
-    def test_autograd(self, params, size, obs, exp_fn):
+    @pytest.mark.parametrize("diff_method", ["parameter-shift", "backprop"])
+    def test_autograd(self, params, obs, exp_fn, diff_method):
         """Test that the expansion works with autograd and is differentiable."""
         params = tuple(pnp.array(p, requires_grad=True) for p in params)
 
+        @qml.transforms.broadcast_expand
+        @qml.qnode(dev, interface="autograd", diff_method=diff_method)
         def cost(*params):
-            tape = make_tape(*params, obs)
-            tapes, fn = qml.transforms.broadcast_expand(tape)
-            if len(tape.measurements) > 1:
-                return qml.math.stack(fn(qml.execute(tapes, dev, qml.gradients.param_shift)))
-            return fn(qml.execute(tapes, dev, qml.gradients.param_shift))
+            make_ops(*params)
+            return qml.math.stack([qml.expval(ob) for ob in obs])
 
         expected = exp_fn(*params)
 
-        if len(obs) > 1 and size == 1:
-            expected = expected.T
-
         assert qml.math.allclose(cost(*params), expected)
 
         jac = qml.jacobian(cost)(*params)
         exp_jac = qml.jacobian(exp_fn)(*params)
 
-        if len(obs) > 1 and size == 1:
-            new_exp_jac = [qml.math.flatten(el) for el in exp_jac]
-            new_jac = [qml.math.flatten(el) for el in jac]
-            assert all(qml.math.allclose(_jac, e_jac) for _jac, e_jac in zip(new_jac, new_exp_jac))
-        else:
-            assert all(qml.math.allclose(_jac, e_jac) for _jac, e_jac in zip(jac, exp_jac))
+        assert all(qml.math.allclose(_jac, e_jac) for _jac, e_jac in zip(jac, exp_jac))
 
     @pytest.mark.jax
-    @pytest.mark.parametrize("params, size", parameters_and_size)
+    @pytest.mark.parametrize("params", parameters)
     @pytest.mark.parametrize("obs, exp_fn", observables_and_exp_fns)
-    def test_jax(self, params, size, obs, exp_fn):
+    @pytest.mark.parametrize("use_jit", [True, False])
+    @pytest.mark.parametrize("diff_method", ["parameter-shift", "backprop"])
+    def test_jax(self, params, obs, exp_fn, use_jit, diff_method):
         """Test that the expansion works with jax and is differentiable."""
+        # pylint: disable=too-many-arguments
         import jax
 
+        jax.config.update("jax_enable_x64", True)
+
         params = tuple(jax.numpy.array(p) for p in params)
 
+        @qml.transforms.broadcast_expand
+        @qml.qnode(dev, interface="jax", diff_method=diff_method)
         def cost(*params):
-            tape = make_tape(*params, obs)
-            tapes, fn = qml.transforms.broadcast_expand(tape)
-            return fn(qml.execute(tapes, dev, qml.gradients.param_shift))
+            make_ops(*params)
+            return tuple(qml.expval(ob) for ob in obs)
 
-        expected = exp_fn(*params)
+        if use_jit:
+            cost = jax.jit(cost)
 
-        if len(obs) > 1 and size == 1:
-            expected = expected.T
+        expected = exp_fn(*params)
 
         assert qml.math.allclose(cost(*params), expected)
 
@@ -230,28 +205,26 @@ def cost(*params):
         exp_jac = jax.jacobian(exp_fn, argnums=[0, 1, 2])(*params)
 
         if len(obs) > 1:
-            exp_jac_0 = jax.jacobian(exp_fn_Z0, argnums=[0, 1, 2])(*params)
-            exp_jac_1 = jax.jacobian(exp_fn_Y1, argnums=[0, 1, 2])(*params)
-
-            assert all(qml.math.allclose(_jac, e_jac) for _jac, e_jac in zip(jac[0], exp_jac_0))
-            assert all(qml.math.allclose(_jac, e_jac) for _jac, e_jac in zip(jac[1], exp_jac_1))
-
+            assert all(qml.math.allclose(_jac, e_jac) for _jac, e_jac in zip(jac[0], exp_jac[0]))
+            assert all(qml.math.allclose(_jac, e_jac) for _jac, e_jac in zip(jac[1], exp_jac[1]))
         else:
             assert all(qml.math.allclose(_jac, e_jac) for _jac, e_jac in zip(jac, exp_jac))
 
+    @pytest.mark.slow
     @pytest.mark.tf
-    @pytest.mark.parametrize("params, size", parameters_and_size)
+    @pytest.mark.parametrize("params", parameters)
     @pytest.mark.parametrize("obs, exp_fn", observables_and_exp_fns)
-    def test_tf(self, params, size, obs, exp_fn):
+    def test_tf(self, params, obs, exp_fn):
         """Test that the expansion works with TensorFlow and is differentiable."""
         import tensorflow as tf
 
         params = tuple(tf.Variable(p, dtype=tf.float64) for p in params)
 
+        @qml.transforms.broadcast_expand
+        @qml.qnode(dev, interface="tensorflow")
         def cost(*params):
-            tape = make_tape(*params, obs)
-            tapes, fn = qml.transforms.broadcast_expand(tape)
-            return fn(qml.execute(tapes, dev, qml.gradients.param_shift))
+            make_ops(*params)
+            return tuple(qml.expval(ob) for ob in obs)
 
         with tf.GradientTape(persistent=True) as t:
             out = tf.stack(cost(*params))
@@ -260,43 +233,44 @@ def cost(*params):
         jac = t.jacobian(out, params)
         exp_jac = t.jacobian(exp, params)
 
-        if len(obs) > 1 and size == 1:
-            new_exp_jac = [qml.math.flatten(el) for el in exp_jac]
-            new_jac = [qml.math.flatten(el) for el in jac]
-            assert all(qml.math.allclose(_jac, e_jac) for _jac, e_jac in zip(new_jac, new_exp_jac))
-        else:
-            for _jac, e_jac in zip(jac, exp_jac):
-                if e_jac is None:
-                    assert qml.math.allclose(_jac, 0.0)
-                else:
-                    assert qml.math.allclose(_jac, e_jac)
+        for _jac, e_jac in zip(jac, exp_jac):
+            if e_jac is None:
+                assert qml.math.allclose(_jac, 0.0)
+            else:
+                assert qml.math.allclose(_jac, e_jac)
 
     @pytest.mark.torch
     @pytest.mark.filterwarnings("ignore:Output seems independent of input")
-    @pytest.mark.parametrize("params, size", parameters_and_size)
+    @pytest.mark.parametrize("params", parameters)
     @pytest.mark.parametrize("obs, exp_fn", observables_and_exp_fns)
-    def test_torch(self, params, size, obs, exp_fn):
+    @pytest.mark.parametrize("diff_method", ["parameter-shift", "backprop"])
+    def test_torch(self, params, obs, exp_fn, diff_method):
         """Test that the expansion works with torch and is differentiable."""
         import torch
 
-        torch_params = tuple(torch.tensor(p, requires_grad=True) for p in params)
+        torch_params = tuple(
+            torch.tensor(p, requires_grad=True, dtype=torch.float64) for p in params
+        )
         params = tuple(pnp.array(p, requires_grad=True) for p in params)
 
+        @qml.transforms.broadcast_expand
+        @qml.qnode(dev, interface="torch", diff_method=diff_method)
         def cost(*params):
-            tape = make_tape(*params, obs)
-            tapes, fn = qml.transforms.broadcast_expand(tape)
-            if len(tape.measurements) > 1:
-                return qml.math.stack(fn(qml.execute(tapes, dev, qml.gradients.param_shift)))
-            return fn(qml.execute(tapes, dev, qml.gradients.param_shift))
+            make_ops(*params)
+            return tuple(qml.expval(ob) for ob in obs)
 
+        res = cost(*torch_params)
         jac = torch.autograd.functional.jacobian(cost, torch_params)
-        exp_jac = qml.jacobian(exp_fn)(*params)
+        exp_jac = torch.autograd.functional.jacobian(exp_fn, torch_params)
 
-        if len(obs) > 1 and size == 1:
-            assert qml.math.allclose(cost(*torch_params), exp_fn(*params).T)
-            new_exp_jac = [qml.math.flatten(el) for el in exp_jac]
-            new_jac = [qml.math.flatten(el) for el in jac]
-            assert all(qml.math.allclose(_jac, e_jac) for _jac, e_jac in zip(new_jac, new_exp_jac))
+        if len(obs) > 1:
+            assert all(qml.math.allclose(r, e) for r, e in zip(res, exp_fn(*params)))
+            # Need to perform a transpose because the broadcast_expand transform pulls out the
+            # broadcasting axis, which the `exp_fn` does not do.
+            jac = tuple(
+                qml.math.stack([jac[i][j] for i in range(len(obs))]) for j in range(len(params))
+            )
         else:
-            assert qml.math.allclose(cost(*torch_params), exp_fn(*params))
-            assert all(qml.math.allclose(_jac, e_jac) for _jac, e_jac in zip(jac, exp_jac))
+            assert qml.math.allclose(res, exp_fn(*params))
+
+        assert all(qml.math.allclose(_jac, e_jac) for _jac, e_jac in zip(jac, exp_jac))