Repository: PennyLaneAI/pennylane
Pull request #1494 (2021-08-04 14:50:15)

Adds a differentiable VJP transform
**Context:** Currently, PennyLane provides a `tape.jacobian()` method, which returns the Jacobian of a quantum circuit. It is then up to the interfaces to take this result, cast it to the correct tensor dtype, and construct the VJP. At the moment, this process is not ideal:

- The Jacobian is always computed prior to the vjp. This sometimes results in redundant quantum evals, especially if `dy = 0`.
- The VJP logic is duplicated across each interface, increasing maintenance burden, and leading to bugs not always being fixed in all places.
- There are no unit-tests for the VJP logic; instead, it is only tested via larger integration tests.
- Finally, the VJP logic is not differentiable, blocking higher order derivatives from being computed.

**Description of the Change:**

- Adds `qml.gradients.vjp` for computing the vector-Jacobian product of a tape
- Adds `qml.gradients.batch_vjp` for computing the vector-Jacobian products of multiple tapes

**Benefits:**

- Users can pass in arbitrary gradient transforms to use in the computation
- Written to be autodifferentiable
- VJP unit tests are now present

**Possible Drawbacks:** None. In a future PR, the vector-Hessian product should also be added:

```python
def vhp(tape, dy, ddy):
    if dy.size > 1:
        vhp = dy @ ddy @ hessian @ dy.T / np.linalg.norm(dy) ** 2
    else:
        vhp = ddy @ hessian
```

**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.
josh146: [ch8004]
josh146: Thanks @glassnotes, have incorporated your suggestions!

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 15500bc8a..0c5204cea 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -2,6 +2,14 @@
 
 <h3>New features since last release</h3>
 
