Repository: PennyLaneAI/pennylane
Pull request #1589 (2021-08-25 13:39:16)

Add a simple API for gradient transforms
**Context:** With a simple batch-tape transform added in #1493, we now have the ability to register custom gradient transforms.

**Description of the Change:**

- Adds a `@qml.gradient_transform` decorator. This simply subclasses from `batch_transform`, with one change: a smart default `expand_fn` is provided, which automatically expands tape operations which are not differentiable.

- The existing gradient transforms `param_shift`, `param_shift_cv`, and `finite_diff` are now decorated with the `qml.gradient_transform` decorator.

**Benefits:**

- It is now possible to check if a gradient transform is a gradient transform via `isinstance(fn, qml.gradient_transform)`.

- When applying a gradient transform, operations without gradient recipes are automatically decomposed.

- gradient transforms can now be applied directly to QNodes, in addition to also acting on tapes:

  ```pycon
  >>> @qml.qnode(dev)
  ... def circuit(x):
  ...     qml.RX(x, wires=0)
  ...     qml.CNOT(wires=[0, 1])
  ...     return qml.expval(qml.PauliZ(0))
  >>> circuit(0.3)
  tensor(0.95533649, requires_grad=True)
  >>> qml.gradients.param_shift(circuit)(0.5)
  array([[-0.47942554]])
  ```

  As before, quantum gradient transforms are fully differentiable, allowing higher order derivatives to be accessed:

  ```pycon
  >>> qml.grad(qml.gradients.param_shift(circuit))(0.5)
  tensor(-0.87758256, requires_grad=True)
  ```

- There is now a documented API for adding custom quantum gradients to PL.

**Possible Drawbacks:**

- The default `expand_fn` only considers decomposing ops that have `grad_method="None"`. But what happens if `grad_method="F"`? How do we know whether the operation decomposes down into ops with `grad_method="A"`?

