Repository: PennyLaneAI/pennylane
Pull request #1634 (2021-09-09 05:55:11)

Differentiable batch execute using JAX using `host_callback`
**Context:**

Extends the batch execution capabilities to support JAX.

**Description of the Change:**

Adds the differentiable batch execution pipeline using JAX.

It is similar to #1573, though it uses the `host_callback` jittable JAX function with certain constraints. This capability supports first derivatives and scalar functions only similar to the current JAX interface.

Includes forward mode differentiation when the gradient has been computed on the forward pass.

**Benefits:**

Differentiable batch execution pipeline with the capabilities of the current JAX interface.

**Possible Drawbacks:**

The `test_matrix_parameter` test case seems to involve differentiation that is not supported by the JAX interface. Just as the current interface, the differentiable batch version raises an error akin to that raised by a current example:

```python
import jax.numpy as jnp
import pennylane as qml
import jax

a = jnp.array([0.1])
U = jnp.array([[ 0.99875026, -0.04997917],
       [ 0.04997917,  0.99875026]])
dev = qml.device("default.qubit.jax", wires=2)

@qml.qnode(dev, interface='jax', diff_method='parameter-shift')
def circuit(U):
    qml.QubitUnitary(U, wires=0)
    return qml.expval(qml.PauliZ(0))

jax.grad(circuit)(U)
```
```
ValueError: Cannot differentiate with respect to parameter(s) {0}
```

**Related GitHub Issues:**
N/A

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [.github/CHANGELOG.md](/PennyLaneAI/pennylane/blob/master/.github/CHANGELOG.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/1634?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1634](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1634?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (8627606) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/16aac9ff7092a94324b27d6b0dfae6f4af43a092?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (16aac9f) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1634/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1634?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@           Coverage Diff            @@
##           master    #1634    +/-   ##
========================================
  Coverage   99.19%   99.19%            
========================================
  Files         199      200     +1     
  Lines       14857    14969   +112     
========================================
+ Hits        14737    14849   +112     
  Misses        120      120            
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1634?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/interfaces/batch/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1634/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1634/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4LnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/tape/tape.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1634/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RhcGUvdGFwZS5weQ==) | `98.91% <100.00%> (+<0.01%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1634?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1634?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [16aac9f...8627606](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1634?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 073c577e3..2c676ad6b 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,6 +4,39 @@
 
 <h3>New features since last release</h3>
 
+* Support for differentiable execution of batches of circuits has been
+  extended to the JAX interface for scalar functions, via the beta
+  `pennylane.interfaces.batch` module.
+  [(#1634)](https://github.com/PennyLaneAI/pennylane/pull/1634)
+
+  For example using the `execute` function from the `pennylane.interfaces.batch` module:
+
+  ```python
+  from pennylane.interfaces.batch import execute
+
+  def cost_fn(x):
+      with qml.tape.JacobianTape() as tape1:
+          qml.RX(x[0], wires=[0])
+          qml.RY(x[1], wires=[1])
+          qml.CNOT(wires=[0, 1])
+          qml.var(qml.PauliZ(0) @ qml.PauliX(1))
+
+      with qml.tape.JacobianTape() as tape2:
+          qml.RX(x[0], wires=0)
+          qml.RY(x[0], wires=1)
+          qml.CNOT(wires=[0, 1])
+          qml.probs(wires=1)
+
+      result = execute(
+        [tape1, tape2], dev,
+        gradient_fn=qml.gradients.param_shift,
+        interface="autograd"
+      )
+      return (result[0] + result[1][0, 0])[0]
+
+  res = jax.grad(cost_fn)(params)
+  ```
+
 * The unitary matrix corresponding to a quantum circuit can now be created using the new
   `get_unitary_matrix()` transform.
   [(#1609)](https://github.com/PennyLaneAI/pennylane/pull/1609)
@@ -263,4 +296,5 @@
 
 This release contains contributions from (in alphabetical order):
 
-Utkarsh Azad, Olivia Di Matteo, Andrew Gardhouse, Josh Izaac, Christina Lee, Ingrid Strandberg, David Wierichs.
+Utkarsh Azad, Olivia Di Matteo, Andrew Gardhouse, Josh Izaac, Christina Lee,
+Ingrid Strandberg, Antal Száva, David Wierichs.
diff --git a/pennylane/interfaces/batch/__init__.py b/pennylane/interfaces/batch/__init__.py
index ee9953655..d5750b47f 100644
--- a/pennylane/interfaces/batch/__init__.py
+++ b/pennylane/interfaces/batch/__init__.py
@@ -29,6 +29,7 @@ import pennylane as qml
 INTERFACE_NAMES = {
     "NumPy": (None,),
     "Autograd": ("autograd", "numpy"),  # for backwards compatibility
+    "JAX": ("jax", "JAX"),
     "PyTorch": ("torch", "pytorch"),
     "TensorFlow": ("tf", "tensorflow"),
 }
@@ -193,7 +194,7 @@ def execute(
     gradient_kwargs=None,
     cache=True,
     cachesize=10000,
-    max_diff=2,
+    max_diff=1,
     override_shots=False,
     expand_fn="device",
     max_expansion=10,
@@ -358,12 +359,13 @@ def execute(
             from .tensorflow import execute as _execute
         elif interface in INTERFACE_NAMES["PyTorch"]:
             from .torch import execute as _execute
+        elif interface in INTERFACE_NAMES["JAX"]:
+            from .jax import execute as _execute
         else:
             raise ValueError(
                 f"Unknown interface {interface}. Supported "
                 f"interfaces are {SUPPORTED_INTERFACES}"
             )
-
     except ImportError as e:
         interface_name = [k for k, v in INTERFACE_NAMES.items() if interface in v][0]
 
diff --git a/pennylane/interfaces/batch/jax.py b/pennylane/interfaces/batch/jax.py
new file mode 100644
index 000000000..a3a8462da
--- /dev/null
+++ b/pennylane/interfaces/batch/jax.py
@@ -0,0 +1,269 @@
+# 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.
+"""
+This module contains functions for adding the JAX interface
+to a PennyLane Device class.
+"""
+
+# pylint: disable=too-many-arguments
+import jax
+import jax.numpy as jnp
+from jax.experimental import host_callback
+
+import numpy as np
+import pennylane as qml
+
+dtype = jnp.float64
+
+
+def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_diff=1):
+    """Execute a batch of tapes with JAX parameters on a device.
+
+    Args:
+        tapes (Sequence[.QuantumTape]): batch of tapes to execute
+        device (.Device): Device to use to execute the batch of tapes.
+            If the device does not provide a ``batch_execute`` method,
+            by default the tapes will be executed in serial.
+        execute_fn (callable): The execution function used to execute the tapes
+            during the forward pass. This function must return a tuple ``(results, jacobians)``.
+            If ``jacobians`` is an empty list, then ``gradient_fn`` is used to
+            compute the gradients during the backwards pass.
+        gradient_kwargs (dict): dictionary of keyword arguments to pass when
+            determining the gradients of tapes
+        gradient_fn (callable): the gradient function to use to compute quantum gradients
+        _n (int): a positive integer used to track nesting of derivatives, for example
+            if the nth-order derivative is requested.
+        max_diff (int): If ``gradient_fn`` is a gradient transform, this option specifies
+            the maximum order of derivatives to support. Increasing this value allows
+            for higher order derivatives to be extracted, at the cost of additional
+            (classical) computational overhead during the backwards pass.
+
+    Returns:
+        list[list[float]]: A nested list of tape results. Each element in
+        the returned list corresponds in order to the provided tapes.
+    """
+    if max_diff > 1:
+        raise ValueError("The JAX interface only supports first order derivatives.")
+
+    for tape in tapes:
+        # set the trainable parameters
+        params = tape.get_parameters(trainable_only=False)
+        tape.trainable_params = qml.math.get_trainable_indices(params)
+
+    parameters = tuple(list(t.get_parameters()) for t in tapes)
+    if gradient_fn is None:
+        return _execute_with_fwd(
+            parameters,
+            tapes=tapes,
+            device=device,
+            execute_fn=execute_fn,
+            gradient_kwargs=gradient_kwargs,
+            _n=_n,
+        )
+
+    return _execute(
+        parameters,
+        tapes=tapes,
+        device=device,
+        execute_fn=execute_fn,
+        gradient_fn=gradient_fn,
+        gradient_kwargs=gradient_kwargs,
+        _n=_n,
+    )
+
+
+def _execute(
+    params,
+    tapes=None,
+    device=None,
+    execute_fn=None,
+    gradient_fn=None,
+    gradient_kwargs=None,
+    _n=1,
+):  # pylint: disable=dangerous-default-value,unused-argument
+
+    # Only have scalar outputs
+    total_size = len(tapes)
+    total_params = np.sum([len(p) for p in params])
+
+    @jax.custom_vjp
+    def wrapped_exec(params):
+        def wrapper(p):
+            """Compute the forward pass."""
+            new_tapes = []
+
+            for t, a in zip(tapes, p):
+                new_tapes.append(t.copy(copy_operations=True))
+                new_tapes[-1].set_parameters(a)
+
+            with qml.tape.Unwrap(*new_tapes):
+                res, _ = execute_fn(new_tapes, **gradient_kwargs)
+
+            return res
+
+        shapes = [jax.ShapeDtypeStruct((1,), dtype) for _ in range(total_size)]
+        res = host_callback.call(wrapper, params, result_shape=shapes)
+        return res
+
+    def wrapped_exec_fwd(params):
+        return wrapped_exec(params), params
+
+    def wrapped_exec_bwd(params, g):
+
+        if isinstance(gradient_fn, qml.gradients.gradient_transform):
+
+            def non_diff_wrapper(args):
+                """Compute the VJP in a non-differentiable manner."""
+                new_tapes = []
+                p = args[:-1]
+                dy = args[-1]
+
+                for t, a in zip(tapes, p):
+                    new_tapes.append(t.copy(copy_operations=True))
+                    new_tapes[-1].set_parameters(a)
+                    new_tapes[-1].trainable_params = t.trainable_params
+
+                vjp_tapes, processing_fn = qml.gradients.batch_vjp(
+                    new_tapes,
+                    dy,
+                    gradient_fn,
+                    reduction="append",
+                    gradient_kwargs=gradient_kwargs,
+                )
+
+                partial_res = execute_fn(vjp_tapes)[0]
+                res = processing_fn(partial_res)
+                return np.concatenate(res)
+
+            args = tuple(params) + (g,)
+            vjps = host_callback.call(
+                non_diff_wrapper,
+                args,
+                result_shape=jax.ShapeDtypeStruct((total_params,), dtype),
+            )
+
+            param_idx = 0
+            res = []
+
+            # Group the vjps based on the parameters of the tapes
+            for p in params:
+                param_vjp = vjps[param_idx : param_idx + len(p)]
+                res.append(param_vjp)
+                param_idx += len(p)
+
+            # Unwrap partial results into ndim=0 arrays to allow
+            # differentiability with JAX
+            # E.g.,
+            # [DeviceArray([-0.9553365], dtype=float32), DeviceArray([0., 0.],
+            # dtype=float32)]
+            # is mapped to
+            # [[DeviceArray(-0.9553365, dtype=float32)], [DeviceArray(0.,
+            # dtype=float32), DeviceArray(0., dtype=float32)]].
+            need_unwrapping = any(r.ndim != 0 for r in res)
+            if need_unwrapping:
+                unwrapped_res = []
+                for r in res:
+                    if r.ndim != 0:
+                        r = [jnp.array(p) for p in r]
+                    unwrapped_res.append(r)
+
+                res = unwrapped_res
+
+            return (tuple(res),)
+
+        # Gradient function is a device method.
+        with qml.tape.Unwrap(*tapes):
+            jacs = gradient_fn(tapes, **gradient_kwargs)
+
+        vjps = [qml.gradients.compute_vjp(d, jac) for d, jac in zip(g, jacs)]
+        res = [[jnp.array(p) for p in v] for v in vjps]
+        return (tuple(res),)
+
+    wrapped_exec.defvjp(wrapped_exec_fwd, wrapped_exec_bwd)
+    return wrapped_exec(params)
+
+
+# The execute function in forward mode
+def _execute_with_fwd(
+    params,
+    tapes=None,
+    device=None,
+    execute_fn=None,
+    gradient_kwargs=None,
+    _n=1,
+):  # pylint: disable=dangerous-default-value,unused-argument
+
+    # Only have scalar outputs
+    total_size = len(tapes)
+
+    @jax.custom_vjp
+    def wrapped_exec(params):
+        def wrapper(p):
+            """Compute the forward pass by returning the jacobian too."""
+            new_tapes = []
+
+            for t, a in zip(tapes, p):
+                new_tapes.append(t.copy(copy_operations=True))
+                new_tapes[-1].set_parameters(a)
+
+            res, jacs = execute_fn(new_tapes, **gradient_kwargs)
+
+            # On the forward execution return the jacobian too
+            return res, jacs
+
+        fwd_shapes = [jax.ShapeDtypeStruct((1,), dtype) for _ in range(total_size)]
+        jacobian_shape = [jax.ShapeDtypeStruct((1, len(p)), dtype) for p in params]
+        res, jacs = host_callback.call(
+            wrapper,
+            params,
+            result_shape=tuple([fwd_shapes, jacobian_shape]),
+        )
+        return res, jacs
+
+    def wrapped_exec_fwd(params):
+        res, jacs = wrapped_exec(params)
+        return res, tuple([jacs, params])
+
+    def wrapped_exec_bwd(params, g):
+
+        # Use the jacobian that was computed on the forward pass
+        jacs, params = params
+
+        # Adjust the structure of how the jacobian is returned to match the
+        # non-forward mode cases
+        # E.g.,
+        # [DeviceArray([[ 0.06695931,  0.01383095, -0.46500877]], dtype=float32)]
+        # is mapped to
+        # [[DeviceArray(0.06695931, dtype=float32), DeviceArray(0.01383095,
+        # dtype=float32), DeviceArray(-0.46500877, dtype=float32)]]
+        res_jacs = []
+        for j in jacs:
+            this_j = []
+            for i in range(j.shape[1]):
+                this_j.append(j[0, i])
+            res_jacs.append(this_j)
+        return tuple([tuple(res_jacs)])
+
+    wrapped_exec.defvjp(wrapped_exec_fwd, wrapped_exec_bwd)
+    res = wrapped_exec(params)
+
+    tracing = any(isinstance(r, jax.interpreters.ad.JVPTracer) for r in res)
+
+    # When there are no tracers (not differentiating), we have the result of
+    # the forward pass and the jacobian, but only need the result of the
+    # forward pass
+    if len(res) == 2 and not tracing:
+        res = res[0]
+
+    return res
diff --git a/pennylane/tape/tape.py b/pennylane/tape/tape.py
index 66391c7f2..20bcf94d7 100644
--- a/pennylane/tape/tape.py
+++ b/pennylane/tape/tape.py
@@ -210,6 +210,7 @@ def expand_tape(tape, depth=1, stop_at=None, expand_measurements=False):
 
     # Update circuit info
     new_tape._update_circuit_info()
+    new_tape._output_dim = tape.output_dim
     return new_tape
 
 
diff --git a/pennylane/transforms/batch_transform.py b/pennylane/transforms/batch_transform.py
index 9c33f3490..731ab5404 100644
--- a/pennylane/transforms/batch_transform.py
+++ b/pennylane/transforms/batch_transform.py
@@ -242,10 +242,7 @@ class batch_transform:
                 gradient_fn = qml.gradients.finite_diff
 
             res = qml.execute(
-                tapes,
-                device=qnode.device,
-                gradient_fn=gradient_fn,
-                interface=interface,
+                tapes, device=qnode.device, gradient_fn=gradient_fn, interface=interface, max_diff=2
             )
 
             return processing_fn(res)


Test patch:
diff --git a/tests/beta/test_beta_qnode.py b/tests/beta/test_beta_qnode.py
index 1791b2a8d..7eb583a3f 100644
--- a/tests/beta/test_beta_qnode.py
+++ b/tests/beta/test_beta_qnode.py
@@ -36,7 +36,7 @@ class TestValidation:
         test_interface = "something"
         expected_error = (
             fr"Unknown interface {test_interface}\. Interface must be "
-            r"one of \[None, 'autograd', 'numpy', 'torch', 'pytorch', 'tf', 'tensorflow'\]\."
+            r"one of \[None, 'autograd', 'numpy', 'jax', 'JAX', 'torch', 'pytorch', 'tf', 'tensorflow'\]\."
         )
 
         with pytest.raises(qml.QuantumFunctionError, match=expected_error):
@@ -55,7 +55,7 @@ class TestValidation:
 
         expected_error = (
             fr"Unknown interface {test_interface}\. Interface must be "
-            r"one of \[None, 'autograd', 'numpy', 'torch', 'pytorch', 'tf', 'tensorflow'\]\."
+            r"one of \[None, 'autograd', 'numpy', 'jax', 'JAX', 'torch', 'pytorch', 'tf', 'tensorflow'\]\."
         )
 
         with pytest.raises(qml.QuantumFunctionError, match=expected_error):
diff --git a/tests/interfaces/test_batch_autograd.py b/tests/interfaces/test_batch_autograd.py
index aca6ec911..9521d9cc2 100644
--- a/tests/interfaces/test_batch_autograd.py
+++ b/tests/interfaces/test_batch_autograd.py
@@ -271,7 +271,7 @@ class TestCaching:
                 qml.CNOT(wires=[0, 1])
                 qml.var(qml.PauliZ(0) @ qml.PauliX(1))
 
-            return execute([tape], dev, gradient_fn=param_shift, cache=cache)[0]
+            return execute([tape], dev, gradient_fn=param_shift, cache=cache, max_diff=2)[0]
 
         # No caching: number of executions is not ideal
         hess1 = qml.jacobian(qml.grad(cost))(params, cache=False)
@@ -763,7 +763,7 @@ class TestHigherOrderDerivatives:
                 qml.CNOT(wires=[0, 1])
                 qml.probs(wires=1)
 
-            result = execute([tape1, tape2], dev, gradient_fn=param_shift)
+            result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2)
             return result[0] + result[1][0, 0]
 
         res = cost_fn(params)
diff --git a/tests/interfaces/test_batch_jax.py b/tests/interfaces/test_batch_jax.py
new file mode 100644
index 000000000..a5e50e24a
--- /dev/null
+++ b/tests/interfaces/test_batch_jax.py
@@ -0,0 +1,612 @@
+# 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 jax interface"""
+import functools
+
+import pytest
+
+jax = pytest.importorskip("jax")
+jnp = pytest.importorskip("jax.numpy")
+import numpy as np
+
+import pennylane as qml
+from pennylane.gradients import param_shift
+from pennylane.interfaces.batch import execute
+
+
+class TestJaxExecuteUnitTests:
+    """Unit tests for jax execution"""
+
+    def test_jacobian_options(self, mocker, tol):
+        """Test setting jacobian options"""
+        spy = mocker.spy(qml.gradients, "param_shift")
+
+        a = jnp.array([0.1, 0.2])
+
+        dev = qml.device("default.qubit", wires=1)
+
+        def cost(a, device):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute(
+                [tape],
+                device,
+                gradient_fn=param_shift,
+                gradient_kwargs={"shift": np.pi / 4},
+                interface="jax",
+            )[0][0]
+
+        res = jax.grad(cost)(a, device=dev)
+
+        for args in spy.call_args_list:
+            assert args[1]["shift"] == np.pi / 4
+
+    @pytest.mark.skip(reason="just because")
+    def test_unknown_gradient_fn_error(self):
+        """Test that an error is raised if an unknown gradient function
+        is passed"""
+        a = jnp.array([0.1, 0.2])
+
+        dev = qml.device("default.qubit", wires=1)
+
+        def cost(a, device):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute([tape], device, gradient_fn=lambda x: x, interface="jax",)[
+                0
+            ][0]
+
+        with pytest.raises(ValueError, match="Unknown gradient function"):
+            res = jax.grad(cost)(a, device=dev)
+
+    def test_incorrect_mode(self):
+        """Test that an error is raised if an gradient transform
+        is used with mode=forward"""
+        a = jnp.array([0.1, 0.2])
+
+        dev = qml.device("default.qubit", wires=1)
+
+        def cost(a, device):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute(
+                [tape],
+                device,
+                gradient_fn=param_shift,
+                mode="forward",
+                interface="jax",
+            )[0]
+
+        with pytest.raises(
+            ValueError, match="Gradient transforms cannot be used with mode='forward'"
+        ):
+            res = jax.grad(cost)(a, device=dev)
+
+    def test_unknown_interface(self):
+        """Test that an error is raised if the interface is unknown"""
+        a = jnp.array([0.1, 0.2])
+
+        dev = qml.device("default.qubit", wires=1)
+
+        def cost(a, device):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute([tape], device, gradient_fn=param_shift, interface="None")[0]
+
+        with pytest.raises(ValueError, match="Unknown interface"):
+            cost(a, device=dev)
+
+    def test_forward_mode(self, mocker):
+        """Test that forward mode uses the `device.execute_and_gradients` pathway"""
+        dev = qml.device("default.qubit", wires=1)
+        spy = mocker.spy(dev, "execute_and_gradients")
+
+        def cost(a):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute(
+                [tape],
+                dev,
+                gradient_fn="device",
+                interface="jax",
+                gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True},
+            )[0]
+
+        a = jnp.array([0.1, 0.2])
+        cost(a)
+
+        # adjoint method only performs a single device execution, but gets both result and gradient
+        assert dev.num_executions == 1
+        spy.assert_called()
+
+    def test_backward_mode(self, mocker):
+        """Test that backward mode uses the `device.batch_execute` and `device.gradients` pathway"""
+        dev = qml.device("default.qubit", wires=1)
+        spy_execute = mocker.spy(qml.devices.DefaultQubit, "batch_execute")
+        spy_gradients = mocker.spy(qml.devices.DefaultQubit, "gradients")
+
+        def cost(a):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute(
+                [tape],
+                dev,
+                gradient_fn="device",
+                mode="backward",
+                gradient_kwargs={"method": "adjoint_jacobian"},
+                interface="jax",
+            )[0][0]
+
+        a = jnp.array([0.1, 0.2])
+        cost(a)
+
+        assert dev.num_executions == 1
+        spy_execute.assert_called()
+        spy_gradients.assert_not_called()
+
+        jax.grad(cost)(a)
+        spy_gradients.assert_called()
+
+    def test_max_diff_error(self):
+        """Test that an error is being raised if max_diff > 1 for the JAX
+        interface."""
+        a = jnp.array([0.1, 0.2])
+
+        dev = qml.device("default.qubit", wires=1)
+
+        with pytest.raises(
+            ValueError, match="The JAX interface only supports first order derivatives."
+        ):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            execute(
+                [tape],
+                dev,
+                gradient_fn=param_shift,
+                gradient_kwargs={"shift": np.pi / 4},
+                interface="jax",
+                max_diff=2,
+            )
+
+
+class TestCaching:
+    """Test for caching behaviour"""
+
+    def test_cache_maxsize(self, mocker):
+        """Test the cachesize property of the cache"""
+        dev = qml.device("default.qubit", wires=1)
+        spy = mocker.spy(qml.interfaces.batch, "cache_execute")
+
+        def cost(a, cachesize):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute(
+                [tape], dev, gradient_fn=param_shift, cachesize=cachesize, interface="jax"
+            )[0][0]
+
+        params = jnp.array([0.1, 0.2])
+        jax.grad(cost)(params, cachesize=2)
+        cache = spy.call_args[0][1]
+
+        assert cache.maxsize == 2
+        assert cache.currsize == 2
+        assert len(cache) == 2
+
+    def test_custom_cache(self, mocker):
+        """Test the use of a custom cache object"""
+        dev = qml.device("default.qubit", wires=1)
+        spy = mocker.spy(qml.interfaces.batch, "cache_execute")
+
+        def cost(a, cache):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute([tape], dev, gradient_fn=param_shift, cache=cache, interface="jax")[0][0]
+
+        custom_cache = {}
+        params = jnp.array([0.1, 0.2])
+        jax.grad(cost)(params, cache=custom_cache)
+
+        cache = spy.call_args[0][1]
+        assert cache is custom_cache
+
+    def test_custom_cache_multiple(self, mocker):
+        """Test the use of a custom cache object with multiple tapes"""
+        dev = qml.device("default.qubit", wires=1)
+        spy = mocker.spy(qml.interfaces.batch, "cache_execute")
+
+        a = jnp.array(0.1)
+        b = jnp.array(0.2)
+
+        def cost(a, b, cache):
+            with qml.tape.JacobianTape() as tape1:
+                qml.RY(a, wires=0)
+                qml.RX(b, wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            with qml.tape.JacobianTape() as tape2:
+                qml.RY(a, wires=0)
+                qml.RX(b, wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            res = execute(
+                [tape1, tape2], dev, gradient_fn=param_shift, cache=cache, interface="jax"
+            )
+            return res[0][0]
+
+        custom_cache = {}
+        jax.grad(cost)(a, b, cache=custom_cache)
+
+        cache = spy.call_args[0][1]
+        assert cache is custom_cache
+
+    def test_caching_param_shift(self, tol):
+        """Test that, when using parameter-shift transform,
+        caching produces the optimum number of evaluations."""
+        dev = qml.device("default.qubit", wires=1)
+
+        def cost(a, cache):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute([tape], dev, gradient_fn=param_shift, cache=cache, interface="jax")[0][0]
+
+        # Without caching, 5 evaluations are required to compute
+        # the Jacobian: 1 (forward pass) + 2 (backward pass) * (2 shifts * 2 params)
+        params = jnp.array([0.1, 0.2])
+        jax.grad(cost)(params, cache=None)
+        assert dev.num_executions == 5
+
+        # With caching, 5 evaluations are required to compute
+        # the Jacobian: 1 (forward pass) + (2 shifts * 2 params)
+        dev._num_executions = 0
+        jac_fn = jax.grad(cost)
+        grad1 = jac_fn(params, cache=True)
+        assert dev.num_executions == 5
+
+        # Check that calling the cost function again
+        # continues to evaluate the device (that is, the cache
+        # is emptied between calls)
+        grad2 = jac_fn(params, cache=True)
+        assert dev.num_executions == 10
+        assert np.allclose(grad1, grad2, atol=tol, rtol=0)
+
+        # Check that calling the cost function again
+        # with different parameters produces a different Jacobian
+        grad2 = jac_fn(2 * params, cache=True)
+        assert dev.num_executions == 15
+        assert not np.allclose(grad1, grad2, atol=tol, rtol=0)
+
+    def test_caching_adjoint_backward(self):
+        """Test that caching produces the optimum number of adjoint evaluations
+        when mode=backward"""
+        dev = qml.device("default.qubit", wires=2)
+        params = jnp.array([0.1, 0.2, 0.3])
+
+        def cost(a, cache):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.RY(a[2], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute(
+                [tape],
+                dev,
+                gradient_fn="device",
+                cache=cache,
+                mode="backward",
+                interface="jax",
+                gradient_kwargs={"method": "adjoint_jacobian"},
+            )[0][0]
+
+        # Without caching, 2 evaluations are required.
+        # 1 for the forward pass, and one per output dimension
+        # on the backward pass.
+        jax.grad(cost)(params, cache=None)
+        assert dev.num_executions == 2
+
+        # With caching, also 2 evaluations are required. One
+        # for the forward pass, and one for the backward pass.
+        dev._num_executions = 0
+        jac_fn = jax.grad(cost)
+        grad1 = jac_fn(params, cache=True)
+        assert dev.num_executions == 2
+
+
+execute_kwargs = [
+    {"gradient_fn": param_shift},
+    {
+        "gradient_fn": "device",
+        "mode": "forward",
+        "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True},
+    },
+    {
+        "gradient_fn": "device",
+        "mode": "backward",
+        "gradient_kwargs": {"method": "adjoint_jacobian"},
+    },
+]
+
+
+@pytest.mark.parametrize("execute_kwargs", execute_kwargs)
+class TestJaxExecuteIntegration:
+    """Test the jax interface execute function
+    integrates well for both forward and backward execution"""
+
+    def test_execution(self, execute_kwargs):
+        """Test execution"""
+        dev = qml.device("default.qubit", wires=1)
+
+        def cost(a, b):
+            with qml.tape.JacobianTape() as tape1:
+                qml.RY(a, wires=0)
+                qml.RX(b, wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            with qml.tape.JacobianTape() as tape2:
+                qml.RY(a, wires=0)
+                qml.RX(b, wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute([tape1, tape2], dev, interface="jax", **execute_kwargs)
+
+        a = jnp.array(0.1)
+        b = jnp.array(0.2)
+        res = cost(a, b)
+
+        assert len(res) == 2
+        assert res[0].shape == (1,)
+        assert res[1].shape == (1,)
+
+    def test_scalar_jacobian(self, execute_kwargs, tol):
+        """Test scalar jacobian calculation"""
+        a = jnp.array(0.1)
+        dev = qml.device("default.qubit", wires=2)
+
+        def cost(a):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a, wires=0)
+                qml.expval(qml.PauliZ(0))
+            return execute([tape], dev, interface="jax", **execute_kwargs)[0][0]
+
+        res = jax.grad(cost)(a)
+        assert res.shape == ()
+
+        # compare to standard tape jacobian
+        with qml.tape.JacobianTape() as tape:
+            qml.RY(a, wires=0)
+            qml.expval(qml.PauliZ(0))
+
+        tape.trainable_params = {0}
+        tapes, fn = param_shift(tape)
+        expected = fn(dev.batch_execute(tapes))
+
+        assert expected.shape == (1, 1)
+        assert np.allclose(res, np.squeeze(expected), atol=tol, rtol=0)
+
+    def test_reusing_quantum_tape(self, execute_kwargs, tol):
+        """Test re-using a quantum tape by passing new parameters"""
+        a = jnp.array(0.1)
+        b = jnp.array(0.2)
+
+        dev = qml.device("default.qubit", wires=2)
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RY(a, wires=0)
+            qml.RX(b, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        assert tape.trainable_params == {0, 1}
+
+        def cost(a, b):
+
+            # An explicit call to _update() is required here to update the
+            # trainable parameters in between tape executions.
+            # This is different from how the autograd interface works.
+            # Unless the update is issued, the validation check related to the
+            # number of provided parameters fails in the tape: (len(params) !=
+            # required_length) and the tape produces incorrect results.
+            tape._update()
+            tape.set_parameters([a, b])
+            return execute([tape], dev, interface="jax", **execute_kwargs)[0][0]
+
+        jac_fn = jax.grad(cost)
+        jac = jac_fn(a, b)
+
+        a = jnp.array(0.54)
+        b = jnp.array(0.8)
+
+        # check that the cost function continues to depend on the
+        # values of the parameters for subsequent calls
+        res2 = cost(2 * a, b)
+        expected = [np.cos(2 * a)]
+        assert np.allclose(res2, expected, atol=tol, rtol=0)
+
+        jac_fn = jax.grad(lambda a, b: cost(2 * a, b))
+        jac = jac_fn(a, b)
+        expected = -2 * np.sin(2 * a)
+        assert np.allclose(jac, expected, atol=tol, rtol=0)
+
+    def test_classical_processing_single_tape(self, execute_kwargs, tol):
+        """Test classical processing within the quantum tape for a single tape"""
+        a = jnp.array(0.1)
+        b = jnp.array(0.2)
+        c = jnp.array(0.3)
+
+        def cost(a, b, c, device):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a * c, wires=0)
+                qml.RZ(b, wires=0)
+                qml.RX(c + c ** 2 + jnp.sin(a), wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return execute([tape], device, interface="jax", **execute_kwargs)[0][0]
+
+        dev = qml.device("default.qubit", wires=2)
+        res = jax.grad(cost, argnums=(0, 1, 2))(a, b, c, device=dev)
+        assert len(res) == 3
+
+    def test_classical_processing_multiple_tapes(self, execute_kwargs, tol):
+        """Test classical processing within the quantum tape for multiple
+        tapes"""
+        dev = qml.device("default.qubit", wires=2)
+        params = jax.numpy.array([0.3, 0.2])
+
+        def cost_fn(x):
+            with qml.tape.JacobianTape() as tape1:
+                qml.Hadamard(0)
+                qml.RY(x[0], wires=[0])
+                qml.CNOT(wires=[0, 1])
+                qml.expval(qml.PauliZ(0))
+
+            with qml.tape.JacobianTape() as tape2:
+                qml.Hadamard(0)
+                qml.CRX(2 * x[0] * x[1], wires=[0, 1])
+                qml.RX(2 * x[1], wires=[1])
+                qml.expval(qml.PauliZ(0))
+
+            result = execute(tapes=[tape1, tape2], device=dev, interface="jax", **execute_kwargs)
+            return (result[0] + result[1] - 7 * result[1])[0]
+
+        res = jax.grad(cost_fn)(params)
+        assert res.shape == (2,)
+
+    def test_multiple_tapes_output(self, execute_kwargs, tol):
+        """Test the output types for the execution of multiple quantum tapes"""
+        dev = qml.device("default.qubit", wires=2)
+        params = jax.numpy.array([0.3, 0.2])
+
+        def cost_fn(x):
+            with qml.tape.JacobianTape() as tape1:
+                qml.Hadamard(0)
+                qml.RY(x[0], wires=[0])
+                qml.CNOT(wires=[0, 1])
+                qml.expval(qml.PauliZ(0))
+
+            with qml.tape.JacobianTape() as tape2:
+                qml.Hadamard(0)
+                qml.CRX(2 * x[0] * x[1], wires=[0, 1])
+                qml.RX(2 * x[1], wires=[1])
+                qml.expval(qml.PauliZ(0))
+
+            return execute(tapes=[tape1, tape2], device=dev, interface="jax", **execute_kwargs)
+
+        res = cost_fn(params)
+        assert isinstance(res, list)
+        assert all(isinstance(r, jnp.ndarray) for r in res)
+        assert all(r.shape == (1,) for r in res)
+
+    @pytest.mark.xfail
+    def test_matrix_parameter(self, execute_kwargs, tol):
+        """Test that the jax interface works correctly
+        with a matrix parameter"""
+        a = jnp.array(0.1)
+        U = qml.RY(a, wires=0).matrix
+
+        def cost(U, device):
+            with qml.tape.JacobianTape() as tape:
+                qml.PauliX(0)
+                qml.QubitUnitary(U, wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            tape.trainable_params = {0}
+            return execute([tape], device, interface="jax", **execute_kwargs)[0][0]
+
+        dev = qml.device("default.qubit", wires=2)
+        res = cost(U, device=dev)
+        assert np.allclose(res, -np.cos(a), atol=tol, rtol=0)
+
+        jac_fn = jax.grad(cost, argnums=(0))
+        res = jac_fn(U, device=dev)
+        assert np.allclose(res, np.sin(a), atol=tol, rtol=0)
+
+    def test_differentiable_expand(self, execute_kwargs, tol):
+        """Test that operation and nested tapes expansion
+        is differentiable"""
+
+        class U3(qml.U3):
+            def expand(self):
+                tape = qml.tape.JacobianTape()
+                theta, phi, lam = self.data
+                wires = self.wires
+                tape._ops += [
+                    qml.Rot(lam, theta, -lam, wires=wires),
+                    qml.PhaseShift(phi + lam, wires=wires),
+                ]
+                return tape
+
+        def cost_fn(a, p, device):
+            tape = qml.tape.JacobianTape()
+
+            with tape:
+                qml.RX(a, wires=0)
+                U3(*p, wires=0)
+                qml.expval(qml.PauliX(0))
+
+            tape = tape.expand(stop_at=lambda obj: device.supports_operation(obj.name))
+            return execute([tape], device, interface="jax", **execute_kwargs)[0][0]
+
+        a = jnp.array(0.1)
+        p = jnp.array([0.1, 0.2, 0.3])
+
+        dev = qml.device("default.qubit", wires=1)
+        res = cost_fn(a, p, device=dev)
+        expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * (
+            np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2])
+        )
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        jac_fn = jax.grad(cost_fn, argnums=(1))
+        res = jac_fn(a, p, device=dev)
+        expected = jnp.array(
+            [
+                np.cos(p[1]) * (np.cos(a) * np.cos(p[0]) - np.sin(a) * np.sin(p[0]) * np.sin(p[2])),
+                np.cos(p[1]) * np.cos(p[2]) * np.sin(a)
+                - np.sin(p[1])
+                * (np.cos(a) * np.sin(p[0]) + np.cos(p[0]) * np.sin(a) * np.sin(p[2])),
+                np.sin(a)
+                * (np.cos(p[0]) * np.cos(p[1]) * np.cos(p[2]) - np.sin(p[1]) * np.sin(p[2])),
+            ]
+        )
+        assert np.allclose(res, expected, atol=tol, rtol=0)
diff --git a/tests/interfaces/test_batch_tensorflow.py b/tests/interfaces/test_batch_tensorflow.py
index 948d0a858..9b6fa692e 100644
--- a/tests/interfaces/test_batch_tensorflow.py
+++ b/tests/interfaces/test_batch_tensorflow.py
@@ -231,7 +231,9 @@ class TestCaching:
                 qml.CNOT(wires=[0, 1])
                 qml.var(qml.PauliZ(0) @ qml.PauliX(1))
 
-            return execute([tape], dev, gradient_fn=param_shift, cache=cache, interface="tf")[0]
+            return execute(
+                [tape], dev, gradient_fn=param_shift, cache=cache, interface="tf", max_diff=2
+            )[0]
 
         # No caching: number of executions is not ideal
         with tf.GradientTape() as t2:
@@ -355,7 +357,7 @@ class TestTensorFlowExecuteIntegration:
                 qml.CNOT(wires=[0, 1])
                 qml.expval(qml.PauliZ(0))
                 qml.expval(qml.PauliY(1))
-            res = execute([tape], dev, **execute_kwargs)[0]
+            res = execute([tape], dev, max_diff=2, **execute_kwargs)[0]
 
         expected = [np.cos(a), -np.cos(a) * np.sin(b)]
         assert np.allclose(res, expected, atol=tol, rtol=0)
@@ -728,7 +730,9 @@ class TestHigherOrderDerivatives:
                     qml.CNOT(wires=[0, 1])
                     qml.probs(wires=1)
 
-                result = execute([tape1, tape2], dev, gradient_fn=param_shift, interface="tf")
+                result = execute(
+                    [tape1, tape2], dev, gradient_fn=param_shift, interface="tf", max_diff=2
+                )
                 res = result[0][0] + result[1][0, 0]
 
             expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y))
@@ -761,7 +765,7 @@ class TestHigherOrderDerivatives:
                     qml.RX(params[1], wires=0)
                     qml.probs(wires=0)
 
-                res = execute([tape], dev, gradient_fn=param_shift, interface="tf")
+                res = execute([tape], dev, gradient_fn=param_shift, interface="tf", max_diff=2)
                 res = tf.stack(res)
 
             g = t1.jacobian(res, params, experimental_use_pfor=False)
diff --git a/tests/interfaces/test_batch_torch.py b/tests/interfaces/test_batch_torch.py
index 2c9db85c4..99f0b7d77 100644
--- a/tests/interfaces/test_batch_torch.py
+++ b/tests/interfaces/test_batch_torch.py
@@ -247,7 +247,9 @@ class TestCaching:
                 qml.CNOT(wires=[0, 1])
                 qml.var(qml.PauliZ(0) @ qml.PauliX(1))
 
-            return execute([tape], dev, gradient_fn=param_shift, cache=cache, interface="torch")[0]
+            return execute(
+                [tape], dev, gradient_fn=param_shift, cache=cache, interface="torch", max_diff=2
+            )[0]
 
         # No caching: number of executions is not ideal
         hess1 = torch.autograd.functional.hessian(lambda x: cost(x, cache=None), params)
@@ -878,7 +880,9 @@ class TestHigherOrderDerivatives:
                 qml.CNOT(wires=[0, 1])
                 qml.probs(wires=1)
 
-            result = execute([tape1, tape2], dev, gradient_fn=param_shift, interface="torch")
+            result = execute(
+                [tape1, tape2], dev, gradient_fn=param_shift, interface="torch", max_diff=2
+            )
             return result[0] + result[1][0, 0]
 
         res = cost_fn(params)
@@ -911,7 +915,9 @@ class TestHigherOrderDerivatives:
                 qml.RX(x[1], wires=0)
                 qml.probs(wires=0)
 
-            return torch.stack(execute([tape], dev, gradient_fn=param_shift, interface="torch"))
+            return torch.stack(
+                execute([tape], dev, gradient_fn=param_shift, interface="torch", max_diff=2)
+            )
 
         x = torch.tensor([1.0, 2.0], requires_grad=True, device=torch_device)
         res = circuit(x)