+* Vector-Jacobian product transforms have been added to the `qml.gradients` package.
+  [(#1494)](https://github.com/PennyLaneAI/pennylane/pull/1494)
+
+  The new transforms include:
+
+  - `qml.gradients.vjp`
+  - `qml.gradients.batch_vjp`
+  
 * The Hamiltonian can now store grouping information, which can be accessed by a device to 
   speed up computations of the expectation value of a Hamiltonian. 
   [(#1515)](https://github.com/PennyLaneAI/pennylane/pull/1515)
@@ -93,7 +101,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Maria Schuld.
+Josh Izaac, Maria Schuld.
 
 # Release 0.17.0 (current release)
 
@@ -265,7 +273,7 @@ Maria Schuld.
   For example,
 
   ```pycon
-  >>> with qml.tape.QuantumTape() as tape:
+  >>> with qml.tape.JacobianTape() as tape:
   ...     qml.RX(params[0], wires=0)
   ...     qml.RY(params[1], wires=0)
   ...     qml.RX(params[2], wires=0)
diff --git a/pennylane/gradients/__init__.py b/pennylane/gradients/__init__.py
index 2db236e9d..bd7162503 100644
--- a/pennylane/gradients/__init__.py
+++ b/pennylane/gradients/__init__.py
@@ -21,3 +21,4 @@ from . import parameter_shift_cv
 from .finite_difference import finite_diff, finite_diff_coeffs, generate_shifted_tapes
 from .parameter_shift import param_shift
 from .parameter_shift_cv import param_shift_cv
+from .vjp import compute_vjp, batch_vjp, vjp
diff --git a/pennylane/gradients/vjp.py b/pennylane/gradients/vjp.py
new file mode 100644
index 000000000..7e68eb3bd
--- /dev/null
+++ b/pennylane/gradients/vjp.py
@@ -0,0 +1,314 @@
+# 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 computing the vector-Jacobian product
+of tapes.
+"""
+import numpy as np
+
+from pennylane import math
+
+
+def compute_vjp(dy, jac):
+    """Convenience function to compute the vector-Jacobian product for a given
+    vector of gradient outputs and a Jacobian.
+
+    Args:
+        dy (tensor_like): vector of gradient outputs
+        jac (tensor_like): Jacobian matrix. For an n-dimensional ``dy``
+            vector, the first n-dimensions of ``jac`` should match
+            the shape of ``dy``.
+
+    Returns:
+        tensor_like: the vector-Jacobian product
+    """
+    if jac is None:
+        return None
+
+    dy_row = math.reshape(dy, [-1])
+    jac = math.reshape(jac, [dy_row.shape[0], -1])
+
+    if math.allclose(dy, 0):
+        # If the dy vector is zero, then the
+        # corresponding element of the VJP will be zero.
+        num_params = jac.shape[1]
+        return math.convert_like(np.zeros([num_params]), dy)
+
+    return math.tensordot(jac, dy_row, [[0], [0]])
+
+
+def vjp(tape, dy, gradient_fn):
+    r"""Generate the gradient tapes and processing function required to compute
+    the vector-Jacobian product of a tape.
+
+    Consider a function :math:`\mathbf{f}(\mathbf{x})`. The Jacobian is given by
+
+    .. math::
+
+        \mathbf{J}_{\mathbf{f}}(\mathbf{x}) = \begin{pmatrix}
+            \frac{\partial f_1}{\partial x_1} &\cdots &\frac{\partial f_1}{\partial x_n}\\
+            \vdots &\ddots &\vdots\\
+            \frac{\partial f_m}{\partial x_1} &\cdots &\frac{\partial f_m}{\partial x_n}\\
+        \end{pmatrix}.
+
+    During backpropagation, the chain rule is applied. For example, consider the
+    cost function :math:`h = y\circ f: \mathbb{R}^n \rightarrow \mathbb{R}`,
+    where :math:`y: \mathbb{R}^m \rightarrow \mathbb{R}`.
+    The gradient is:
+
+    .. math::
+
+        \nabla h(\mathbf{x}) = \frac{\partial y}{\partial \mathbf{f}} \frac{\partial \mathbf{f}}{\partial \mathbf{x}}
+        = \frac{\partial y}{\partial \mathbf{f}} \mathbf{J}_{\mathbf{f}}(\mathbf{x}).
+
+    Denote :math:`d\mathbf{y} = \frac{\partial y}{\partial \mathbf{f}}`; we can write this in the form
+    of a matrix multiplication:
+
+    .. math:: \left[\nabla h(\mathbf{x})\right]_{j} = \sum_{i=0}^m d\mathbf{y}_i ~ \mathbf{J}_{ij}.
+
+    Thus, we can see that the gradient of the cost function is given by the so-called
+    **vector-Jacobian product**; the product of the row-vector :math:`d\mathbf{y}`, representing
+    the gradient of subsequent components of the cost function, and :math:`\mathbf{J}`,
+    the Jacobian of the current node of interest.
+
+    Args:
+        tape (.QuantumTape): quantum tape to differentiate
+        dy (tensor_like): Gradient-output vector. Must have shape
+            matching the output shape of the corresponding tape.
+        gradient_fn (callable): the gradient transform to use to differentiate
+            the tape
+
+    Returns:
+        tensor_like or None: Vector-Jacobian product. Returns None if the tape
+        has no trainable parameters.
+
+    **Example**
+
+    Consider the following Torch-compatible quantum tape:
+
+    .. code-block:: python
+
+        import torch
+        from pennylane.interfaces.torch import TorchInterface
+
+        x = torch.tensor([[0.1, 0.2, 0.3],
+                          [0.4, 0.5, 0.6]], requires_grad=True, dtype=torch.float64)
+
+        with TorchInterface.apply(qml.tape.JacobianTape()) as tape:
+            qml.RX(x[0, 0], wires=0)
+            qml.RY(x[0, 1], wires=1)
+            qml.RZ(x[0, 2], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RX(x[1, 0], wires=1)
+            qml.RY(x[1, 1], wires=0)
+            qml.RZ(x[1, 2], wires=1)
+            qml.expval(qml.PauliZ(0))
+            qml.probs(wires=1)
+
+    We can use the ``vjp`` function to compute the vector-Jacobian product,
+    given a gradient-output vector ``dy``:
+
+    >>> dy = torch.tensor([1., 1., 1.], dtype=torch.float64)
+    >>> vjp_tapes, fn = qml.gradients.vjp(tape, dy, qml.gradients.param_shift)
+
+    Note that ``dy`` has shape ``(3,)``, matching the output dimension of the tape
+    (1 expectation and 2 probability values).
+
+    Executing the VJP tapes, and applying the processing function:
+
+    >>> dev = qml.device("default.qubit", wires=2)
+    >>> vjp = fn([t.execute(dev) for t in vjp_tapes])
+    >>> vjp
+    tensor([-0.6069, -0.0451,  0.0451, -0.0139, -0.2809,  0.2809],
+           dtype=torch.float64, grad_fn=<ViewBackward>)
+
+    The output VJP is also differentiable with respect to the tape parameters:
+
+    >>> cost = torch.sum(vjp)
+    >>> cost.backward()
+    >>> x.grad
+    tensor([[-1.1025e+00, -2.0554e-01, -1.4917e-01],
+            [-1.9429e-09, -9.1580e-01,  1.3878e-09]], dtype=torch.float64)
+    """
+    # t._par_info = {}
+    # t._update()
+    num_params = len(tape.trainable_params)
+
+    if num_params == 0:
+        # The tape has no trainable parameters; the VJP
+        # is simply none.
+        return [], lambda _: None
+
+    if math.allclose(dy, 0):
+        # If the dy vector is zero, then the
+        # corresponding element of the VJP will be zero,
+        # and we can avoid a quantum computation.
+        return [], lambda _: math.convert_like(np.zeros([num_params]), dy)
+
+    gradient_tapes, fn = gradient_fn(tape)
+
+    def processing_fn(results):
+        # postprocess results to compute the Jacobian
+        jac = fn(results)
+        return compute_vjp(dy, jac)
+
+    return gradient_tapes, processing_fn
+
+
+def batch_vjp(tapes, dys, gradient_fn, reduction="append"):
+    r"""Generate the gradient tapes and processing function required to compute
+    the vector-Jacobian products of a batch of tapes.
+
+    Consider a function :math:`\mathbf{f}(\mathbf{x})`. The Jacobian is given by
+
+    .. math::
+
+        \mathbf{J}_{\mathbf{f}}(\mathbf{x}) = \begin{pmatrix}
+            \frac{\partial f_1}{\partial x_1} &\cdots &\frac{\partial f_1}{\partial x_n}\\
+            \vdots &\ddots &\vdots\\
+            \frac{\partial f_m}{\partial x_1} &\cdots &\frac{\partial f_m}{\partial x_n}\\
+        \end{pmatrix}.
+
+    During backpropagation, the chain rule is applied. For example, consider the
+    cost function :math:`h = y\circ f: \mathbb{R}^n \rightarrow \mathbb{R}`,
+    where :math:`y: \mathbb{R}^m \rightarrow \mathbb{R}`.
+    The gradient is:
+
+    .. math::
+
+        \nabla h(\mathbf{x}) = \frac{\partial y}{\partial \mathbf{f}} \frac{\partial \mathbf{f}}{\partial \mathbf{x}}
+        = \frac{\partial y}{\partial \mathbf{f}} \mathbf{J}_{\mathbf{f}}(\mathbf{x}).
+
+    Denote :math:`d\mathbf{y} = \frac{\partial y}{\partial \mathbf{f}}`; we can write this in the form
+    of a matrix multiplication:
+
+    .. math:: \left[\nabla h(\mathbf{x})\right]_{j} = \sum_{i=0}^m d\mathbf{y}_i ~ \mathbf{J}_{ij}.
+
+    Thus, we can see that the gradient of the cost function is given by the so-called
+    **vector-Jacobian product**; the product of the row-vector :math:`d\mathbf{y}`, representing
+    the gradient of subsequent components of the cost function, and :math:`\mathbf{J}`,
+    the Jacobian of the current node of interest.
+
+    Args:
+        tapes (Sequence[.QuantumTape]): sequence of quantum tapes to differentiate
+        dys (Sequence[tensor_like]): Sequence of gradient-output vectors ``dy``. Must be the
+            same length as ``tapes``. Each ``dy`` tensor should have shape
+            matching the output shape of the corresponding tape.
+        gradient_fn (callable): the gradient transform to use to differentiate
+            the tapes
+        reduction (str): Determines how the vector-Jacobian products are returned.
+            If ``append``, then the output of the function will be of the form
+            ``List[tensor_like]``, with each element corresponding to the VJP of each
+            input tape. If ``extend``, then the output VJPs will be concatenated.
+
+    Returns:
+        List[tensor_like or None]: list of vector-Jacobian products. ``None`` elements corresponds
+        to tapes with no trainable parameters.
+
+    **Example**
+
+    Consider the following Torch-compatible quantum tapes:
+
+    .. code-block:: python
+
+        x = torch.tensor([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], requires_grad=True, dtype=torch.float64)
+
+        def ansatz(x):
+            qml.RX(x[0, 0], wires=0)
+            qml.RY(x[0, 1], wires=1)
+            qml.RZ(x[0, 2], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RX(x[1, 0], wires=1)
+            qml.RY(x[1, 1], wires=0)
+            qml.RZ(x[1, 2], wires=1)
+
+        with TorchInterface.apply(qml.tape.JacobianTape()) as tape1:
+            ansatz(x)
+            qml.expval(qml.PauliZ(0))
+            qml.probs(wires=1)
+
+        with TorchInterface.apply(qml.tape.JacobianTape()) as tape2:
+            ansatz(x)
+            qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        tapes = [tape1, tape2]
+
+    Both tapes share the same circuit ansatz, but have different measurement outputs.
+
+    We can use the ``batch_vjp`` function to compute the vector-Jacobian product,
+    given a list of gradient-output vectors ``dys`` per tape:
+
+    >>> dys = [torch.tensor([1., 1., 1.], dtype=torch.float64),
+    ...  torch.tensor([1.], dtype=torch.float64)]
+    >>> vjp_tapes, fn = qml.gradients.batch_vjp(tapes, dys, qml.gradients.param_shift)
+
+    Note that each ``dy`` has shape matching the output dimension of the tape
+    (``tape1`` has 1 expectation and 2 probability values --- 3 outputs --- and ``tape2``
+    has 1 expectation value).
+
+    Executing the VJP tapes, and applying the processing function:
+
+    >>> dev = qml.device("default.qubit", wires=2)
+    >>> vjps = fn([t.execute(dev) for t in vjp_tapes])
+    >>> vjps
+    [tensor([-0.6069, -0.0451,  0.0451, -0.0139, -0.2809,  0.2809],
+       dtype=torch.float64, grad_fn=<ViewBackward>),
+       tensor([ 0.1739, -0.1641, -0.0054, -0.2937, -0.4008,  0.0000],
+       dtype=torch.float64, grad_fn=<ViewBackward>)]
+
+    We have two VJPs; one per tape. Each one corresponds to the number of parameters
+    on the tapes (6).
+
+    The output VJPs are also differentiable with respect to the tape parameters:
+
+    >>> cost = torch.sum(vjps[0] + vjps[1])
+    >>> cost.backward()
+    >>> x.grad
+    tensor([[-4.7924e-01, -9.0857e-01, -2.4198e-01],
+            [-9.2973e-02, -1.0772e+00,  4.7184e-09]], dtype=torch.float64)
+    """
+    reshape_info = []
+    gradient_tapes = []
+    processing_fns = []
+
+    # Loop through the tapes and dys vector
+    for tape, dy in zip(tapes, dys):
+        g_tapes, fn = vjp(tape, dy, gradient_fn)
+
+        reshape_info.append(len(g_tapes))
+        processing_fns.append(fn)
+        gradient_tapes.extend(g_tapes)
+
+    def processing_fn(results):
+        vjps = []
+        start = 0
+
+        for t_idx in range(len(tapes)):
+            # extract the correct results from the flat list
+            res_len = reshape_info[t_idx]
+            res_t = results[start : start + res_len]
+            start += res_len
+
+            # postprocess results to compute the VJP
+            vjp_ = processing_fns[t_idx](res_t)
+
+            if vjp_ is None:
+                vjps.append(None)
+                continue
+
+            getattr(vjps, reduction)(vjp_)
+
+        return vjps
+
+    return gradient_tapes, processing_fn
diff --git a/pennylane/math/utils.py b/pennylane/math/utils.py
index 10d23cde5..985c92d78 100644
--- a/pennylane/math/utils.py
+++ b/pennylane/math/utils.py
@@ -55,9 +55,23 @@ def allequal(tensor1, tensor2, **kwargs):
 def allclose(a, b, rtol=1e-05, atol=1e-08, **kwargs):
     """Wrapper around np.allclose, allowing tensors ``a`` and ``b``
     to differ in type"""
-    t1 = ar.to_numpy(a)
-    t2 = ar.to_numpy(b)
-    return np.allclose(t1, t2, rtol=rtol, atol=atol, **kwargs)
+    try:
+        # Some frameworks may provide their own allclose implementation.
+        # Try and use it if available.
+        res = np.allclose(a, b, rtol=rtol, atol=atol, **kwargs)
+    except (TypeError, AttributeError):
+        # Otherwise, convert the input to NumPy arrays.
+        #
+        # TODO: replace this with a bespoke, framework agnostic
+        # low-level implementation to avoid the NumPy conversion:
+        #
+        #    np.abs(a - b) <= atol + rtol * np.abs(b)
+        #
+        t1 = ar.to_numpy(a)
+        t2 = ar.to_numpy(b)
+        res = np.allclose(t1, t2, rtol=rtol, atol=atol, **kwargs)
+
+    return res
 
 
 allclose.__doc__ = _np.allclose.__doc__


Test patch:
diff --git a/tests/gradients/test_vjp.py b/tests/gradients/test_vjp.py
new file mode 100644
index 000000000..4b1627e85
--- /dev/null
+++ b/tests/gradients/test_vjp.py
@@ -0,0 +1,458 @@
+# 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.
+"""Tests for the gradients.vjp module."""
+from functools import partial
+
+import pytest
+
+import pennylane as qml
+from pennylane import numpy as np
+from pennylane.gradients import param_shift
+
+
+class TestComputeVJP:
+    """Tests for the numeric computation of VJPs"""
+
+    def test_computation(self):
+        """Test that the correct VJP is returned"""
+        dy = np.array([[1.0, 2.0], [3.0, 4.0]])
+        jac = np.array([[[1.0, 0.1, 0.2], [0.2, 0.6, 0.1]], [[0.4, -0.7, 1.2], [-0.5, -0.6, 0.7]]])
+
+        vjp = qml.gradients.compute_vjp(dy, jac)
+
+        assert vjp.shape == (3,)
+        assert np.all(vjp == np.tensordot(dy, jac, axes=[[0, 1], [0, 1]]))
+
+    def test_jacobian_is_none(self):
+        """A None Jacobian returns a None VJP"""
+
+        dy = np.array([[1.0, 2.0], [3.0, 4.0]])
+        jac = None
+
+        vjp = qml.gradients.compute_vjp(dy, jac)
+        assert vjp is None
+
+    def test_zero_dy(self):
+        """A zero dy vector will return a zero matrix"""
+        dy = np.zeros([2, 2])
+        jac = np.array([[[1.0, 0.1, 0.2], [0.2, 0.6, 0.1]], [[0.4, -0.7, 1.2], [-0.5, -0.6, 0.7]]])
+
+        vjp = qml.gradients.compute_vjp(dy, jac)
+        assert np.all(vjp == np.zeros([3]))
+
+
+class TestVJP:
+    """Tests for the vjp function"""
+
+    def test_no_trainable_parameters(self):
+        """A tape with no trainable parameters will simply return None"""
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(0.4, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        tape.trainable_params = {}
+        dy = np.array([1.0])
+        tapes, fn = qml.gradients.vjp(tape, dy, param_shift)
+
+        assert not tapes
+        assert fn(tapes) is None
+
+    def test_zero_dy(self):
+        """A zero dy vector will return no tapes and a zero matrix"""
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(0.4, wires=0)
+            qml.RX(0.6, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        tape.trainable_params = {0, 1}
+        dy = np.array([0.0])
+        tapes, fn = qml.gradients.vjp(tape, dy, param_shift)
+
+        assert not tapes
+        assert np.all(fn(tapes) == np.zeros([len(tape.trainable_params)]))
+
+    def test_single_expectation_value(self, tol):
+        """Tests correct output shape and evaluation for a tape
+        with a single expval output"""
+        dev = qml.device("default.qubit", wires=2)
+        x = 0.543
+        y = -0.654
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        tape.trainable_params = {0, 1}
+        dy = np.array([1.0])
+
+        tapes, fn = qml.gradients.vjp(tape, dy, param_shift)
+        assert len(tapes) == 4
+
+        res = fn(dev.batch_execute(tapes))
+        assert res.shape == (2,)
+
+        expected = np.array([-np.sin(y) * np.sin(x), np.cos(y) * np.cos(x)])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_multiple_expectation_values(self, tol):
+        """Tests correct output shape and evaluation for a tape
+        with multiple expval outputs"""
+        dev = qml.device("default.qubit", wires=2)
+        x = 0.543
+        y = -0.654
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+            qml.expval(qml.PauliX(1))
+
+        tape.trainable_params = {0, 1}
+        dy = np.array([1.0, 2.0])
+
+        tapes, fn = qml.gradients.vjp(tape, dy, param_shift)
+        assert len(tapes) == 4
+
+        res = fn(dev.batch_execute(tapes))
+        assert res.shape == (2,)
+
+        expected = np.array([-np.sin(x), 2 * np.cos(y)])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_prob_expectation_values(self, tol):
+        """Tests correct output shape and evaluation for a tape
+        with prob and expval outputs"""
+        dev = qml.device("default.qubit", wires=2)
+        x = 0.543
+        y = -0.654
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+            qml.probs(wires=[0, 1])
+
+        tape.trainable_params = {0, 1}
+        dy = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
+
+        tapes, fn = qml.gradients.vjp(tape, dy, param_shift)
+        assert len(tapes) == 4
+
+        res = fn(dev.batch_execute(tapes))
+        assert res.shape == (2,)
+
+        expected = (
+            np.array(
+                [
+                    [-2 * np.sin(x), 0],
+                    [
+                        -(np.cos(y / 2) ** 2 * np.sin(x)),
+                        -(np.cos(x / 2) ** 2 * np.sin(y)),
+                    ],
+                    [
+                        -(np.sin(x) * np.sin(y / 2) ** 2),
+                        (np.cos(x / 2) ** 2 * np.sin(y)),
+                    ],
+                    [
+                        (np.sin(x) * np.sin(y / 2) ** 2),
+                        (np.sin(x / 2) ** 2 * np.sin(y)),
+                    ],
+                    [
+                        (np.cos(y / 2) ** 2 * np.sin(x)),
+                        -(np.sin(x / 2) ** 2 * np.sin(y)),
+                    ],
+                ]
+            )
+            / 2
+        )
+
+        assert np.allclose(res, dy @ expected, atol=tol, rtol=0)
+
+
+def expected(params):
+    x, y = 1.0 * params
+    return (
+        np.array(
+            [
+                (np.cos(y / 2) ** 2 * np.sin(x)) + (np.cos(y / 2) ** 2 * np.sin(x)),
+                (np.cos(x / 2) ** 2 * np.sin(y)) - (np.sin(x / 2) ** 2 * np.sin(y)),
+            ]
+        )
+        / 2
+    )
+
+
+def ansatz(x, y):
+    qml.RX(x, wires=[0])
+    qml.RY(y, wires=[1])
+    qml.CNOT(wires=[0, 1])
+    qml.probs(wires=[0, 1])
+
+
+class TestVJPGradients:
+    """Gradient tests for the vjp function"""
+
+    def test_autograd(self, tol):
+        """Tests that the output of the VJP transform
+        can be differentiated using autograd."""
+        dev = qml.device("default.qubit.autograd", wires=2)
+        params = np.array([0.543, -0.654], requires_grad=True)
+
+        def cost_fn(x, dy):
+            with qml.tape.JacobianTape() as tape:
+                ansatz(x[0], x[1])
+
+            tape.trainable_params = {0, 1}
+            tapes, fn = qml.gradients.vjp(tape, dy, param_shift)
+            vjp = fn(dev.batch_execute(tapes))
+            return vjp
+
+        dy = np.array([-1.0, 0.0, 0.0, 1.0], requires_grad=False)
+        res = cost_fn(params, dy)
+        assert np.allclose(res, expected(params), atol=tol, rtol=0)
+
+        res = qml.jacobian(cost_fn)(params, dy)
+        assert np.allclose(res, qml.jacobian(expected)(params), atol=tol, rtol=0)
+
+    def test_torch(self, tol):
+        """Tests that the output of the VJP transform
+        can be differentiated using Torch."""
+        torch = pytest.importorskip("torch")
+        from pennylane.interfaces.torch import TorchInterface
+
+        dev = qml.device("default.qubit", wires=2)
+
+        params = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64)
+        dy = torch.tensor([-1.0, 0.0, 0.0, 1.0], dtype=torch.float64)
+
+        with TorchInterface.apply(qml.tape.QubitParamShiftTape()) as tape:
+            ansatz(params[0], params[1])
+
+        tape.trainable_params = {0, 1}
+        tapes, fn = qml.gradients.vjp(tape, dy, param_shift)
+        vjp = fn([t.execute(dev) for t in tapes])
+
+        assert np.allclose(vjp.detach(), expected(params.detach()), atol=tol, rtol=0)
+
+        cost = vjp[0]
+        cost.backward()
+
+        exp = qml.jacobian(lambda x: expected(x)[0])(params.detach().numpy())
+        assert np.allclose(params.grad, exp, atol=tol, rtol=0)
+
+    def test_tf(self, tol):
+        """Tests that the output of the VJP transform
+        can be differentiated using TF."""
+        tf = pytest.importorskip("tensorflow")
+
+        dev = qml.device("default.qubit.tf", wires=2)
+
+        params = tf.Variable([0.543, -0.654], dtype=tf.float64)
+        dy = tf.constant([-1.0, 0.0, 0.0, 1.0], dtype=tf.float64)
+
+        with tf.GradientTape() as t:
+            with qml.tape.JacobianTape() as tape:
+                ansatz(params[0], params[1])
+
+            tape.trainable_params = {0, 1}
+            tapes, fn = qml.gradients.vjp(tape, dy, param_shift)
+            vjp = fn(dev.batch_execute(tapes))
+
+        assert np.allclose(vjp, expected(params), atol=tol, rtol=0)
+
+        res = t.jacobian(vjp, params)
+        assert np.allclose(res, qml.jacobian(expected)(params.numpy()), atol=tol, rtol=0)
+
+    def test_jax(self, tol):
+        """Tests that the output of the VJP transform
+        can be differentiated using JAX."""
+        jax = pytest.importorskip("jax")
+        from jax import numpy as jnp
+
+        dev = qml.device("default.qubit.jax", wires=2)
+        params = jnp.array([0.543, -0.654])
+
+        @partial(jax.jit, static_argnums=1)
+        def cost_fn(x, dy):
+            with qml.tape.JacobianTape() as tape:
+                ansatz(x[0], x[1])
+
+            tape.trainable_params = {0, 1}
+            tapes, fn = qml.gradients.vjp(tape, dy, param_shift)
+            vjp = fn(dev.batch_execute(tapes))
+            return vjp
+
+        dy = (-1.0, 0.0, 0.0, 1.0)
+        res = cost_fn(params, dy)
+        assert np.allclose(res, expected(params), atol=tol, rtol=0)
+
+        res = jax.jacobian(cost_fn, argnums=0)(params, dy)
+        exp = qml.jacobian(expected)(np.array(params))
+        assert np.allclose(res, exp, atol=tol, rtol=0)
+
+
+class TestBatchVJP:
+    """Tests for the batch VJP function"""
+
+    def test_one_tape_no_trainable_parameters(self):
+        """A tape with no trainable parameters will simply return None"""
+        dev = qml.device("default.qubit", wires=2)
+
+        with qml.tape.QuantumTape() as tape1:
+            qml.RX(0.4, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        with qml.tape.JacobianTape() as tape2:
+            qml.RX(0.4, wires=0)
+            qml.RX(0.6, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        tape1.trainable_params = {}
+        tape2.trainable_params = {0, 1}
+
+        tapes = [tape1, tape2]
+        dys = [np.array([1.0]), np.array([1.0])]
+
+        v_tapes, fn = qml.gradients.batch_vjp(tapes, dys, param_shift)
+        assert len(v_tapes) == 4
+
+        # Even though there are 3 parameters, only two contribute
+        # to the VJP, so only 2*2=4 quantum evals
+        res = fn(dev.batch_execute(v_tapes))
+        assert res[0] is None
+        assert res[1] is not None
+
+    def test_all_tapes_no_trainable_parameters(self):
+        """If all tapes have no trainable parameters all outputs will be None"""
+        dev = qml.device("default.qubit", wires=2)
+
+        with qml.tape.QuantumTape() as tape1:
+            qml.RX(0.4, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        with qml.tape.QuantumTape() as tape2:
+            qml.RX(0.4, wires=0)
+            qml.RX(0.6, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        tape1.trainable_params = set()
+        tape2.trainable_params = set()
+
+        tapes = [tape1, tape2]
+        dys = [np.array([1.0]), np.array([1.0])]
+
+        v_tapes, fn = qml.gradients.batch_vjp(tapes, dys, param_shift)
+
+        assert v_tapes == []
+        assert fn([]) == [None, None]
+
+    def test_zero_dy(self):
+        """A zero dy vector will return no tapes and a zero matrix"""
+        dev = qml.device("default.qubit", wires=2)
+
+        with qml.tape.QuantumTape() as tape1:
+            qml.RX(0.4, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        with qml.tape.JacobianTape() as tape2:
+            qml.RX(0.4, wires=0)
+            qml.RX(0.6, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        tape1.trainable_params = {0}
+        tape2.trainable_params = {0, 1}
+
+        tapes = [tape1, tape2]
+        dys = [np.array([0.0]), np.array([1.0])]
+
+        v_tapes, fn = qml.gradients.batch_vjp(tapes, dys, param_shift)
+        res = fn(dev.batch_execute(v_tapes))
+
+        # Even though there are 3 parameters, only two contribute
+        # to the VJP, so only 2*2=4 quantum evals
+        assert len(v_tapes) == 4
+        assert np.allclose(res[0], 0)
+
+    def test_reduction_append(self):
+        """Test the 'append' reduction strategy"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        with qml.tape.JacobianTape() as tape1:
+            qml.RX(0.4, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        with qml.tape.JacobianTape() as tape2:
+            qml.RX(0.4, wires=0)
+            qml.RX(0.6, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        tape1.trainable_params = {0}
+        tape2.trainable_params = {0, 1}
+
+        tapes = [tape1, tape2]
+        dys = [np.array([1.0]), np.array([1.0])]
+
+        v_tapes, fn = qml.gradients.batch_vjp(tapes, dys, param_shift, reduction="append")
+        res = fn(dev.batch_execute(v_tapes))
+
+        # Returned VJPs will be appended to a list, one vjp per tape
+        assert len(res) == 2
+        assert all(isinstance(r, np.ndarray) for r in res)
+        assert all(len(r) == len(t.trainable_params) for t, r in zip(tapes, res))
+
+    def test_reduction_extend(self):
+        """Test the 'extend' reduction strategy"""
+        dev = qml.device("default.qubit", wires=2)
+
+        with qml.tape.JacobianTape() as tape1:
+            qml.RX(0.4, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        with qml.tape.JacobianTape() as tape2:
+            qml.RX(0.4, wires=0)
+            qml.RX(0.6, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        tape1.trainable_params = {0}
+        tape2.trainable_params = {0, 1}
+
+        tapes = [tape1, tape2]
+        dys = [np.array([1.0]), np.array([1.0])]
+
+        v_tapes, fn = qml.gradients.batch_vjp(tapes, dys, param_shift, reduction="extend")
+        res = fn(dev.batch_execute(v_tapes))
+
+        # Returned VJPs will be extended into a list. Each element of the returned
+        # list will correspond to a single input parameter of the combined
+        # tapes.
+        assert len(res) == sum(len(t.trainable_params) for t in tapes)