**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/1589?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1589](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1589?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (f168130) into [custom_gradient](https://codecov.io/gh/PennyLaneAI/pennylane/commit/0ccb881d2e522b58916c7d2d2a13c0ac92470ea0?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (0ccb881) will **decrease** coverage by `0.01%`.
> The diff coverage is `91.66%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1589/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/1589?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@                 Coverage Diff                 @@
##           custom_gradient    #1589      +/-   ##
===================================================
- Coverage            99.17%   99.15%   -0.02%     
===================================================
  Files                  191      192       +1     
  Lines                13771    13793      +22     
===================================================
+ Hits                 13657    13677      +20     
- Misses                 114      116       +2     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1589?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/gradients/gradient\_transform.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1589/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9ncmFkaWVudF90cmFuc2Zvcm0ucHk=) | `85.71% <85.71%> (ø)` | |
| [pennylane/gradients/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1589/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9fX2luaXRfXy5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/gradients/finite\_difference.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1589/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9maW5pdGVfZGlmZmVyZW5jZS5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/gradients/parameter\_shift.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1589/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9wYXJhbWV0ZXJfc2hpZnQucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/gradients/parameter\_shift\_cv.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1589/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9wYXJhbWV0ZXJfc2hpZnRfY3YucHk=) | `99.13% <100.00%> (+<0.01%)` | :arrow_up: |
| [pennylane/transforms/batch\_transform.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1589/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvYmF0Y2hfdHJhbnNmb3JtLnB5) | `100.00% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1589?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/1589?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [0ccb881...f168130](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1589?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).

anthayes92: Looks good @josh146 ! Just a couple of small things (and a few points I didn't understand so wanted to clarify). Another awesome transform :rocket: 

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 7d16275a3..269f71d8f 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -2,6 +2,52 @@
 
 <h3>New features since last release</h3>
 
+* Custom gradient transforms can now be created using the new
+  `@qml.gradients.gradient_transform` decorator on a batch-tape transform.
+  [(#1589)](https://github.com/PennyLaneAI/pennylane/pull/1589)
+
+  Quantum gradient transforms are a specific case of `qml.batch_transform`.
+  To create a quantum gradient transform, simply write a function that accepts a tape,
+  and returns a batch of tapes to be independently executed on a quantum device, alongside
+  a post-processing function that processes the tape results into the gradient.
+
+  Furthermore, a smart default expansion function is provided, which automatically expands tape
+  operations which are not differentiable prior to applying the quantum gradient.
+  All gradient transforms in `qml.gradients` are now decorated with this decorator.
+
+  Supported gradient transforms must be of the following form:
+
+  ```python
+  @qml.gradients.gradient_transform
+  def my_custom_gradient(tape, argnum=None, **kwargs):
+      ...
+      return gradient_tapes, processing_fn
+  ```
+
+  Various built-in quantum gradient transforms are provided within the
+  `qml.gradients` module, including `qml.gradients.param_shift`.
+  Once defined, quantum gradient transforms can be applied directly
+  to QNodes:
+
+  ```pycon
+  >>> @qml.qnode(dev)
+  ... def circuit(x):
+  ...     qml.RX(x, wires=0)
+  ...     qml.CNOT(wires=[0, 1])
+  ...     return qml.expval(qml.PauliZ(0))
+  >>> circuit(0.3)
+  tensor(0.95533649, requires_grad=True)
+  >>> qml.gradients.param_shift(circuit)(0.5)
+  array([[-0.47942554]])
+  ```
+
+  Quantum gradient transforms are fully differentiable, allowing higher order derivatives to be
+  accessed:
+
+  ```pycon
+  >>> qml.grad(qml.gradients.param_shift(circuit))(0.5)
+  tensor(-0.87758256, requires_grad=True)
+  ```
 
 * A new pytorch device, `qml.device('default.qubit.torch', wires=wires)`, supports
   backpropogation with the torch interface.
diff --git a/pennylane/_grad.py b/pennylane/_grad.py
index d65c9ef79..0c21491e7 100644
--- a/pennylane/_grad.py
+++ b/pennylane/_grad.py
@@ -180,7 +180,14 @@ def jacobian(func, argnum=None):
         if len(argnum) == 1:
             return _jacobian(func, argnum[0])(*args, **kwargs)
 
-        return np.stack([_jacobian(func, arg)(*args, **kwargs) for arg in argnum]).T
+        jacobians = [_jacobian(func, arg)(*args, **kwargs) for arg in argnum]
+
+        try:
+            return np.stack(jacobians).T
+        except ValueError:
+            # The Jacobian of each argument is a different shape and cannot
+            # be stacked; simply return the tuple of argument Jacobians.
+            return tuple(jacobians)
 
     return _jacobian_function
 
diff --git a/pennylane/gradients/__init__.py b/pennylane/gradients/__init__.py
index b2bc5c2d7..b95b24f3d 100644
--- a/pennylane/gradients/__init__.py
+++ b/pennylane/gradients/__init__.py
@@ -18,6 +18,7 @@ from . import finite_difference
 from . import parameter_shift
 from . import parameter_shift_cv
 
+from .gradient_transform import gradient_transform
 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
diff --git a/pennylane/gradients/finite_difference.py b/pennylane/gradients/finite_difference.py
index f823d761f..098fe5e91 100644
--- a/pennylane/gradients/finite_difference.py
+++ b/pennylane/gradients/finite_difference.py
@@ -23,6 +23,8 @@ from scipy.special import factorial
 
 import pennylane as qml
 
+from .gradient_transform import gradient_transform
+
 
 @functools.lru_cache(maxsize=None)
 def finite_diff_coeffs(n, approx_order, strategy):
@@ -179,6 +181,7 @@ def generate_shifted_tapes(tape, idx, shifts, multipliers=None):
     return tapes
 
 
+@gradient_transform
 def finite_diff(tape, argnum=None, h=1e-7, approx_order=1, n=1, strategy="forward", f0=None):
     r"""Generate the finite-difference tapes and postprocessing methods required
     to compute the gradient of a gate parameter with respect to its outputs.
diff --git a/pennylane/gradients/gradient_transform.py b/pennylane/gradients/gradient_transform.py
new file mode 100644
index 000000000..d70ef25f5
--- /dev/null
+++ b/pennylane/gradients/gradient_transform.py
@@ -0,0 +1,179 @@
+# 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 utilities for defining custom gradient transforms,
+including a decorator for specifying gradient expansions."""
+# pylint: disable=too-few-public-methods
+import pennylane as qml
+
+
+unsupported_op = lambda op: op.grad_method is None
+supported_op = lambda op: op.grad_method is not None
+trainable_op = lambda op: any(qml.math.requires_grad(p) for p in op.parameters)
+
+
+def gradient_expand(tape, depth=10):
+    """Expand out a tape so that it supports differentiation
+    of requested operations.
+
+    This is achieved by decomposing all trainable operations that have
+    ``Operation.grad_method=None`` until all resulting operations
+    have a defined gradient method, up to maximum depth ``depth``. Note that this
+    might not be possible, in which case the gradient rule will fail to apply.
+
+    Args:
+        tape (.QuantumTape): the input tape to expand
+        depth (int) : the maximum expansion depth
+
+    Returns:
+        .QuantumTape: the expanded tape
+    """
+
+    # check if the tape contains unsupported trainable operations
+    if any(unsupported_op(op) and trainable_op(op) for op in tape.operations):
+
+        # Define the stopping condition for the expansion
+        stop_cond = lambda obj: (
+            not isinstance(obj, qml.measure.MeasurementProcess)
+            and ((supported_op(obj) and trainable_op(obj)) or not trainable_op(obj))
+        )
+
+        return tape.expand(depth=depth, stop_at=stop_cond)
+
+    return tape
+
+
+class gradient_transform(qml.batch_transform):
+    """Decorator for defining quantum gradient transforms.
+
+    Quantum gradient transforms are a specific case of :class:`~.batch_transform`.
+    All quantum gradient transforms accept a tape, and output
+    a batch of tapes to be independently executed on a quantum device, alongside
+    a post-processing function that returns the result.
+
+    Args:
+        expand_fn (function): An expansion function (if required) to be applied to the
+            input tape before the gradient computation takes place. If not provided,
+            the default expansion function simply expands all operations that
+            have ``Operation.grad_method=None`` until all resulting operations
+            have a defined gradient method.
+        differentiable (bool): Specifies whether the gradient transform is differentiable or
+            not. A transform may be non-differentiable if it does not use an
+            autodiff framework for its tensor manipulations. In such a case, setting
+            ``differentiable=False`` instructs the decorator
+            to mark the output as 'constant', reducing potential overhead.
+        hybrid (bool): Specifies whether classical processing inside a QNode
+            should be taken into account when transforming a QNode.
+
+            - If ``True``, and classical processing is detected and this
+              option is set to ``True``, the Jacobian of the classical
+              processing will be computed and included. When evaluated, the
+              returned Jacobian will be with respect to the QNode arguments.
+
+            - If ``False``, any internal QNode classical processing will be
+              **ignored**. When evaluated, the returned Jacobian will be with
+              respect to the **gate** arguments, and not the QNode arguments.
+
+    Supported gradient transforms must be of the following form:
+
+    .. code-block:: python
+
+        @gradient_transform
+        def my_custom_gradient(tape, argnum=None, **kwargs):
+            ...
+            return gradient_tapes, processing_fn
+
+    where:
+
+    - ``tape`` (*QuantumTape*): the input quantum tape to compute the gradient of
+
+    - ``argnum`` (*int* or *list[int]* or *None*): Which trainable parameters of the tape
+      to differentiate with respect to. If not provided, the derivatives with respect to all
+      trainable inputs of the tape should be returned (``tape.trainable_params``).
+
+    - ``gradient_tapes`` (*List[QuantumTape]*): is a list of output tapes to be evaluated.
+      If this list is empty, no quantum evaluations will be made.
+
+    - ``processing_fn`` is a processing function to be applied to the output of the evaluated
+      ``gradient_tapes``. It should accept a list of numeric results with length ``len(gradient_tapes)``,
+      and return the Jacobian matrix.
+
+    Once defined, the quantum gradient transform can be used as follows:
+
+    >>> gradient_tapes, processing_fn = my_custom_gradient(tape, *gradient_kwargs)
+    >>> res = execute(tapes, dev, interface="autograd", gradient_fn=qml.gradients.param_shift)
+    >>> jacobian = processing_fn(res)
+
+    Alternatively, gradient transforms can be applied directly to QNodes,
+    in which case the execution is implicit:
+
+    >>> fn = my_custom_gradient(qnode, *gradient_kwargs)
+    >>> fn(weights) # transformed function takes the same arguments as the QNode
+    1.2629730888100839
+
+    .. note::
+
+        The input tape might have parameters of various types, including
+        NumPy arrays, JAX DeviceArrays, and TensorFlow and PyTorch tensors.
+
+        If the gradient transform is written in a autodiff-compatible manner, either by
+        using a framework such as Autograd or TensorFlow, or by using ``qml.math`` for
+        tensor manipulation, then higher-order derivatives will also be supported.
+
+        Alternatively, you may use the ``tape.unwrap()`` context manager to temporarily
+        convert all tape parameters to NumPy arrays and floats:
+
+        >>> with tape.unwrap():
+        ...     params = tape.get_parameters()  # list of floats
+    """
+
+    def __init__(self, transform_fn, expand_fn=gradient_expand, differentiable=True, hybrid=True):
+        self.hybrid = hybrid
+        super().__init__(transform_fn, expand_fn=expand_fn, differentiable=differentiable)
+
+    def qnode_execution_wrapper(self, qnode, targs, tkwargs):
+        # Here, we overwrite the QNode execution wrapper in order
+        # to take into account that classical processing may be present
+        # inside the QNode.
+        hybrid = tkwargs.pop("hybrid", self.hybrid)
+        _wrapper = super().qnode_execution_wrapper(qnode, targs, tkwargs)
+        cjac_fn = qml.transforms.classical_jacobian(qnode)
+
+        def jacobian_wrapper(*args, **kwargs):
+            qjac = _wrapper(*args, **kwargs)
+            cjac = cjac_fn(*args, **kwargs)
+
+            if any(m.return_type is qml.operation.Probability for m in qnode.qtape.measurements):
+                qjac = qml.math.squeeze(qjac)
+
+            if isinstance(cjac, tuple):
+                # Classical processing of multiple arguments is present. Return qjac @ cjac.
+                jacs = [
+                    qml.math.squeeze(qml.math.tensordot(c, qjac, [[0], [-1]]))
+                    for c in cjac
+                    if c is not None
+                ]
+                return jacs
+
+            is_square = cjac.shape == (1,) or (cjac.ndim == 2 and cjac.shape[0] == cjac.shape[1])
+
+            if not hybrid or (is_square and qml.math.allclose(cjac, qml.numpy.eye(cjac.shape[0]))):
+                # Classical Jacobian is the identity. No classical processing
+                # is present inside the QNode.
+                return qjac
+
+            # Classical processing of a single argument is present. Return qjac @ cjac.
+            jac = qml.math.squeeze(qml.math.tensordot(qml.math.T(cjac), qjac, [[-1], [-1]]))
+            return qml.math.T(jac)
+
+        return jacobian_wrapper
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
index 66de8551a..18f506ee6 100644
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -20,6 +20,7 @@ import numpy as np
 
 import pennylane as qml
 
+from .gradient_transform import gradient_transform
 from .finite_difference import finite_diff, generate_shifted_tapes
 
 
@@ -349,6 +350,7 @@ def var_param_shift(tape, argnum, shift=np.pi / 2, gradient_recipes=None, f0=Non
     return gradient_tapes, processing_fn
 
 
+@gradient_transform
 def param_shift(
     tape, argnum=None, shift=np.pi / 2, gradient_recipes=None, fallback_fn=finite_diff, f0=None
 ):
diff --git a/pennylane/gradients/parameter_shift_cv.py b/pennylane/gradients/parameter_shift_cv.py
index 5db6b6b45..73e932961 100644
--- a/pennylane/gradients/parameter_shift_cv.py
+++ b/pennylane/gradients/parameter_shift_cv.py
@@ -23,6 +23,7 @@ import numpy as np
 
 import pennylane as qml
 
+from .gradient_transform import gradient_transform
 from .finite_difference import finite_diff, generate_shifted_tapes
 from .parameter_shift import expval_param_shift, _get_operation_recipe, _process_gradient_recipe
 
@@ -460,6 +461,7 @@ def second_order_param_shift(tape, dev_wires, argnum=None, shift=np.pi / 2, grad
     return gradient_tapes, processing_fn
 
 
+@gradient_transform
 def param_shift_cv(
     tape,
     dev,
diff --git a/pennylane/interfaces/batch/autograd.py b/pennylane/interfaces/batch/autograd.py
index 4cb012573..56bf3f276 100644
--- a/pennylane/interfaces/batch/autograd.py
+++ b/pennylane/interfaces/batch/autograd.py
@@ -109,7 +109,7 @@ def _execute(
     for i, r in enumerate(res):
         res[i] = np.tensor(r)
 
-        if r.dtype == np.dtype("object"):
+        if res[i].dtype == np.dtype("object"):
             # For backwards compatibility, we flatten ragged tape outputs
             res[i] = np.hstack(r)
 
diff --git a/pennylane/interfaces/batch/torch.py b/pennylane/interfaces/batch/torch.py
index 8a34d234d..6786569e7 100644
--- a/pennylane/interfaces/batch/torch.py
+++ b/pennylane/interfaces/batch/torch.py
@@ -95,16 +95,14 @@ class ExecuteTapes(torch.autograd.Function):
                 break
 
         for i, r in enumerate(res):
-            if r.dtype == np.dtype("object"):
+            if r.dtype is np.dtype("object"):
                 # For backwards compatibility, we flatten ragged tape outputs
                 r = np.hstack(r)
 
-            res[i] = torch.as_tensor(torch.from_numpy(r), device=ctx.torch_device)
+            res[i] = torch.as_tensor(r, device=ctx.torch_device)
 
             if ctx.jacs:
-                ctx.jacs[i] = torch.as_tensor(
-                    torch.from_numpy(ctx.jacs[i]), device=ctx.torch_device
-                )
+                ctx.jacs[i] = torch.as_tensor(ctx.jacs[i], device=ctx.torch_device)
 
         return tuple(res)
 
diff --git a/pennylane/transforms/batch_transform.py b/pennylane/transforms/batch_transform.py
index 6e05619eb..d5261a1ff 100644
--- a/pennylane/transforms/batch_transform.py
+++ b/pennylane/transforms/batch_transform.py
@@ -20,38 +20,6 @@ import pennylane as qml
 from pennylane.interfaces.batch import execute
 
 
-def _create_qnode_internal_wrapper(qnode, transform, targs, tkwargs):
-    def _wrapper(*args, **kwargs):
-        qnode.construct(args, kwargs)
-        tapes, processing_fn = transform.construct(qnode.qtape, *targs, **tkwargs)
-
-        # TODO: work out what to do for backprop
-        interface = qnode.interface
-
-        # TODO: extract gradient_fn from QNode
-        gradient_fn = qnode.diff_method
-
-        if interface is None or not transform.differentiable:
-            gradient_fn = None
-
-        elif gradient_fn in ("best", "parameter-shift"):
-            gradient_fn = qml.gradients.param_shift
-
-        elif gradient_fn == "finite-diff":
-            gradient_fn = qml.gradients.finite_diff
-
-        res = execute(
-            tapes,
-            device=qnode.device,
-            gradient_fn=gradient_fn,
-            interface=interface,
-        )
-
-        return processing_fn(res)
-
-    return _wrapper
-
-
 class batch_transform:
     r"""Class for registering a tape transform that takes a tape, and outputs
     a batch of tapes to be independently executed on a quantum device.
@@ -195,16 +163,55 @@ class batch_transform:
         self.differentiable = differentiable
         functools.update_wrapper(self, transform_fn)
 
+    def qnode_execution_wrapper(self, qnode, targs, tkwargs):
+        """A wrapper method that takes a QNode and transform arguments,
+        and returns a function that 'wraps' the QNode execution.
+
+        The returned function should accept the same keyword arguments as
+        the QNode, and return the output of the applying the tape transform
+        to the QNode's constructed tape.
+        """
+
+        def _wrapper(*args, **kwargs):
+            qnode.construct(args, kwargs)
+            tapes, processing_fn = self.construct(qnode.qtape, *targs, **tkwargs)
+
+            # TODO: work out what to do for backprop
+            interface = qnode.interface
+
+            # TODO: extract gradient_fn from QNode
+            gradient_fn = qnode.diff_method
+
+            if interface is None or not self.differentiable:
+                gradient_fn = None
+
+            elif gradient_fn in ("best", "parameter-shift"):
+                gradient_fn = qml.gradients.param_shift
+
+            elif gradient_fn == "finite-diff":
+                gradient_fn = qml.gradients.finite_diff
+
+            res = execute(
+                tapes,
+                device=qnode.device,
+                gradient_fn=gradient_fn,
+                interface=interface,
+            )
+
+            return processing_fn(res)
+
+        return _wrapper
+
     def __call__(self, qnode, *targs, **tkwargs):
         if isinstance(qnode, qml.tape.QuantumTape):
             # Input is a quantum tape.
             # tapes, fn = some_transform(tape, *transform_args)
-            return self.construct(qnode, *targs, *tkwargs)
+            return self.construct(qnode, *targs, **tkwargs)
 
         if isinstance(qnode, qml.QNode):
             # Input is a QNode:
             # result = some_transform(qnode, *transform_args)(*qnode_args)
-            wrapper = _create_qnode_internal_wrapper(qnode, self, targs, tkwargs)
+            wrapper = self.qnode_execution_wrapper(qnode, targs, tkwargs)
             wrapper = functools.wraps(qnode)(wrapper)
 
         else:
@@ -226,7 +233,7 @@ class batch_transform:
             targs = (qnode,) + targs
 
             def wrapper(qnode):
-                _wrapper = _create_qnode_internal_wrapper(qnode, self, targs, tkwargs)
+                _wrapper = self.qnode_execution_wrapper(qnode, targs, tkwargs)
                 _wrapper = functools.wraps(qnode)(_wrapper)
                 return _wrapper
 
@@ -247,7 +254,9 @@ class batch_transform:
             tuple[list[tapes], callable]: list of transformed tapes
             to execute and a post-processing function.
         """
-        if self.expand_fn is not None:
+        expand = kwargs.pop("_expand", True)
+
+        if expand and self.expand_fn is not None:
             tape = self.expand_fn(tape)
 
         tapes, processing_fn = self.transform_fn(tape, *args, **kwargs)


Test patch:
diff --git a/tests/gradients/test_finite_difference.py b/tests/gradients/test_finite_difference.py
index 0c4ce036b..04302f18c 100644
--- a/tests/gradients/test_finite_difference.py
+++ b/tests/gradients/test_finite_difference.py
@@ -134,7 +134,7 @@ class TestFiniteDiff:
     def test_non_differentiable_error(self):
         """Test error raised if attempting to differentiate with
         respect to a non-differentiable argument"""
-        psi = np.array([1, 0, 1, 0]) / np.sqrt(2)
+        psi = np.array([1, 0, 1, 0], requires_grad=False) / np.sqrt(2)
 
         with qml.tape.JacobianTape() as tape:
             qml.QubitStateVector(psi, wires=[0, 1])
@@ -147,7 +147,7 @@ class TestFiniteDiff:
         with pytest.raises(
             ValueError, match=r"Cannot differentiate with respect to parameter\(s\) {0}"
         ):
-            finite_diff(tape)
+            finite_diff(tape, _expand=False)
 
         # setting trainable parameters avoids this
         tape.trainable_params = {1, 2}
diff --git a/tests/gradients/test_gradient_transform.py b/tests/gradients/test_gradient_transform.py
new file mode 100644
index 000000000..260446f4f
--- /dev/null
+++ b/tests/gradients/test_gradient_transform.py
@@ -0,0 +1,403 @@
+# 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.gradient_transform module."""
+import pytest
+
+import pennylane as qml
+from pennylane import numpy as np
+from pennylane.gradients.gradient_transform import gradient_transform, gradient_expand
+
+
+class TestGradientExpand:
+    """Tests for the gradient expand function"""
+
+    def test_no_expansion(self, mocker):
+        """Test that a circuit with differentiable
+        operations is not expanded"""
+        x = np.array(0.2, requires_grad=True)
+        y = np.array(0.1, requires_grad=True)
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(x, wires=0)
+            qml.RY(y, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        spy = mocker.spy(tape, "expand")
+        new_tape = gradient_expand(tape)
+
+        assert new_tape is tape
+        spy.assert_not_called()
+
+    def test_trainable_nondiff_expansion(self, mocker):
+        """Test that a circuit with non-differentiable
+        trainable operations is expanded"""
+        x = np.array(0.2, requires_grad=True)
+        y = np.array(0.1, requires_grad=True)
+
+        class NonDiffPhaseShift(qml.PhaseShift):
+            grad_method = None
+
+        with qml.tape.QuantumTape() as tape:
+            NonDiffPhaseShift(x, wires=0)
+            qml.RY(y, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        spy = mocker.spy(tape, "expand")
+        new_tape = gradient_expand(tape)
+
+        assert new_tape is not tape
+        spy.assert_called()
+
+        new_tape.operations[0].name == "RZ"
+        new_tape.operations[0].grad_method == "A"
+        new_tape.operations[1].name == "RY"
+        new_tape.operations[2].name == "CNOT"
+
+    def test_nontrainable_nondiff(self, mocker):
+        """Test that a circuit with non-differentiable
+        non-trainable operations is not expanded"""
+        x = np.array(0.2, requires_grad=False)
+        y = np.array(0.1, requires_grad=True)
+
+        class NonDiffPhaseShift(qml.PhaseShift):
+            grad_method = None
+
+        with qml.tape.QuantumTape() as tape:
+            NonDiffPhaseShift(x, wires=0)
+            qml.RY(y, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        spy = mocker.spy(tape, "expand")
+        new_tape = gradient_expand(tape)
+
+        assert new_tape is tape
+        spy.assert_not_called()
+
+    def test_trainable_numeric(self, mocker):
+        """Test that a circuit with numeric differentiable
+        trainable operations is *not* expanded"""
+        x = np.array(0.2, requires_grad=True)
+        y = np.array(0.1, requires_grad=True)
+
+        class NonDiffPhaseShift(qml.PhaseShift):
+            grad_method = "F"
+
+        with qml.tape.QuantumTape() as tape:
+            NonDiffPhaseShift(x, wires=0)
+            qml.RY(y, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        spy = mocker.spy(tape, "expand")
+        new_tape = gradient_expand(tape)
+
+        assert new_tape is tape
+        spy.assert_not_called()
+
+
+class TestGradientTransformIntegration:
+    """Test integration of the gradient transform decorator"""
+
+    def test_acting_on_qnodes(self, tol):
+        """Test that a gradient transform acts on QNodes
+        correctly"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(weights):
+            qml.RX(weights[0], wires=[0])
+            qml.RY(weights[1], wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliX(1))
+
+        grad_fn = qml.gradients.param_shift(circuit)
+
+        w = np.array([0.543, -0.654], requires_grad=True)
+        res = grad_fn(w)
+
+        x, y = w
+        expected = np.array([[-np.sin(x), 0], [0, -2 * np.cos(y) * np.sin(y)]])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_decorator(self, tol):
+        """Test that a gradient transform decorating a QNode
+        acts correctly"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.gradients.param_shift
+        @qml.qnode(dev)
+        def circuit(weights):
+            qml.RX(weights[0], wires=[0])
+            qml.RY(weights[1], wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliX(1))
+
+        w = np.array([0.543, -0.654], requires_grad=True)
+        res = circuit(w)
+
+        x, y = w
+        expected = np.array([[-np.sin(x), 0], [0, -2 * np.cos(y) * np.sin(y)]])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_passing_arguments(self, mocker, tol):
+        """Test that a gradient transform correctly
+        passes arguments"""
+        dev = qml.device("default.qubit", wires=2)
+        spy = mocker.spy(qml.gradients.parameter_shift, "expval_param_shift")
+
+        @qml.qnode(dev)
+        def circuit(weights):
+            qml.RX(weights[0], wires=[0])
+            qml.RY(weights[1], wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliX(1))
+
+        grad_fn = qml.gradients.param_shift(circuit, shift=np.pi / 4)
+
+        w = np.array([0.543, -0.654], requires_grad=True)
+        res = grad_fn(w)
+
+        x, y = w
+        expected = np.array([[-np.sin(x), 0], [0, -2 * np.cos(y) * np.sin(y)]])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        assert spy.call_args[0][2] == np.pi / 4
+
+    def test_expansion(self, mocker, tol):
+        """Test that a gradient transform correctly
+        expands gates with no gradient recipe"""
+        dev = qml.device("default.qubit", wires=2)
+        spy = mocker.spy(qml.gradients.parameter_shift, "expval_param_shift")
+
+        class NonDiffRXGate(qml.PhaseShift):
+            grad_method = "F"
+
+            @staticmethod
+            def decomposition(x, wires):
+                return [qml.RX(x, wires=wires)]
+
+        @qml.qnode(dev)
+        def circuit(weights):
+            NonDiffRXGate(weights[0], wires=[0])
+            qml.RY(weights[1], wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliX(1))
+
+        grad_fn = qml.gradients.param_shift(circuit)
+
+        w = np.array([0.543, -0.654], requires_grad=True)
+        res = grad_fn(w)
+
+        x, y = w
+        expected = np.array([[-np.sin(x), 0], [0, -2 * np.cos(y) * np.sin(y)]])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert spy.call_args[0][0].operations[0].name == "RX"
+
+    def test_permuted_arguments(self, tol):
+        """Test that a gradient transform acts on QNodes
+        correctly when the QNode arguments are permuted"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(weights):
+            qml.RX(weights[1], wires=[0])
+            qml.RY(weights[0], wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliX(1))
+
+        w = np.array([-0.654, 0.543], requires_grad=True)
+        res = qml.gradients.param_shift(circuit)(w)
+
+        expected = qml.jacobian(circuit)(w)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_classical_processing_arguments(self, mocker, tol):
+        """Test that a gradient transform acts on QNodes
+        correctly when the QNode arguments are classically processed"""
+        dev = qml.device("default.qubit", wires=2)
+        spy = mocker.spy(qml.transforms, "classical_jacobian")
+
+        @qml.qnode(dev)
+        def circuit(weights):
+            qml.RX(weights[0] ** 2, wires=[0])
+            qml.RY(weights[1], wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0))
+
+        w = np.array([0.543, -0.654], requires_grad=True)
+        res = qml.gradients.param_shift(circuit)(w)
+
+        classical_jac = spy.spy_return(w)
+        assert isinstance(classical_jac, np.ndarray)
+        assert np.allclose(classical_jac, np.array([[2 * w[0], 0], [0, 1]]))
+
+        x, y = w
+        expected = [-2 * x * np.sin(x ** 2), 0]
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_classical_processing_multiple_arguments(self, mocker, tol):
+        """Test that a gradient transform acts on QNodes
+        correctly when multiple QNode arguments are classically processed"""
+        dev = qml.device("default.qubit", wires=2)
+        spy = mocker.spy(qml.transforms, "classical_jacobian")
+
+        @qml.qnode(dev)
+        def circuit(data, weights):
+            qml.RY(np.cos(data), wires=0)
+            qml.RX(weights[0] ** 2, wires=[0])
+            qml.RY(weights[1], wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0))
+
+        # set d as non-differentiable
+        d = np.array(0.56, requires_grad=False)
+        w = np.array([0.543, -0.654], requires_grad=True)
+        x, y = w
+
+        res = qml.gradients.param_shift(circuit)(d, w)
+        classical_jac = spy.spy_return(d, w)
+        assert np.allclose(classical_jac, np.array([[0, 2 * w[0], 0], [0, 0, 1]]).T)
+
+        expected = np.array([-2 * x * np.cos(np.cos(d)) * np.sin(x ** 2), 0])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        # set d as differentiable
+        d = np.array(0.56, requires_grad=True)
+        w = np.array([0.543, -0.654], requires_grad=True)
+
+        res = qml.gradients.param_shift(circuit)(d, w)
+        classical_jac = spy.spy_return(d, w)
+        assert isinstance(classical_jac, tuple)
+        assert np.allclose(classical_jac[0], [-np.sin(d), 0, 0])
+        assert np.allclose(classical_jac[1], np.array([[0, 2 * w[0], 0], [0, 0, 1]]).T)
+
+        expected_dd = np.cos(x ** 2) * np.sin(d) * np.sin(np.cos(d))
+        expected_dw = np.array([-2 * x * np.cos(np.cos(d)) * np.sin(x ** 2), 0])
+        assert np.allclose(res[0], expected_dd, atol=tol, rtol=0)
+        assert np.allclose(res[1], expected_dw, atol=tol, rtol=0)
+
+    def test_advanced_classical_processing_arguments(self, tol):
+        """Test that a gradient transform acts on QNodes
+        correctly when the QNode arguments are classically processed,
+        and the input weights and the output weights have weird shape."""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(weights):
+            qml.RX(weights[0, 0] ** 2, wires=[0])
+            qml.RY(weights[0, 1], wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        w = np.array([[0.543, -0.654], [0.0, 0.0]], requires_grad=True)
+        res = qml.gradients.param_shift(circuit)(w)
+        assert res.shape == (4, 2, 2)
+
+        expected = qml.jacobian(circuit)(w)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        # when executed with hybrid=False, only the quantum jacobian is returned
+        res = qml.gradients.param_shift(circuit, hybrid=False)(w)
+        assert res.shape == (4, 2)
+
+        @qml.qnode(dev)
+        def circuit(weights):
+            qml.RX(weights[0], wires=[0])
+            qml.RY(weights[1], wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        w = np.array([0.543 ** 2, -0.654], requires_grad=True)
+        expected = qml.jacobian(circuit)(w)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+
+class TestInterfaceIntegration:
+    """Test that the gradient transforms are differentiable
+    using each interface"""
+
+    def test_autograd(self, tol):
+        """Test that a gradient transform remains differentiable
+        with autograd"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.gradients.param_shift
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RY(x ** 2, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.var(qml.PauliX(1))
+
+        x = np.array(-0.654, requires_grad=True)
+
+        res = circuit(x)
+        expected = -4 * x * np.cos(x ** 2) * np.sin(x ** 2)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        res = qml.grad(circuit)(x)
+        expected = -2 * (4 * x ** 2 * np.cos(2 * x ** 2) + np.sin(2 * x ** 2))
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_tf(self, tol):
+        """Test that a gradient transform remains differentiable
+        with TF"""
+        tf = pytest.importorskip("tensorflow")
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.gradients.param_shift
+        @qml.qnode(dev, interface="tf", diff_method="parameter-shift")
+        def circuit(x):
+            qml.RY(x ** 2, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.var(qml.PauliX(1))
+
+        x_ = -0.654
+        x = tf.Variable(x_, dtype=tf.float64)
+
+        with tf.GradientTape() as tape:
+            res = circuit(x)
+
+        expected = -4 * x_ * np.cos(x_ ** 2) * np.sin(x_ ** 2)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        res = tape.gradient(res, x)
+        expected = -2 * (4 * x_ ** 2 * np.cos(2 * x_ ** 2) + np.sin(2 * x_ ** 2))
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_torch(self, tol):
+        """Test that a gradient transform remains differentiable
+        with PyTorch"""
+        torch = pytest.importorskip("torch")
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.gradients.param_shift
+        @qml.qnode(dev, interface="torch")
+        def circuit(x):
+            qml.RY(x, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.var(qml.PauliX(1))
+
+        x_ = -0.654
+        x = torch.tensor(x_, dtype=torch.float64, requires_grad=True)
+        res = circuit(x)[0]
+
+        expected = -2 * np.cos(x_) * np.sin(x_)
+        assert np.allclose(res.detach(), expected, atol=tol, rtol=0)
+
+        res.backward()
+        expected = -2 * np.cos(2 * x_)
+        assert np.allclose(x.grad.detach(), expected, atol=tol, rtol=0)
diff --git a/tests/gradients/test_parameter_shift.py b/tests/gradients/test_parameter_shift.py
index ad2cc2bd0..753710095 100644
--- a/tests/gradients/test_parameter_shift.py
+++ b/tests/gradients/test_parameter_shift.py
@@ -276,7 +276,7 @@ class TestParameterShiftRule:
         dev = qml.device("default.qubit", wires=1)
 
         with qml.tape.JacobianTape() as tape:
-            qml.QubitStateVector(np.array([1.0, -1.0]) / np.sqrt(2), wires=0)
+            qml.QubitStateVector(np.array([1.0, -1.0], requires_grad=False) / np.sqrt(2), wires=0)
             G(theta, wires=[0])
             qml.expval(qml.PauliZ(0))
 
@@ -308,7 +308,7 @@ class TestParameterShiftRule:
         params = np.array([theta, theta ** 3, np.sqrt(2) * theta])
 
         with qml.tape.JacobianTape() as tape:
-            qml.QubitStateVector(np.array([1.0, -1.0]) / np.sqrt(2), wires=0)
+            qml.QubitStateVector(np.array([1.0, -1.0], requires_grad=False) / np.sqrt(2), wires=0)
             qml.Rot(*params, wires=[0])
             qml.expval(qml.PauliZ(0))
 
@@ -344,7 +344,7 @@ class TestParameterShiftRule:
         b = 0.123
 
         with qml.tape.JacobianTape() as tape:
-            qml.QubitStateVector(np.array([1.0, -1.0]) / np.sqrt(2), wires=0)
+            qml.QubitStateVector(np.array([1.0, -1.0], requires_grad=False) / np.sqrt(2), wires=0)
             G(b, wires=[0, 1])
             qml.expval(qml.PauliX(0))
 
@@ -371,7 +371,7 @@ class TestParameterShiftRule:
         a, b, c = np.array([theta, theta ** 3, np.sqrt(2) * theta])
 
         with qml.tape.JacobianTape() as tape:
-            qml.QubitStateVector(np.array([1.0, -1.0]) / np.sqrt(2), wires=0)
+            qml.QubitStateVector(np.array([1.0, -1.0], requires_grad=False) / np.sqrt(2), wires=0)
             qml.CRot(a, b, c, wires=[0, 1])
             qml.expval(qml.PauliX(0))
 
diff --git a/tests/gradients/test_parameter_shift_cv.py b/tests/gradients/test_parameter_shift_cv.py
index c43548a6b..c64cf5a6e 100644
--- a/tests/gradients/test_parameter_shift_cv.py
+++ b/tests/gradients/test_parameter_shift_cv.py
@@ -705,7 +705,8 @@ class TestExpectationQuantumGradients:
             [
                 [0.51310276 + 0.81702166j, 0.13649626 + 0.22487759j],
                 [0.26300233 + 0.00556194j, -0.96414101 - 0.03508489j],
-            ]
+            ],
+            requires_grad=False,
         )
 
         with qml.tape.JacobianTape() as tape: