Repository: PennyLaneAI/pennylane
Pull request #3431 (2022-11-28 20:39:46)

[Return-types Fix] Update `batch_params` and `batch_input` transforms
**Context:**
Part of the ongoing effort to update the `batch_transforms` for the new return types system.

**Description of the Change:**
Update the `batch_params` and `batch_input` transforms to stack the results at the leaf nodes of the nested tuple structure. More work needs to be done to integrate parameter broadcasting.



github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3431?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3431](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3431?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (51dcc8e) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/79738a10001bff2f205ef36379bc99368edb748d?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (79738a1) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #3431   +/-   ##
=======================================
  Coverage   99.70%   99.70%           
=======================================
  Files         299      299           
  Lines       26637    26653   +16     
=======================================
+ Hits        26558    26574   +16     
  Misses         79       79           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3431?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/transforms/batch\_input.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3431/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvYmF0Y2hfaW5wdXQucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/transforms/batch\_params.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3431/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvYmF0Y2hfcGFyYW1zLnB5) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


Patch:
diff --git a/.pylintrc b/.pylintrc
index 21ccd82aa..5058594ec 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -2,7 +2,7 @@
 # A comma-separated list of package or module names from where C extensions may
 # be loaded. Extensions are loading into the active Python interpreter and may
 # run arbitrary code
-extension-pkg-whitelist=numpy,scipy,autograd,toml,appdir,autograd.numpy,autograd.numpy.linalg,semantic_version,torch,tensorflow,tensorflow.contrib,tensorflow.contrib.eager,LazyLoader,networkx,networkx.dag
+extension-pkg-whitelist=numpy,scipy,autograd,toml,appdir,autograd.numpy,autograd.numpy.linalg,autograd.numpy.builtins,semantic_version,torch,tensorflow,tensorflow.contrib,tensorflow.contrib.eager,LazyLoader,networkx,networkx.dag
 
 [TYPECHECK]
 
@@ -10,12 +10,12 @@ extension-pkg-whitelist=numpy,scipy,autograd,toml,appdir,autograd.numpy,autograd
 # (useful for modules/projects where namespaces are manipulated during runtime
 # and thus existing member attributes cannot be deduced by static analysis. It
 # supports qualified module names, as well as Unix pattern matching.
-ignored-modules=numpy,scipy,autograd,toml,appdir,autograd.numpy,autograd.numpy.linalg,semantic_version,torch,tensorflow,tensorflow.contrib,tensorflow.contrib.eager,LazyLoader,networkx,networkx.dag,math,pennylane.numpy
+ignored-modules=numpy,scipy,autograd,toml,appdir,autograd.numpy,autograd.numpy.linalg,autograd.numpy.builtins,semantic_version,torch,tensorflow,tensorflow.contrib,tensorflow.contrib.eager,LazyLoader,networkx,networkx.dag,math,pennylane.numpy
 
 # List of classes names for which member attributes should not be checked
 # (useful for classes with attributes dynamically set). This supports can work
 # with qualified names.
-ignored-classes=numpy,scipy,autograd,toml,appdir,autograd.numpy,autograd.numpy.linalg,semantic_version,torch,tensorflow,tensorflow.contrib,tensorflow.contrib.eager,LazyLoader,networkx,networkx.dag,math,pennylane.numpy,pennylane.numpy.random,pennylane.numpy.linalg,pennylane.operation,retworkx,kahypar
+ignored-classes=numpy,scipy,autograd,toml,appdir,autograd.numpy,autograd.numpy.linalg,autograd.numpy.builtins,semantic_version,torch,tensorflow,tensorflow.contrib,tensorflow.contrib.eager,LazyLoader,networkx,networkx.dag,math,pennylane.numpy,pennylane.numpy.random,pennylane.numpy.linalg,pennylane.numpy.builtins,pennylane.operation,retworkx,kahypar
 
 [MESSAGES CONTROL]
 
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index e389fbd04..6357d73ba 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -466,6 +466,9 @@
   `qml.qinfo.classical_fisher`, and `qml.qinfo.quantum_fisher` to support the new return types.
   [(#3449)](https://github.com/PennyLaneAI/pennylane/pull/3449)
 
+* Updated `qml.transforms.batch_params` and `qml.transforms.batch_input` to support the new return types
+  [(#3431)](https://github.com/PennyLaneAI/pennylane/pull/3431)
+
 * Updated `qml.transforms.cut_circuit` and `qml.transforms.cut_circuit_mc` to
   support the new return types.
   [(#3346)](https://github.com/PennyLaneAI/pennylane/pull/3346)
diff --git a/pennylane/interfaces/tensorflow.py b/pennylane/interfaces/tensorflow.py
index 1292a90df..b53ed9d95 100644
--- a/pennylane/interfaces/tensorflow.py
+++ b/pennylane/interfaces/tensorflow.py
@@ -112,7 +112,7 @@ def _jac_restructured(jacs, tapes):
     """
     start = 0
     jacs_nested = []
-    for i, tape in enumerate(tapes):
+    for tape in tapes:
         num_meas = len(tape.measurements)
         num_params = len(tape.trainable_params)
 
@@ -125,6 +125,7 @@ def _jac_restructured(jacs, tapes):
             tape_jacs = tape_jacs[0]
 
         jacs_nested.append(tape_jacs)
+        start += num_meas * num_params
 
     return tuple(jacs_nested)
 
diff --git a/pennylane/transforms/batch_input.py b/pennylane/transforms/batch_input.py
index b678b5ec9..f02b43697 100644
--- a/pennylane/transforms/batch_input.py
+++ b/pennylane/transforms/batch_input.py
@@ -19,6 +19,7 @@ import pennylane as qml
 from pennylane import numpy as np
 from pennylane.tape import QuantumTape
 from pennylane.transforms.batch_transform import batch_transform
+from pennylane.transforms.batch_params import _nested_stack
 
 
 @batch_transform
@@ -78,10 +79,13 @@ def batch_input(
     argnum_params = [all_parameters[i] for i in argnum]
 
     if any(num in tape.trainable_params for num in argnum):
-        raise ValueError(
-            "Batched inputs must be non-trainable. Please make sure that the parameters indexed by "
-            + "'argnum' are not marked as trainable."
-        )
+        # JAX arrays can't be marked as non-trainable, so don't raise this error
+        # if the interface is JAX
+        if qml.math.get_interface(*argnum_params) != "jax":
+            raise ValueError(
+                "Batched inputs must be non-trainable. Please make sure that the parameters indexed by "
+                + "'argnum' are not marked as trainable."
+            )
 
     batch_dims = np.unique([qml.math.shape(x)[0] for x in argnum_params])
     if len(batch_dims) != 1:
@@ -107,4 +111,10 @@ def batch_input(
         new_tape.set_parameters(params, trainable_only=False)
         output_tapes.append(new_tape)
 
-    return output_tapes, lambda x: qml.math.squeeze(qml.math.stack(x))
+    def processing_fn(res):
+        if qml.active_return():
+            return _nested_stack(res)
+
+        return qml.math.squeeze(qml.math.stack(res))
+
+    return output_tapes, processing_fn
diff --git a/pennylane/transforms/batch_params.py b/pennylane/transforms/batch_params.py
index 95af35b67..94d96ff6f 100644
--- a/pennylane/transforms/batch_params.py
+++ b/pennylane/transforms/batch_params.py
@@ -21,6 +21,22 @@ import pennylane as qml
 from .batch_transform import batch_transform
 
 
+def _nested_stack(res):
+    """
+    Given a list of identical nested tuple structures, stack the arrays at the leaves
+    """
+    # for some reason pylint thinks qml.numpy.builtins is a dict
+    # pylint: disable=no-member
+    if not isinstance(res[0], (tuple, qml.numpy.builtins.SequenceBox)):
+        return qml.math.stack(res)
+
+    stacked_results = []
+    for i in range(len(res[0])):
+        stacked_results.append(_nested_stack([r[i] for r in res]))
+
+    return tuple(stacked_results)
+
+
 @batch_transform
 def batch_params(tape, all_operations=False):
     """Transform a QNode to support an initial batch dimension
@@ -139,4 +155,10 @@ def batch_params(tape, all_operations=False):
         new_tape.set_parameters(p, trainable_only=not all_operations)
         output_tapes.append(new_tape)
 
-    return output_tapes, lambda x: qml.math.squeeze(qml.math.stack(x))
+    def processing_fn(res):
+        if qml.active_return():
+            return _nested_stack(res)
+
+        return qml.math.squeeze(qml.math.stack(res))
+
+    return output_tapes, processing_fn


Test patch:
diff --git a/tests/returntypes/transforms/test_batch_input_new.py b/tests/returntypes/transforms/test_batch_input_new.py
new file mode 100644
index 000000000..9d75c5ffb
--- /dev/null
+++ b/tests/returntypes/transforms/test_batch_input_new.py
@@ -0,0 +1,763 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Unit tests for the batch inputs transform.
+"""
+import pytest
+
+import pennylane as qml
+from pennylane import numpy as np
+
+
+def test_simple_circuit():
+    """Test that batching works for a simple circuit"""
+    dev = qml.device("default.qubit", wires=2)
+
+    @qml.batch_input(argnum=1)
+    @qml.qnode(dev, diff_method="parameter-shift")
+    def circuit(inputs, weights):
+        qml.RY(weights[0], wires=0)
+        qml.AngleEmbedding(inputs, wires=range(2), rotation="Y")
+        qml.RY(weights[1], wires=1)
+        return qml.expval(qml.PauliZ(1))
+
+    batch_size = 5
+    inputs = np.random.uniform(0, np.pi, (batch_size, 2), requires_grad=False)
+    weights = np.random.uniform(-np.pi, np.pi, (2,))
+
+    res = circuit(inputs, weights)
+    assert res.shape == (batch_size,)
+
+
+def test_simple_circuit_one_batch():
+    """Test that batching works for a simple circuit when the batch size is 1"""
+    dev = qml.device("default.qubit", wires=2)
+
+    @qml.batch_input(argnum=1)
+    @qml.qnode(dev, diff_method="parameter-shift")
+    def circuit(inputs, weights):
+        qml.RY(weights[0], wires=0)
+        qml.AngleEmbedding(inputs, wires=range(2), rotation="Y")
+        qml.RY(weights[1], wires=1)
+        return qml.expval(qml.PauliZ(1))
+
+    batch_size = 1
+    inputs = np.random.uniform(0, np.pi, (batch_size, 2), requires_grad=False)
+    weights = np.random.uniform(-np.pi, np.pi, (2,))
+
+    res = circuit(inputs, weights)
+    assert res.shape == (batch_size,)
+
+
+def test_circuit_non_param_operator_before_batched_operator():
+    """Test a circuit where a non-parametric operation is located before a batched operator."""
+    dev = qml.device("default.qubit", wires=2)
+
+    @qml.batch_input(argnum=0)
+    @qml.qnode(dev)
+    def circuit(input):
+        qml.CNOT(wires=[0, 1])
+        qml.RY(input, wires=1)
+        qml.RX(0.1, wires=0)
+        return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+    batch_size = 3
+
+    input = np.linspace(0.1, 0.5, batch_size, requires_grad=False)
+
+    res = circuit(input)
+
+    assert res.shape == (batch_size,)
+
+
+def test_value_error():
+    """Test if the batch_input raises relevant errors correctly"""
+
+    dev = qml.device("default.qubit", wires=2)
+
+    class Embedding(qml.AngleEmbedding):
+        """Variant of qml.AngleEmbedding that does not provide fixed
+        ``ndim_params`` in order to allow for the detection of inconsistent
+        batching in ``batch_input``."""
+
+        @property
+        def ndim_params(self):
+            return self._ndim_params
+
+    @qml.batch_input(argnum=[0, 2])
+    @qml.qnode(dev, diff_method="parameter-shift")
+    def circuit(input1, input2, weights):
+        Embedding(input1, wires=range(2), rotation="Y")
+        qml.RY(weights[0], wires=0)
+        qml.RY(input2[0], wires=0)
+        qml.RY(weights[1], wires=1)
+        return qml.expval(qml.PauliZ(1))
+
+    np.random.seed(42)
+    batch_size = 5
+    input1 = np.random.uniform(0, np.pi, (batch_size, 2), requires_grad=False)
+    input2 = np.random.uniform(0, np.pi, (4, 1), requires_grad=False)
+    weights = np.random.uniform(-np.pi, np.pi, (2,))
+
+    with pytest.raises(ValueError, match="Batch dimension for all gate arguments"):
+        res = circuit(input1, input2, weights)
+
+
+def test_batch_input_with_trainable_parameters_raises_error():
+    """Test that using the batch_input method with trainable parameters raises a ValueError."""
+    dev = qml.device("default.qubit", wires=2)
+
+    @qml.batch_input(argnum=0)
+    @qml.qnode(dev)
+    def circuit(input):
+        qml.RY(input, wires=1)
+        qml.CNOT(wires=[0, 1])
+        qml.RX(0.1, wires=0)
+        return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+    batch_size = 3
+
+    input = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
+
+    with pytest.raises(
+        ValueError,
+        match="Batched inputs must be non-trainable."
+        + " Please make sure that the parameters indexed by "
+        + "'argnum' are not marked as trainable.",
+    ):
+        circuit(input)
+
+
+def test_mottonenstate_preparation(mocker):
+    """Test that batching works for MottonenStatePreparation"""
+    dev = qml.device("default.qubit", wires=3)
+
+    @qml.batch_input(argnum=0)
+    @qml.qnode(dev)
+    def circuit(data, weights):
+        qml.templates.MottonenStatePreparation(data, wires=[0, 1, 2])
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+        return qml.probs(wires=[0, 1, 2])
+
+    batch_size = 3
+
+    # create a batched input statevector
+    data = np.random.random((batch_size, 2**3), requires_grad=False)
+    data /= np.linalg.norm(data, axis=1).reshape(-1, 1)  # normalize
+
+    # weights is not batched
+    weights = np.random.random((10, 3, 3), requires_grad=True)
+
+    spy = mocker.spy(circuit.device, "batch_execute")
+    res = circuit(data, weights)
+    assert res.shape == (batch_size, 2**3)
+    assert len(spy.call_args[0][0]) == batch_size
+
+    # check the results against individually executed circuits (no batching)
+    @qml.qnode(dev)
+    def circuit2(data, weights):
+        qml.templates.MottonenStatePreparation(data, wires=[0, 1, 2])
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+        return qml.probs(wires=[0, 1, 2])
+
+    indiv_res = []
+    for state in data:
+        indiv_res.append(circuit2(state, weights))
+    assert np.allclose(res, indiv_res)
+
+
+def test_multi_returns():
+    """Test that batching works for a simple circuit with multiple returns"""
+    dev = qml.device("default.qubit", wires=2)
+
+    @qml.batch_input(argnum=1)
+    @qml.qnode(dev, diff_method="parameter-shift")
+    def circuit(inputs, weights):
+        qml.RY(weights[0], wires=0)
+        qml.AngleEmbedding(inputs, wires=range(2), rotation="Y")
+        qml.RY(weights[1], wires=1)
+        return qml.expval(qml.PauliZ(1)), qml.probs(wires=[0, 1])
+
+    batch_size = 6
+    inputs = np.random.uniform(0, np.pi, (batch_size, 2), requires_grad=False)
+    weights = np.random.uniform(-np.pi, np.pi, (2,))
+
+    res = circuit(inputs, weights)
+    assert isinstance(res, tuple)
+    assert len(res) == 2
+
+    assert res[0].shape == (batch_size,)
+    assert res[1].shape == (batch_size, 4)
+
+
+def test_shot_vector():
+    """Test that batching works for a simple circuit with a shot vector"""
+    dev = qml.device("default.qubit", wires=2, shots=(100, (200, 3), 300))
+
+    @qml.batch_input(argnum=1)
+    @qml.qnode(dev, diff_method="parameter-shift")
+    def circuit(inputs, weights):
+        qml.RY(weights[0], wires=0)
+        qml.AngleEmbedding(inputs, wires=range(2), rotation="Y")
+        qml.RY(weights[1], wires=1)
+        return qml.probs(wires=[0, 1])
+
+    batch_size = 6
+    inputs = np.random.uniform(0, np.pi, (batch_size, 2), requires_grad=False)
+    weights = np.random.uniform(-np.pi, np.pi, (2,))
+
+    res = circuit(inputs, weights)
+
+    assert isinstance(res, tuple)
+    assert len(res) == 5
+    assert all(shot_res.shape == (batch_size, 4) for shot_res in res)
+
+
+def test_multi_returns_shot_vector():
+    """Test that batching works for a simple circuit with multiple returns
+    and with a shot vector"""
+    dev = qml.device("default.qubit", wires=2, shots=(100, (200, 3), 300))
+
+    @qml.batch_input(argnum=1)
+    @qml.qnode(dev, diff_method="parameter-shift")
+    def circuit(inputs, weights):
+        qml.RY(weights[0], wires=0)
+        qml.AngleEmbedding(inputs, wires=range(2), rotation="Y")
+        qml.RY(weights[1], wires=1)
+        return qml.expval(qml.PauliZ(1)), qml.probs(wires=[0, 1])
+
+    batch_size = 6
+    inputs = np.random.uniform(0, np.pi, (batch_size, 2), requires_grad=False)
+    weights = np.random.uniform(-np.pi, np.pi, (2,))
+
+    res = circuit(inputs, weights)
+
+    assert isinstance(res, tuple)
+    assert len(res) == 5
+    assert all(isinstance(shot_res, tuple) for shot_res in res)
+    assert all(len(shot_res) == 2 for shot_res in res)
+    assert all(shot_res[0].shape == (batch_size,) for shot_res in res)
+    assert all(shot_res[1].shape == (batch_size, 4) for shot_res in res)
+
+
+class TestDiffSingle:
+    """Test gradients for a single measurement"""
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"])
+    def test_autograd(self, diff_method, tol):
+        """Test derivatives when using autograd"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_input(argnum=0)
+        @qml.qnode(dev, diff_method=diff_method)
+        def circuit(input, x):
+            qml.RY(input, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.RX(x, wires=0)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        batch_size = 3
+
+        def cost(input, x):
+            return np.sum(circuit(input, x))
+
+        input = np.linspace(0.1, 0.5, batch_size, requires_grad=False)
+        x = np.array(0.1, requires_grad=True)
+
+        res = qml.grad(cost)(input, x)
+        expected = -np.sin(0.1) * sum(np.sin(input))
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"])
+    def test_jax(self, diff_method, tol):
+        """Test derivatives when using JAX"""
+        import jax
+        import jax.numpy as jnp
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_input(argnum=0)
+        @qml.qnode(dev, diff_method=diff_method, interface="jax")
+        def circuit(input, x):
+            qml.RY(input, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.RX(x, wires=0)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        batch_size = 3
+
+        def cost(input, x):
+            return jnp.sum(circuit(input, x))
+
+        input = jnp.linspace(0.1, 0.5, batch_size)
+        x = jnp.array(0.1)
+
+        res = jax.grad(cost, argnums=1)(input, x)
+        expected = -np.sin(0.1) * sum(np.sin(input))
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("diff_method", ["adjoint", "parameter-shift"])
+    def test_jax_jit(self, diff_method, tol):
+        """Test derivatives when using JAX"""
+        import jax
+        import jax.numpy as jnp
+
+        jax.config.update("jax_enable_x64", True)
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @jax.jit
+        @qml.batch_input(argnum=0)
+        @qml.qnode(dev, diff_method=diff_method, interface="jax-jit")
+        def circuit(input, x):
+            qml.RY(input, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.RX(x, wires=0)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        batch_size = 3
+
+        def cost(input, x):
+            return jnp.sum(circuit(input, x))
+
+        input = jnp.linspace(0.1, 0.5, batch_size)
+        x = jnp.array(0.1)
+
+        res = jax.grad(cost, argnums=1)(input, x)
+        expected = -np.sin(0.1) * sum(np.sin(input))
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"])
+    def test_torch(self, diff_method, tol):
+        """Test derivatives when using torch"""
+        import torch
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_input(argnum=0)
+        @qml.qnode(dev, diff_method=diff_method, interface="torch")
+        def circuit(input, x):
+            qml.RY(input, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.RX(x, wires=0)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        batch_size = 3
+
+        def cost(input, x):
+            return torch.sum(circuit(input, x))
+
+        input = torch.linspace(0.1, 0.5, batch_size, requires_grad=False)
+        x = torch.tensor(0.1, requires_grad=True)
+
+        loss = cost(input, x)
+        loss.backward()
+
+        res = x.grad
+        expected = -np.sin(0.1) * torch.sum(torch.sin(input))
+        assert qml.math.allclose(res, expected, atol=tol, rtol=0)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"])
+    def test_tf(self, diff_method, tol):
+        """Test derivatives when using TF"""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_input(argnum=0)
+        @qml.qnode(dev, diff_method=diff_method, interface="tf")
+        def circuit(input, x):
+            qml.RY(input, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.RX(x, wires=0)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        batch_size = 3
+        input = tf.Variable(np.linspace(0.1, 0.5, batch_size), trainable=False)
+        x = tf.Variable(0.1, trainable=True)
+
+        with tf.GradientTape() as tape:
+            loss = tf.reduce_sum(circuit(input, x))
+
+        res = tape.gradient(loss, x)
+        expected = -np.sin(0.1) * tf.reduce_sum(tf.sin(input))
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"])
+    def test_tf_autograph(self, diff_method, tol):
+        """Test derivatives when using TF and autograph"""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @tf.function
+        @qml.batch_input(argnum=0)
+        @qml.qnode(dev, diff_method=diff_method, interface="tf")
+        def circuit(input, x):
+            qml.RY(input, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.RX(x, wires=0)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        batch_size = 3
+        input = tf.Variable(np.linspace(0.1, 0.5, batch_size), trainable=False)
+        x = tf.Variable(0.1, trainable=True, dtype=tf.float64)
+
+        with tf.GradientTape() as tape:
+            loss = tf.reduce_sum(circuit(input, x))
+
+        res = tape.gradient(loss, x)
+        expected = -np.sin(0.1) * tf.reduce_sum(tf.sin(input))
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+
+class TestDiffMulti:
+    """Test gradients for multiple measurements"""
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
+    def test_autograd(self, diff_method, tol):
+        """Test derivatives when using autograd"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_input(argnum=0)
+        @qml.qnode(dev, diff_method=diff_method)
+        def circuit(input, x):
+            qml.RY(input, wires=0)
+            qml.RY(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        def cost(input, x):
+            res = circuit(input, x)
+            return qml.math.concatenate([qml.math.expand_dims(res[0], 1), res[1]], axis=1)
+
+        batch_size = 3
+        input = np.linspace(0.1, 0.5, batch_size, requires_grad=False)
+        x = np.array(0.1, requires_grad=True)
+
+        res = cost(input, x)
+        expected = qml.math.transpose(
+            qml.math.stack(
+                [
+                    np.cos(input + x),
+                    np.cos((input + x) / 2) ** 2,
+                    np.zeros_like(input),
+                    np.zeros_like(input),
+                    np.sin((input + x) / 2) ** 2,
+                ]
+            )
+        )
+        assert qml.math.allclose(res, expected, atol=tol)
+
+        grad = qml.jacobian(lambda x: cost(input, x))(x)
+        expected = qml.math.transpose(
+            qml.math.stack(
+                [
+                    -np.sin(input + x),
+                    -np.sin(input + x) / 2,
+                    np.zeros_like(input),
+                    np.zeros_like(input),
+                    np.sin(input + x) / 2,
+                ]
+            )
+        )
+        assert qml.math.allclose(grad, expected, atol=tol, rtol=0)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
+    def test_jax(self, diff_method, tol):
+        """Test derivatives when using JAX"""
+        import jax
+        import jax.numpy as jnp
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_input(argnum=0)
+        @qml.qnode(dev, diff_method=diff_method, interface="jax")
+        def circuit(input, x):
+            qml.RY(input, wires=0)
+            qml.RY(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        batch_size = 3
+        input = jnp.linspace(0.1, 0.5, batch_size)
+        x = jnp.array(0.1)
+
+        res = circuit(input, x)
+        expected = (
+            jnp.cos(input + x),
+            qml.math.transpose(
+                qml.math.stack(
+                    [
+                        np.cos((input + x) / 2) ** 2,
+                        np.zeros_like(input),
+                        np.zeros_like(input),
+                        np.sin((input + x) / 2) ** 2,
+                    ]
+                )
+            ),
+        )
+
+        assert isinstance(res, tuple)
+        assert len(res) == 2
+        for r, exp in zip(res, expected):
+            assert qml.math.allclose(r, exp, atol=tol)
+
+        grad = jax.jacobian(circuit, argnums=1)(input, x)
+        expected = (
+            -jnp.sin(input + x),
+            qml.math.transpose(
+                qml.math.stack(
+                    [
+                        -jnp.sin(input + x) / 2,
+                        jnp.zeros_like(input),
+                        jnp.zeros_like(input),
+                        jnp.sin(input + x) / 2,
+                    ]
+                )
+            ),
+        )
+
+        assert isinstance(grad, tuple)
+        assert len(grad) == 2
+        for g, exp in zip(grad, expected):
+            assert qml.math.allclose(g, exp, atol=tol, rtol=0)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("diff_method", ["parameter-shift"])
+    def test_jax_jit(self, diff_method, tol):
+        """Test derivatives when using JAX and jitting"""
+        import jax
+        import jax.numpy as jnp
+
+        jax.config.update("jax_enable_x64", True)
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @jax.jit
+        @qml.batch_input(argnum=0)
+        @qml.qnode(dev, diff_method=diff_method, interface="jax-jit")
+        def circuit(input, x):
+            qml.RY(input, wires=0)
+            qml.RY(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        batch_size = 3
+        input = jnp.linspace(0.1, 0.5, batch_size)
+        x = jnp.array(0.1)
+
+        res = circuit(input, x)
+        expected = (
+            jnp.cos(input + x),
+            qml.math.transpose(
+                qml.math.stack(
+                    [
+                        np.cos((input + x) / 2) ** 2,
+                        np.zeros_like(input),
+                        np.zeros_like(input),
+                        np.sin((input + x) / 2) ** 2,
+                    ]
+                )
+            ),
+        )
+
+        assert isinstance(res, tuple)
+        assert len(res) == 2
+        for r, exp in zip(res, expected):
+            assert qml.math.allclose(r, exp, atol=tol)
+
+        grad = jax.jacobian(circuit, argnums=1)(input, x)
+        expected = (
+            -jnp.sin(input + x),
+            qml.math.transpose(
+                qml.math.stack(
+                    [
+                        -jnp.sin(input + x) / 2,
+                        jnp.zeros_like(input),
+                        jnp.zeros_like(input),
+                        jnp.sin(input + x) / 2,
+                    ]
+                )
+            ),
+        )
+
+        assert isinstance(grad, tuple)
+        assert len(grad) == 2
+        for g, exp in zip(grad, expected):
+            assert qml.math.allclose(g, exp, atol=tol, rtol=0)
+
+    @pytest.mark.xfail(reason="Torch interface for new returns")
+    @pytest.mark.torch
+    @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
+    def test_torch(self, diff_method, tol):
+        """Test derivatives when using torch"""
+        import torch
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_input(argnum=0)
+        @qml.qnode(dev, diff_method=diff_method, interface="torch")
+        def circuit(input, x):
+            qml.RY(input, wires=0)
+            qml.RY(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        def cost(input, x):
+            res = circuit(input, x)
+            return qml.math.concatenate([qml.math.expand_dims(res[0], 1), res[1]], axis=1)
+
+        batch_size = 3
+        input = torch.tensor(np.linspace(0.1, 0.5, batch_size), requires_grad=False)
+        x = torch.tensor(0.1, requires_grad=True)
+
+        res = cost(input, x)
+        expected = qml.math.transpose(
+            qml.math.stack(
+                [
+                    torch.cos(input + x),
+                    torch.cos((input + x) / 2) ** 2,
+                    torch.zeros_like(input),
+                    torch.zeros_like(input),
+                    torch.sin((input + x) / 2) ** 2,
+                ]
+            )
+        )
+        assert qml.math.allclose(res, expected, atol=tol)
+
+        grad = torch.autograd.functional.jacobian(lambda x: cost(input, x), x)
+        expected = qml.math.transpose(
+            qml.math.stack(
+                [
+                    -torch.sin(input + x),
+                    -torch.sin(input + x) / 2,
+                    torch.zeros_like(input),
+                    torch.zeros_like(input),
+                    torch.sin(input + x) / 2,
+                ]
+            )
+        )
+        assert qml.math.allclose(grad, expected, atol=tol, rtol=0)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
+    def test_tf(self, diff_method, tol):
+        """Test derivatives when using TF"""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_input(argnum=0)
+        @qml.qnode(dev, diff_method=diff_method, interface="tf")
+        def circuit(input, x):
+            qml.RY(input, wires=0)
+            qml.RY(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        batch_size = 3
+        input = tf.Variable(np.linspace(0.1, 0.5, batch_size), trainable=False)
+        x = tf.Variable(0.1, trainable=True, dtype=tf.float64)
+
+        with tf.GradientTape() as tape:
+            res = circuit(input, x)
+            res = qml.math.concatenate([qml.math.expand_dims(res[0], 1), res[1]], axis=1)
+
+        expected = qml.math.transpose(
+            qml.math.stack(
+                [
+                    np.cos(input + x),
+                    np.cos((input + x) / 2) ** 2,
+                    np.zeros_like(input),
+                    np.zeros_like(input),
+                    np.sin((input + x) / 2) ** 2,
+                ]
+            )
+        )
+        assert qml.math.allclose(res, expected, atol=tol)
+
+        grad = tape.jacobian(res, x)
+        expected = qml.math.transpose(
+            qml.math.stack(
+                [
+                    -np.sin(input + x),
+                    -np.sin(input + x) / 2,
+                    np.zeros_like(input),
+                    np.zeros_like(input),
+                    np.sin(input + x) / 2,
+                ]
+            )
+        )
+        assert qml.math.allclose(grad, expected, atol=tol, rtol=0)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
+    def test_tf_autograph(self, diff_method, tol):
+        """Test derivatives when using TF and autograph"""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @tf.function
+        @qml.batch_input(argnum=0)
+        @qml.qnode(dev, diff_method=diff_method, interface="tf")
+        def circuit(input, x):
+            qml.RY(input, wires=0)
+            qml.RY(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        batch_size = 3
+        input = tf.Variable(np.linspace(0.1, 0.5, batch_size), trainable=False)
+        x = tf.Variable(0.1, trainable=True, dtype=tf.float64)
+
+        with tf.GradientTape() as tape:
+            res = circuit(input, x)
+            res = qml.math.concatenate([qml.math.expand_dims(res[0], 1), res[1]], axis=1)
+
+        expected = qml.math.transpose(
+            qml.math.stack(
+                [
+                    np.cos(input + x),
+                    np.cos((input + x) / 2) ** 2,
+                    np.zeros_like(input),
+                    np.zeros_like(input),
+                    np.sin((input + x) / 2) ** 2,
+                ]
+            )
+        )
+        assert qml.math.allclose(res, expected, atol=tol)
+
+        grad = tape.jacobian(res, x)
+        expected = qml.math.transpose(
+            qml.math.stack(
+                [
+                    -np.sin(input + x),
+                    -np.sin(input + x) / 2,
+                    np.zeros_like(input),
+                    np.zeros_like(input),
+                    np.sin(input + x) / 2,
+                ]
+            )
+        )
+        assert qml.math.allclose(grad, expected, atol=tol, rtol=0)
diff --git a/tests/returntypes/transforms/test_batch_params_new.py b/tests/returntypes/transforms/test_batch_params_new.py
new file mode 100644
index 000000000..206703dfa
--- /dev/null
+++ b/tests/returntypes/transforms/test_batch_params_new.py
@@ -0,0 +1,786 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Unit tests for the batch params transform.
+"""
+import functools
+import pytest
+
+import pennylane as qml
+from pennylane import numpy as np
+
+
+def test_simple_circuit(mocker):
+    """Test that batching works for a simple circuit"""
+    dev = qml.device("default.qubit", wires=3)
+
+    @qml.batch_params
+    @qml.qnode(dev)
+    def circuit(data, x, weights):
+        qml.templates.AmplitudeEmbedding(data, wires=[0, 1, 2], normalize=True)
+        qml.RX(x, wires=0)
+        qml.RY(0.2, wires=1)
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+        return qml.probs(wires=[0, 2])
+
+    batch_size = 5
+    data = np.random.random((batch_size, 8))
+    x = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
+    weights = np.ones((batch_size, 10, 3, 3), requires_grad=True)
+
+    spy = mocker.spy(circuit.device, "batch_execute")
+    res = circuit(data, x, weights)
+    assert res.shape == (batch_size, 4)
+    assert len(spy.call_args[0][0]) == batch_size
+
+
+def test_simple_circuit_one_batch(mocker):
+    """Test that batching works for a simple circuit when the batch size is 1"""
+    dev = qml.device("default.qubit", wires=3)
+
+    @qml.batch_params
+    @qml.qnode(dev)
+    def circuit(data, x, weights):
+        qml.templates.AmplitudeEmbedding(data, wires=[0, 1, 2], normalize=True)
+        qml.RX(x, wires=0)
+        qml.RY(0.2, wires=1)
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+        return qml.probs(wires=[0, 2])
+
+    batch_size = 1
+    data = np.random.random((batch_size, 8))
+    x = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
+    weights = np.ones((batch_size, 10, 3, 3), requires_grad=True)
+
+    spy = mocker.spy(circuit.device, "batch_execute")
+    res = circuit(data, x, weights)
+    assert res.shape == (batch_size, 4)
+    assert len(spy.call_args[0][0]) == batch_size
+
+
+def test_basic_entangler_layers(mocker):
+    """Test that batching works for BasicEngtanglerLayers"""
+    dev = qml.device("default.qubit", wires=2)
+
+    @qml.batch_params
+    @qml.qnode(dev)
+    def circuit(weights):
+        qml.templates.BasicEntanglerLayers(weights, wires=[0, 1])
+        qml.RY(0.2, wires=1)
+        return qml.probs(wires=[0, 1])
+
+    batch_size = 5
+    weights = np.random.random((batch_size, 2, 2))
+
+    spy = mocker.spy(circuit.device, "batch_execute")
+    res = circuit(weights)
+    assert res.shape == (batch_size, 4)
+    assert len(spy.call_args[0][0]) == batch_size
+
+
+def test_angle_embedding(mocker):
+    """Test that batching works for AngleEmbedding"""
+    dev = qml.device("default.qubit", wires=3)
+
+    @qml.batch_params
+    @qml.qnode(dev)
+    def circuit(data):
+        qml.templates.AngleEmbedding(data, wires=[0, 1, 2])
+        qml.RY(0.2, wires=1)
+        return qml.probs(wires=[0, 2])
+
+    batch_size = 5
+    data = np.random.random((batch_size, 3))
+
+    spy = mocker.spy(circuit.device, "batch_execute")
+    res = circuit(data)
+    assert res.shape == (batch_size, 4)
+    assert len(spy.call_args[0][0]) == batch_size
+
+
+def test_mottonenstate_preparation(mocker):
+    """Test that batching works for MottonenStatePreparation"""
+    dev = qml.device("default.qubit", wires=3)
+
+    @qml.batch_params
+    @qml.qnode(dev)
+    def circuit(data, weights):
+        qml.templates.MottonenStatePreparation(data, wires=[0, 1, 2])
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+        return qml.probs(wires=[0, 1, 2])
+
+    batch_size = 3
+
+    # create a batched input statevector
+    data = np.random.random((batch_size, 2**3))
+    data /= np.linalg.norm(data, axis=1).reshape(-1, 1)  # normalize
+    weights = np.random.random((batch_size, 10, 3, 3))
+
+    spy = mocker.spy(circuit.device, "batch_execute")
+    res = circuit(data, weights)
+    assert res.shape == (batch_size, 2**3)
+    assert len(spy.call_args[0][0]) == batch_size
+
+    # check the results against individually executed circuits (no batching)
+    @qml.qnode(dev)
+    def circuit2(data, weights):
+        qml.templates.MottonenStatePreparation(data, wires=[0, 1, 2])
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+        return qml.probs(wires=[0, 1, 2])
+
+    indiv_res = []
+    for state, weight in zip(data, weights):
+        indiv_res.append(circuit2(state, weight))
+    assert np.allclose(res, indiv_res)
+
+
+def test_basis_state_preparation(mocker):
+    """Test that batching works for BasisStatePreparation"""
+    dev = qml.device("default.qubit", wires=4)
+
+    @qml.batch_params
+    @qml.qnode(dev)
+    def circuit(data, weights):
+        qml.templates.BasisStatePreparation(data, wires=[0, 1, 2, 3])
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2, 3])
+        return qml.probs(wires=[0, 1, 2, 3])
+
+    batch_size = 5
+
+    # create random batched basis states
+    data = np.random.randint(2, size=(batch_size, 4))
+    weights = np.random.random((batch_size, 10, 4, 3))
+
+    spy = mocker.spy(circuit.device, "batch_execute")
+    res = circuit(data, weights)
+    assert res.shape == (batch_size, 2**4)
+    assert len(spy.call_args[0][0]) == batch_size
+
+    # check the results against individually executed circuits (no batching)
+    @qml.qnode(dev)
+    def circuit2(data, weights):
+        qml.templates.BasisStatePreparation(data, wires=[0, 1, 2, 3])
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2, 3])
+        return qml.probs(wires=[0, 1, 2, 3])
+
+    indiv_res = []
+    for state, weight in zip(data, weights):
+        indiv_res.append(circuit2(state, weight))
+    assert np.allclose(res, indiv_res)
+
+
+def test_multi_returns():
+    """Test that batching works for a simple circuit with multiple returns"""
+    dev = qml.device("default.qubit", wires=3)
+
+    @qml.batch_params
+    @qml.qnode(dev)
+    def circuit(data, x, weights):
+        qml.templates.AmplitudeEmbedding(data, wires=[0, 1, 2], normalize=True)
+        qml.RX(x, wires=0)
+        qml.RY(0.2, wires=1)
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+        return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 2])
+
+    batch_size = 6
+    data = np.random.random((batch_size, 8))
+    x = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
+    weights = np.ones((batch_size, 10, 3, 3), requires_grad=True)
+
+    res = circuit(data, x, weights)
+
+    assert isinstance(res, tuple)
+    assert len(res) == 2
+
+    assert res[0].shape == (batch_size,)
+    assert res[1].shape == (batch_size, 4)
+
+
+def test_shot_vector():
+    """Test that batching works for a simple circuit with a shot vector"""
+    dev = qml.device("default.qubit", wires=3, shots=(100, (200, 3), 300))
+
+    @qml.batch_params
+    @qml.qnode(dev)
+    def circuit(data, x, weights):
+        qml.templates.AmplitudeEmbedding(data, wires=[0, 1, 2], normalize=True)
+        qml.RX(x, wires=0)
+        qml.RY(0.2, wires=1)
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+        return qml.probs(wires=[0, 2])
+
+    batch_size = 6
+    data = np.random.random((batch_size, 8))
+    x = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
+    weights = np.ones((batch_size, 10, 3, 3), requires_grad=True)
+
+    res = circuit(data, x, weights)
+
+    assert isinstance(res, tuple)
+    assert len(res) == 5
+    assert all(shot_res.shape == (batch_size, 4) for shot_res in res)
+
+
+def test_multi_returns_shot_vector():
+    """Test that batching works for a simple circuit with multiple returns
+    and with a shot vector"""
+    dev = qml.device("default.qubit", wires=3, shots=(100, (200, 3), 300))
+
+    @qml.batch_params
+    @qml.qnode(dev)
+    def circuit(data, x, weights):
+        qml.templates.AmplitudeEmbedding(data, wires=[0, 1, 2], normalize=True)
+        qml.RX(x, wires=0)
+        qml.RY(0.2, wires=1)
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+        return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 2])
+
+    batch_size = 6
+    data = np.random.random((batch_size, 8))
+    x = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
+    weights = np.ones((batch_size, 10, 3, 3), requires_grad=True)
+
+    res = circuit(data, x, weights)
+
+    assert isinstance(res, tuple)
+    assert len(res) == 5
+    assert all(isinstance(shot_res, tuple) for shot_res in res)
+    assert all(len(shot_res) == 2 for shot_res in res)
+    assert all(shot_res[0].shape == (batch_size,) for shot_res in res)
+    assert all(shot_res[1].shape == (batch_size, 4) for shot_res in res)
+
+
+class TestDiffSingle:
+    """Test gradients for a single measurement"""
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"])
+    def test_autograd(self, diff_method, tol):
+        """Test derivatives when using autograd"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_params
+        @qml.qnode(dev, diff_method=diff_method)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            qml.RY(0.1, wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        def cost(x):
+            return np.sum(circuit(x))
+
+        batch_size = 3
+        x = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
+
+        res = qml.grad(cost)(x)
+        expected = -np.sin(0.1) * np.sin(x)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"])
+    def test_jax(self, diff_method, tol):
+        """Test derivatives when using JAX."""
+        import jax
+
+        jnp = jax.numpy
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_params
+        @qml.qnode(dev, interface="jax", diff_method=diff_method)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            qml.RY(0.1, wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        def cost(x):
+            return jnp.sum(circuit(x))
+
+        batch_size = 3
+        x = jnp.linspace(0.1, 0.5, batch_size)
+
+        res = jax.grad(cost)(x)
+        expected = -np.sin(0.1) * np.sin(x)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("diff_method", ["adjoint", "parameter-shift"])
+    @pytest.mark.parametrize("interface", ["jax", "jax-jit"])
+    def test_jax_jit(self, diff_method, interface, tol):
+        """Test derivatives when using JAX and JIT."""
+        import jax
+        import jax.numpy as jnp
+
+        jax.config.update("jax_enable_x64", True)
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_params
+        @qml.qnode(dev, interface=interface, diff_method=diff_method)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            qml.RY(0.1, wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        @jax.jit
+        def cost(x):
+            return jnp.sum(circuit(x))
+
+        batch_size = 3
+        x = jnp.linspace(0.1, 0.5, batch_size)
+
+        res = jax.grad(cost)(x)
+        expected = -np.sin(0.1) * np.sin(x)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"])
+    def test_torch(self, diff_method, tol):
+        """Test derivatives when using Torch"""
+        import torch
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_params
+        @qml.qnode(dev, interface="torch", diff_method=diff_method)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            qml.RY(0.1, wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        def cost(x):
+            return torch.sum(circuit(x))
+
+        batch_size = 3
+        x = torch.linspace(0.1, 0.5, batch_size, requires_grad=True)
+
+        loss = cost(x)
+        loss.backward()
+
+        res = x.grad
+        expected = -np.sin(0.1) * torch.sin(x)
+        assert torch.allclose(res, expected, atol=tol, rtol=0)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"])
+    def test_tf(self, diff_method, tol):
+        """Test derivatives when using TF"""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_params
+        @qml.qnode(dev, interface="tf", diff_method=diff_method)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            qml.RY(0.1, wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        def cost(x):
+            return tf.reduce_sum(circuit(x))
+
+        batch_size = 3
+        x = tf.Variable(tf.linspace(0.1, 0.5, batch_size))
+
+        with tf.GradientTape() as tape:
+            loss = cost(x)
+
+        res = tape.gradient(loss, x)
+        expected = -np.sin(0.1) * tf.sin(x)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    @pytest.mark.tf
+    def test_tf_autograph(self, tol):
+        """Test derivatives when using TF and autograph"""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_params
+        @qml.qnode(dev, interface="tf", diff_method="backprop")
+        def circuit(x):
+            qml.RX(x, wires=0)
+            qml.RY(0.1, wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        @tf.function
+        def cost(x):
+            return tf.reduce_sum(circuit(x))
+
+        batch_size = 3
+        x = tf.Variable(tf.linspace(0.1, 0.5, batch_size))
+
+        with tf.GradientTape() as tape:
+            loss = cost(x)
+
+        res = tape.gradient(loss, x)
+        expected = -np.sin(0.1) * tf.sin(x)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+
+class TestDiffMulti:
+    """Test gradients for multiple measurements"""
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
+    def test_autograd(self, diff_method, tol):
+        """Test derivatives when using autograd"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_params
+        @qml.qnode(dev, diff_method=diff_method)
+        def circuit(x):
+            qml.RY(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        def cost(x):
+            res = circuit(x)
+            return qml.math.concatenate([qml.math.expand_dims(res[0], 1), res[1]], axis=1)
+
+        batch_size = 3
+        x = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
+
+        res = cost(x)
+        expected = qml.math.transpose(
+            qml.math.stack(
+                [
+                    np.cos(x),
+                    np.cos(x / 2) ** 2,
+                    np.zeros_like(x),
+                    np.zeros_like(x),
+                    np.sin(x / 2) ** 2,
+                ]
+            )
+        )
+        assert qml.math.allclose(res, expected, atol=tol)
+
+        grad = qml.jacobian(cost)(x)
+        expected = qml.math.stack(
+            [-np.sin(x), -np.sin(x) / 2, np.zeros_like(x), np.zeros_like(x), np.sin(x) / 2]
+        ) * qml.math.expand_dims(np.eye(batch_size), 1)
+
+        assert np.allclose(grad, expected, atol=tol, rtol=0)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
+    def test_jax(self, diff_method, tol):
+        """Test derivatives when using JAX"""
+        import jax
+        import jax.numpy as jnp
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_params
+        @qml.qnode(dev, diff_method=diff_method, interface="jax")
+        def circuit(x):
+            qml.RY(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        batch_size = 3
+        x = jnp.linspace(0.1, 0.5, batch_size)
+
+        res = circuit(x)
+        expected = (
+            jnp.cos(x),
+            qml.math.transpose(
+                qml.math.stack(
+                    [jnp.cos(x / 2) ** 2, jnp.zeros_like(x), jnp.zeros_like(x), jnp.sin(x / 2) ** 2]
+                )
+            ),
+        )
+
+        assert isinstance(res, tuple)
+        assert len(res) == 2
+        for r, exp in zip(res, expected):
+            assert qml.math.allclose(r, exp, atol=tol)
+
+        grad = jax.jacobian(circuit)(x)
+        expected = (
+            -np.sin(x) * np.eye(batch_size),
+            qml.math.stack([-np.sin(x) / 2, np.zeros_like(x), np.zeros_like(x), np.sin(x) / 2])
+            * qml.math.expand_dims(np.eye(batch_size), 1),
+        )
+
+        assert isinstance(grad, tuple)
+        assert len(grad) == 2
+        for g, exp in zip(grad, expected):
+            assert qml.math.allclose(g, exp, atol=tol, rtol=0)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("diff_method", ["parameter-shift"])
+    def test_jax_jit(self, diff_method, tol):
+        """Test derivatives when using JAX"""
+        import jax
+        import jax.numpy as jnp
+
+        jax.config.update("jax_enable_x64", True)
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @jax.jit
+        @qml.batch_params
+        @qml.qnode(dev, diff_method=diff_method, interface="jax-jit")
+        def circuit(x):
+            qml.RY(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        batch_size = 3
+        x = jnp.linspace(0.1, 0.5, batch_size)
+
+        res = circuit(x)
+        expected = (
+            jnp.cos(x),
+            qml.math.transpose(
+                qml.math.stack(
+                    [jnp.cos(x / 2) ** 2, jnp.zeros_like(x), jnp.zeros_like(x), jnp.sin(x / 2) ** 2]
+                )
+            ),
+        )
+
+        assert isinstance(res, tuple)
+        assert len(res) == 2
+        for r, exp in zip(res, expected):
+            assert qml.math.allclose(r, exp, atol=tol)
+
+        grad = jax.jacobian(circuit)(x)
+        expected = (
+            -np.sin(x) * np.eye(batch_size),
+            qml.math.stack([-np.sin(x) / 2, np.zeros_like(x), np.zeros_like(x), np.sin(x) / 2])
+            * qml.math.expand_dims(np.eye(batch_size), 1),
+        )
+
+        assert isinstance(grad, tuple)
+        assert len(grad) == 2
+        for g, exp in zip(grad, expected):
+            assert qml.math.allclose(g, exp, atol=tol, rtol=0)
+
+    @pytest.mark.xfail(reason="Torch interface for new returns")
+    @pytest.mark.torch
+    @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
+    def test_torch(self, diff_method, tol):
+        """Test derivatives when using torch"""
+        import torch
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_params
+        @qml.qnode(dev, diff_method=diff_method, interface="torch")
+        def circuit(x):
+            qml.RY(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        def cost(x):
+            res = circuit(x)
+            return qml.math.concatenate([qml.math.expand_dims(res[0], 1), res[1]], axis=1)
+
+        batch_size = 3
+        x = torch.tensor(np.linspace(0.1, 0.5, batch_size), requires_grad=True)
+
+        res = cost(x)
+        expected = qml.math.transpose(
+            qml.math.stack(
+                [
+                    torch.cos(x),
+                    torch.cos(x / 2) ** 2,
+                    torch.zeros_like(x),
+                    torch.zeros_like(x),
+                    torch.sin(x / 2) ** 2,
+                ]
+            )
+        )
+        assert qml.math.allclose(res, expected, atol=tol)
+
+        grad = torch.autograd.functional.jacobian(cost, x)
+        expected = qml.math.stack(
+            [
+                -torch.sin(x),
+                -torch.sin(x) / 2,
+                torch.zeros_like(x),
+                torch.zeros_like(x),
+                torch.sin(x) / 2,
+            ]
+        ) * qml.math.expand_dims(torch.eye(batch_size), 1)
+
+        assert qml.math.allclose(grad, expected, atol=tol, rtol=0)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
+    def test_tf(self, diff_method, tol):
+        """Test derivatives when using TF"""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_params
+        @qml.qnode(dev, diff_method=diff_method, interface="tf")
+        def circuit(x):
+            qml.RY(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        batch_size = 3
+        x = tf.Variable(np.linspace(0.1, 0.5, batch_size))
+
+        with tf.GradientTape() as tape:
+            res = circuit(x)
+            res = qml.math.concatenate([qml.math.expand_dims(res[0], 1), res[1]], axis=1)
+
+        expected = qml.math.transpose(
+            qml.math.stack(
+                [
+                    np.cos(x),
+                    np.cos(x / 2) ** 2,
+                    np.zeros_like(x),
+                    np.zeros_like(x),
+                    np.sin(x / 2) ** 2,
+                ]
+            )
+        )
+        assert qml.math.allclose(res, expected, atol=tol)
+
+        grad = tape.jacobian(res, x)
+        expected = qml.math.stack(
+            [-np.sin(x), -np.sin(x) / 2, np.zeros_like(x), np.zeros_like(x), np.sin(x) / 2]
+        ) * qml.math.expand_dims(np.eye(batch_size), 1)
+
+        assert np.allclose(grad, expected, atol=tol, rtol=0)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
+    def test_tf_autograph(self, diff_method, tol):
+        """Test derivatives when using TF"""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @tf.function
+        @qml.batch_params
+        @qml.qnode(dev, diff_method=diff_method, interface="tf")
+        def circuit(x):
+            qml.RY(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        batch_size = 3
+        x = tf.Variable(np.linspace(0.1, 0.5, batch_size))
+
+        with tf.GradientTape() as tape:
+            res = circuit(x)
+            res = qml.math.concatenate([qml.math.expand_dims(res[0], 1), res[1]], axis=1)
+
+        expected = qml.math.transpose(
+            qml.math.stack(
+                [
+                    np.cos(x),
+                    np.cos(x / 2) ** 2,
+                    np.zeros_like(x),
+                    np.zeros_like(x),
+                    np.sin(x / 2) ** 2,
+                ]
+            )
+        )
+        assert qml.math.allclose(res, expected, atol=tol)
+
+        grad = tape.jacobian(res, x)
+        expected = qml.math.stack(
+            [-np.sin(x), -np.sin(x) / 2, np.zeros_like(x), np.zeros_like(x), np.sin(x) / 2]
+        ) * qml.math.expand_dims(np.eye(batch_size), 1)
+
+        assert np.allclose(grad, expected, atol=tol, rtol=0)
+
+
+def test_all_operations(mocker):
+    """Test that a batch dimension can be added to all operations"""
+    dev = qml.device("default.qubit", wires=3)
+
+    @functools.partial(qml.batch_params, all_operations=True)
+    @qml.qnode(dev)
+    def circuit(x, weights):
+        qml.RX(x, wires=0)
+        qml.RY([0.2, 0.3, 0.3], wires=1)
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+        return qml.probs(wires=[0, 2])
+
+    batch_size = 3
+    x = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
+    weights = np.ones((batch_size, 10, 3, 3), requires_grad=False)
+
+    spy = mocker.spy(circuit.device, "batch_execute")
+    res = circuit(x, weights)
+    assert res.shape == (batch_size, 4)
+    assert len(spy.call_args[0][0]) == batch_size
+
+
+def test_unbatched_parameter():
+    """Test that an exception is raised if a parameter
+    is not batched"""
+
+    dev = qml.device("default.qubit", wires=1)
+
+    @qml.batch_params
+    @qml.qnode(dev)
+    def circuit(x, y):
+        qml.RY(x, wires=[0])
+        qml.RX(y, wires=[0])
+        return qml.expval(qml.PauliZ(0))
+
+    x = np.array([0.3, 0.4, 0.5])
+    y = np.array(0.2)
+
+    with pytest.raises(ValueError, match="0.2 has incorrect batch dimension"):
+        circuit(x, y)
+
+
+def test_initial_unbatched_parameter():
+    """Test that an exception is raised if an initial parameter
+    is not batched"""
+
+    dev = qml.device("default.qubit", wires=1)
+
+    @qml.batch_params
+    @qml.qnode(dev)
+    def circuit(x, y):
+        qml.RY(x, wires=[0])
+        qml.RX(y, wires=[0])
+        return qml.expval(qml.PauliZ(0))
+
+    x = np.array(0.2)
+    y = np.array([0.3, 0.4, 0.5])
+
+    with pytest.raises(ValueError, match="Parameter 0.2 does not contain a batch"):
+        circuit(x, y)
+
+
+def test_no_batch_param_error():
+    """Test that the right error is thrown when there is nothing to batch"""
+    dev = qml.device("default.qubit", wires=1)
+
+    @qml.batch_params
+    @qml.qnode(dev)
+    def circuit(x):
+        qml.RY(x, wires=0)
+        return qml.expval(qml.PauliZ(0))
+
+    x = [0.2, 0.6, 3]
+    with pytest.raises(ValueError, match="There are no operations to transform"):
+        circuit(x)