Repository: PennyLaneAI/pennylane
Pull request #1618 (2021-09-01 16:58:26)

Final touches to get `batch.execute` ready for QNode integration
**Context:** I have begun the process of integrating the new batch execution pipeline into the QNode. On my first attempt, however, I noticed there some changes to the existing batch execution function is required in order to support all features of the QNode, including:

- Raising useful error messages if the interface is incorrect or not available
- Supporting setting of shots dynamically when calling the QNode.

I also spotted a 'todo' still in the source code.

Rather than include this as part of a large, separate PR, I have separated the low-level batch execution changes into a separate PR for ease of review.

**Description of the Change:**

- The execute function, since being proposed, now supports no-interface/no-gradient, as well as backpropagation. It has been made a top-level import to allow it to be easily used for backpropagation without causing dependency issues from requiring it be imported from the `interfaces` package.

- A private context manager, `set_shots`, has been created, to _safely_ modify shots at execution time.

- Better error handling has been added for import errors.

- A small bugfix was made in the TF interface.

- A redundant validation test was removed.

- A temporary 'hack', where the _location_ of the `gradient_fn` was used to check if the `gradient_fn` was a gradient transform, has been removed, and replaced with a simple `isinstance(gradient_fn, qml.gradients.gradient_transform)`.

**Benefits:**

- The main benefit is that `execute()` can now be called with a custom `override_shots` keyword argument. This will temporarily override the device shot number during execution. Note that validation is already present in the device to ensure that setting shots to an invalid value for a device will raise errors.

- The `set_shots` context manager is significantly safer than the approach currently available in the QNode, which sets `device.shots` without a try-except statement.

- `@contextlib` provides tools for creating context managers without the boilerplate, which is nice. Another benefit is that context managers created this way can also be used as decorators/transforms --- the transformed function will always be executed within the context manager. This is an _extremely_ important ability; coupled with function closure, this ensures that *side effects* of the `execute()` function will continue to use the temporary shots value, even once the execution function has returned.

  In fact, this fixes a bug in TF/Torch currently in the master branch, where dynamically setting the shots does not work with gradient computations. This bug has prevented us from advertising `qnode(weights, shots=100)` as a main user-interface so far.

- Checking if the gradient function is a gradient transform is now cleaner and much less error prone.

**Possible Drawbacks:** n/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/1618?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1618](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1618?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (a3552e7) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/90baa1d0cc509ce2e0936ff0ccd2e8303c35f457?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (90baa1d) will **decrease** coverage by `0.02%`.
> The diff coverage is `91.42%`.

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

```diff
@@            Coverage Diff             @@
##           master    #1618      +/-   ##
==========================================
- Coverage   99.13%   99.11%   -0.03%     
==========================================
  Files         196      196              
  Lines       14162    14166       +4     
==========================================
+ Hits        14040    14041       +1     
- Misses        122      125       +3     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1618?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/gradients/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1618/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% <ø> (ø)` | |
| [pennylane/gradients/finite\_difference.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1618/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% <ø> (ø)` | |
| [pennylane/gradients/parameter\_shift.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1618/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% <ø> (ø)` | |
| [pennylane/gradients/parameter\_shift\_cv.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1618/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% <ø> (ø)` | |
| [pennylane/interfaces/batch/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1618/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvX19pbml0X18ucHk=) | `96.59% <88.88%> (-3.41%)` | :arrow_down: |
| [pennylane/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1618/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/batch/autograd.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1618/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvYXV0b2dyYWQucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/batch/tensorflow.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1618/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvdGVuc29yZmxvdy5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/batch/torch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1618/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvdG9yY2gucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/transforms/batch\_transform.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1618/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/1618?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/1618?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [90baa1d...a3552e7](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1618?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: Nice touches @josh146! Good to go :rocket:  
josh146: [ch8734]
josh146: [ch8735]

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 7c7c17cd9..e40de97da 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -272,6 +272,7 @@
   [(#1542)](https://github.com/PennyLaneAI/pennylane/pull/1542)
   [(#1549)](https://github.com/PennyLaneAI/pennylane/pull/1549)
   [(#1608)](https://github.com/PennyLaneAI/pennylane/pull/1608)
+  [(#1618)](https://github.com/PennyLaneAI/pennylane/pull/1618)
 
   For example:
 
diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index a61d3b38e..6c806cded 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -67,6 +67,7 @@ from pennylane.vqe import ExpvalCost, VQECost
 from .collections import QNodeCollection, dot, map, sum
 import pennylane.grouping  # pylint:disable=wrong-import-order
 import pennylane.gradients  # pylint:disable=wrong-import-order
+from pennylane.interfaces.batch import execute  # pylint:disable=wrong-import-order
 
 # Look for an existing configuration file
 default_config = Configuration("config.toml")
diff --git a/pennylane/gradients/__init__.py b/pennylane/gradients/__init__.py
index 0feec6572..14d8e9a74 100644
--- a/pennylane/gradients/__init__.py
+++ b/pennylane/gradients/__init__.py
@@ -247,8 +247,7 @@ The output tapes can then be evaluated and post-processed to retrieve
 the gradient:
 
 >>> dev = qml.device("default.qubit", wires=2)
->>> from pennylane.interfaces.batch import execute
->>> fn(execute(gradient_tapes, dev, None))
+>>> fn(qml.execute(gradient_tapes, dev, None))
 [[-0.09347337 -0.18884787 -0.28818254]]
 
 Note that the post-processing function ``fn`` returned by the
diff --git a/pennylane/gradients/finite_difference.py b/pennylane/gradients/finite_difference.py
index fffb8438c..3a0b5b2b9 100644
--- a/pennylane/gradients/finite_difference.py
+++ b/pennylane/gradients/finite_difference.py
@@ -274,8 +274,7 @@ def finite_diff(tape, argnum=None, h=1e-7, approx_order=1, n=1, strategy="forwar
         the gradient:
 
         >>> dev = qml.device("default.qubit", wires=2)
-        >>> from pennylane.interfaces.batch import execute
-        >>> fn(execute(gradient_tapes, dev, None))
+        >>> fn(qml.execute(gradient_tapes, dev, None))
         [[-0.38751721 -0.18884787 -0.38355704]
          [ 0.69916862  0.34072424  0.69202359]]
     """
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
index 61dce8ef7..d76554206 100644
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -517,8 +517,7 @@ def param_shift(
         the gradient:
 
         >>> dev = qml.device("default.qubit", wires=2)
-        >>> from pennylane.interfaces.batch import execute
-        >>> fn(execute(gradient_tapes, dev, None))
+        >>> fn(qml.execute(gradient_tapes, dev, None))
         [[-0.38751721 -0.18884787 -0.38355704]
          [ 0.69916862  0.34072424  0.69202359]]
     """
diff --git a/pennylane/gradients/parameter_shift_cv.py b/pennylane/gradients/parameter_shift_cv.py
index 5d65f29f6..ec38d11a3 100644
--- a/pennylane/gradients/parameter_shift_cv.py
+++ b/pennylane/gradients/parameter_shift_cv.py
@@ -630,8 +630,7 @@ def param_shift_cv(
         the gradient:
 
         >>> dev = qml.device("default.gaussian", wires=2)
-        >>> from pennylane.interfaces.batch import execute
-        >>> fn(execute(gradient_tapes, dev, None))
+        >>> fn(qml.execute(gradient_tapes, dev, None))
         array([[-0.32487113, -0.4054074 , -0.87049853,  0.4054074 ]])
     """
 
diff --git a/pennylane/interfaces/batch/__init__.py b/pennylane/interfaces/batch/__init__.py
index f468a1155..6f4f1adb6 100644
--- a/pennylane/interfaces/batch/__init__.py
+++ b/pennylane/interfaces/batch/__init__.py
@@ -16,7 +16,9 @@ This subpackage defines functions for interfacing devices' batch execution
 capabilities with different machine learning libraries.
 """
 # pylint: disable=import-outside-toplevel,too-many-arguments,too-many-branches
+import contextlib
 from functools import wraps
+import itertools
 
 from cachetools import LRUCache
 import numpy as np
@@ -24,6 +26,49 @@ import numpy as np
 import pennylane as qml
 
 
+INTERFACE_NAMES = {
+    "NumPy": (None,),
+    "Autograd": ("autograd", "numpy"),  # for backwards compatibility
+    "PyTorch": ("torch", "pytorch"),
+    "TensorFlow": ("tf", "tensorflow"),
+}
+"""dict[str, str]: maps allowed interface strings to the name of the interface"""
+
+SUPPORTED_INTERFACES = list(itertools.chain(*INTERFACE_NAMES.values()))
+
+
+@contextlib.contextmanager
+def set_shots(device, shots):
+    """Context manager to temporarily change the shots
+    of a device.
+
+    This context manager can be used in two ways.
+
+    As a standard context manager:
+
+    >>> dev = qml.device("default.qubit", wires=2, shots=None)
+    >>> with set_shots(dev, shots=100):
+    ...     print(dev.shots)
+    100
+    >>> print(dev.shots)
+    None
+
+    Or as a decorator that acts on a function that uses the device:
+
+    >>> set_shots(dev, shots=100)(lambda: dev.shots)()
+    100
+    """
+    original_shots = device.shots
+
+    try:
+        if shots is not False and device.shots != shots:
+            device.shots = shots
+        yield
+    finally:
+        if device.shots != original_shots:
+            device.shots = original_shots
+
+
 def cache_execute(fn, cache, pass_kwargs=False, return_tuple=True):
     """Decorator that adds caching to a function that executes
     multiple tapes on a device.
@@ -143,6 +188,7 @@ def execute(
     cache=True,
     cachesize=10000,
     max_diff=2,
+    override_shots=False,
 ):
     """Execute a batch of tapes on a device in an autodifferentiable-compatible manner.
 
@@ -236,17 +282,19 @@ def execute(
         # cache=True: create a LRUCache object
         cache = LRUCache(maxsize=cachesize, getsizeof=lambda x: qml.math.shape(x)[0])
 
+    batch_execute = set_shots(device, override_shots)(device.batch_execute)
+
     if gradient_fn is None:
         with qml.tape.Unwrap(*tapes):
-            res = cache_execute(device.batch_execute, cache, return_tuple=False)(tapes)
+            res = cache_execute(batch_execute, cache, return_tuple=False)(tapes)
 
         return res
 
-    if gradient_fn == "backprop":
-        return cache_execute(device.batch_execute, cache, return_tuple=False)(tapes)
+    if gradient_fn == "backprop" or interface is None:
+        return cache_execute(batch_execute, cache, return_tuple=False)(tapes)
 
-    # the default execution function is device.batch_execute
-    execute_fn = cache_execute(device.batch_execute, cache)
+    # the default execution function is batch_execute
+    execute_fn = cache_execute(batch_execute, cache)
 
     if gradient_fn == "device":
         # gradient function is a device method
@@ -254,16 +302,19 @@ def execute(
         if mode in ("forward", "best"):
             # replace the forward execution function to return
             # both results and gradients
-            execute_fn = device.execute_and_gradients
+            execute_fn = set_shots(device, override_shots)(device.execute_and_gradients)
             gradient_fn = None
 
         elif mode == "backward":
             # disable caching on the forward pass
-            execute_fn = cache_execute(device.batch_execute, cache=None)
+            execute_fn = cache_execute(batch_execute, cache=None)
 
             # replace the backward gradient computation
             gradient_fn = cache_execute(
-                device.gradients, cache, pass_kwargs=True, return_tuple=False
+                set_shots(device, override_shots)(device.gradients),
+                cache,
+                pass_kwargs=True,
+                return_tuple=False,
             )
 
     elif mode == "forward":
@@ -272,14 +323,26 @@ def execute(
         # in this case would have ambiguous behaviour.
         raise ValueError("Gradient transforms cannot be used with mode='forward'")
 
-    if interface == "autograd":
-        from .autograd import execute as _execute
-    elif interface in ("tf", "tensorflow"):
-        from .tensorflow import execute as _execute
-    elif interface in ("torch", "pytorch"):
-        from .torch import execute as _execute
-    else:
-        raise ValueError(f"Unknown interface {interface}")
+    try:
+        if interface in INTERFACE_NAMES["Autograd"]:
+            from .autograd import execute as _execute
+        elif interface in INTERFACE_NAMES["TensorFlow"]:
+            from .tensorflow import execute as _execute
+        elif interface in INTERFACE_NAMES["PyTorch"]:
+            from .torch 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]
+
+        raise qml.QuantumFunctionError(
+            f"{interface_name} not found. Please install the latest "
+            f"version of {interface_name} to enable the '{interface}' interface."
+        ) from e
 
     res = _execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_diff=max_diff)
 
diff --git a/pennylane/interfaces/batch/autograd.py b/pennylane/interfaces/batch/autograd.py
index 98e05dc2f..9c3060a8c 100644
--- a/pennylane/interfaces/batch/autograd.py
+++ b/pennylane/interfaces/batch/autograd.py
@@ -16,8 +16,6 @@ This module contains functions for adding the Autograd interface
 to a PennyLane Device class.
 """
 # pylint: disable=too-many-arguments
-import inspect
-
 import autograd
 from autograd.numpy.numpy_boxes import ArrayBox
 
@@ -171,13 +169,8 @@ def vjp(
         else:
             # Need to compute the Jacobians on the backward pass (accumulation="backward")
 
-            # Temporary: check if the gradient function is a differentiable transform.
-            # For the moment, simply check if it is part of the `qml.gradients` package.
-            # Longer term, we should have a way of checking this directly
-            # (e.g., isinstance(gradient_fn, GradientTransform))
-            module_name = getattr(inspect.getmodule(gradient_fn), "__name__", "")
-
-            if "pennylane.gradients" in module_name:
+            if isinstance(gradient_fn, qml.gradients.gradient_transform):
+                # Gradient function is a gradient transform.
 
                 # Generate and execute the required gradient tapes
                 if _n == max_diff:
@@ -212,27 +205,20 @@ def vjp(
                         )
                     )
 
-            elif (
-                hasattr(gradient_fn, "fn")
-                and inspect.ismethod(gradient_fn.fn)
-                and gradient_fn.fn.__self__ is device
-            ):
-                # Gradient function is a device method.
+            else:
+                # Gradient function is not a gradient transform
+                # (e.g., it might be a device method).
                 # Note that unlike the previous branch:
                 #
                 # - there is no recursion here
                 # - gradient_fn is not differentiable
                 #
                 # so we cannot support higher-order derivatives.
-
                 with qml.tape.Unwrap(*tapes):
                     jacs = gradient_fn(tapes, **gradient_kwargs)
 
                 vjps = [qml.gradients.compute_vjp(d, jac) for d, jac in zip(dy, jacs)]
 
-            else:
-                raise ValueError("Unknown gradient function.")
-
         return [qml.math.to_numpy(v, max_depth=_n) if isinstance(v, ArrayBox) else v for v in vjps]
 
     return grad_fn
diff --git a/pennylane/interfaces/batch/tensorflow.py b/pennylane/interfaces/batch/tensorflow.py
index 8b1ad9943..3678a5cd1 100644
--- a/pennylane/interfaces/batch/tensorflow.py
+++ b/pennylane/interfaces/batch/tensorflow.py
@@ -16,8 +16,6 @@ This module contains functions for adding the TensorFlow interface
 to a PennyLane Device class.
 """
 # pylint: disable=too-many-arguments,too-many-branches
-import inspect
-
 import numpy as np
 import tensorflow as tf
 from tensorflow.python.eager import context
@@ -102,13 +100,8 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
             else:
                 # Need to compute the Jacobians on the backward pass (accumulation="backward")
 
-                # Temporary: check if the gradient function is a differentiable transform.
-                # For the moment, simply check if it is part of the `qml.gradients` package.
-                # Longer term, we should have a way of checking this directly
-                # (e.g., isinstance(gradient_fn, GradientTransform))
-                module_name = getattr(inspect.getmodule(gradient_fn), "__name__", "")
-
-                if "pennylane.gradients" in module_name:
+                if isinstance(gradient_fn, qml.gradients.gradient_transform):
+                    # Gradient function is a gradient transform.
 
                     # Generate and execute the required gradient tapes
                     if _n == max_diff or not context.executing_eagerly():
@@ -148,12 +141,9 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
                             )
                         )
 
-                elif (
-                    hasattr(gradient_fn, "fn")
-                    and inspect.ismethod(gradient_fn.fn)
-                    and gradient_fn.fn.__self__ is device
-                ):
-                    # Gradient function is a device method.
+                else:
+                    # Gradient function is not a gradient transform
+                    # (e.g., it might be a device method).
                     # Note that unlike the previous branch:
                     #
                     # - there is no recursion here
@@ -163,9 +153,6 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
                     with qml.tape.Unwrap(*tapes, params=params_unwrapped, set_trainable=False):
                         vjps = _compute_vjp(dy, gradient_fn(tapes, **gradient_kwargs))
 
-                else:
-                    raise ValueError("Unknown gradient function.")
-
             variables = tfkwargs.get("variables", None)
             return (vjps, variables) if variables is not None else vjps
 
diff --git a/pennylane/interfaces/batch/torch.py b/pennylane/interfaces/batch/torch.py
index 6786569e7..151cdd591 100644
--- a/pennylane/interfaces/batch/torch.py
+++ b/pennylane/interfaces/batch/torch.py
@@ -16,8 +16,6 @@ This module contains functions for adding the PyTorch interface
 to a PennyLane Device class.
 """
 # pylint: disable=too-many-arguments,protected-access
-import inspect
-
 import numpy as np
 import torch
 
@@ -119,13 +117,8 @@ class ExecuteTapes(torch.autograd.Function):
         else:
             # Need to compute the Jacobians on the backward pass (accumulation="backward")
 
-            # Temporary: check if the gradient function is a differentiable transform.
-            # For the moment, simply check if it is part of the `qml.gradients` package.
-            # Longer term, we should have a way of checking this directly
-            # (e.g., isinstance(gradient_fn, GradientTransform))
-            module_name = getattr(inspect.getmodule(ctx.gradient_fn), "__name__", "")
-
-            if "pennylane.gradients" in module_name:
+            if isinstance(ctx.gradient_fn, qml.gradients.gradient_transform):
+                # Gradient function is a gradient transform.
 
                 # Generate and execute the required gradient tapes
                 if ctx._n < ctx.max_diff:
@@ -172,12 +165,9 @@ class ExecuteTapes(torch.autograd.Function):
 
                     vjps = processing_fn(ctx.execute_fn(vjp_tapes)[0])
 
-            elif (
-                hasattr(ctx.gradient_fn, "fn")
-                and inspect.ismethod(ctx.gradient_fn.fn)
-                and ctx.gradient_fn.fn.__self__ is ctx.device
-            ):
-                # Gradient function is a device method.
+            else:
+                # Gradient function is not a gradient transform
+                # (e.g., it might be a device method).
                 # Note that unlike the previous branch:
                 #
                 # - there is no recursion here
@@ -190,9 +180,6 @@ class ExecuteTapes(torch.autograd.Function):
 
                 vjps = _compute_vjp(dy, jacs, device=ctx.torch_device)
 
-            else:
-                raise ValueError("Unknown gradient function.")
-
         # The output of backward must match the input of forward.
         # Therefore, we return `None` for the gradient of `kwargs`.
         return (None,) + tuple(vjps)
diff --git a/pennylane/transforms/batch_transform.py b/pennylane/transforms/batch_transform.py
index d5261a1ff..ba4a47309 100644
--- a/pennylane/transforms/batch_transform.py
+++ b/pennylane/transforms/batch_transform.py
@@ -17,8 +17,6 @@ import functools
 
 import pennylane as qml
 
-from pennylane.interfaces.batch import execute
-
 
 class batch_transform:
     r"""Class for registering a tape transform that takes a tape, and outputs
@@ -110,9 +108,8 @@ class batch_transform:
 
     We can execute these tapes manually:
 
-    >>> from pennylane.interfaces.batch import execute
     >>> dev = qml.device("default.qubit", wires=1)
-    >>> res = execute(tapes, dev, interface="autograd", gradient_fn=qml.gradients.param_shift)
+    >>> res = qml.execute(tapes, dev, interface="autograd", gradient_fn=qml.gradients.param_shift)
     >>> print(res)
     [tensor([0.94765073], requires_grad=True), tensor([0.31532236], requires_grad=True)]
 
@@ -191,7 +188,7 @@ class batch_transform:
             elif gradient_fn == "finite-diff":
                 gradient_fn = qml.gradients.finite_diff
 
-            res = execute(
+            res = qml.execute(
                 tapes,
                 device=qnode.device,
                 gradient_fn=gradient_fn,


Test patch:
diff --git a/tests/interfaces/test_batch_autograd.py b/tests/interfaces/test_batch_autograd.py
index 0dfa04f85..50f7cef69 100644
--- a/tests/interfaces/test_batch_autograd.py
+++ b/tests/interfaces/test_batch_autograd.py
@@ -13,7 +13,10 @@
 # limitations under the License.
 """Unit tests for the autograd interface"""
 import functools
+import importlib
+import sys
 
+import autograd
 import pytest
 from pennylane import numpy as np
 
@@ -25,6 +28,29 @@ from pennylane.interfaces.batch import execute
 class TestAutogradExecuteUnitTests:
     """Unit tests for autograd execution"""
 
+    def test_import_error(self, mocker):
+        """Test that an exception is caught on import error"""
+
+        mock = mocker.patch.object(autograd.extend, "defvjp")
+        mock.side_effect = ImportError()
+
+        try:
+            del sys.modules["pennylane.interfaces.batch.autograd"]
+        except:
+            pass
+
+        dev = qml.device("default.qubit", wires=2, shots=None)
+
+        with qml.tape.JacobianTape() as tape:
+            qml.expval(qml.PauliY(1))
+
+        with pytest.raises(
+            qml.QuantumFunctionError,
+            match="Autograd not found. Please install the latest version "
+            "of Autograd to enable the 'autograd' interface",
+        ):
+            qml.execute([tape], dev, gradient_fn=param_shift, interface="autograd")
+
     def test_jacobian_options(self, mocker, tol):
         """Test setting jacobian options"""
         spy = mocker.spy(qml.gradients, "param_shift")
@@ -48,26 +74,8 @@ class TestAutogradExecuteUnitTests:
         for args in spy.call_args_list:
             assert args[1]["shift"] == np.pi / 4
 
-    def test_unknown_gradient_fn_error(self):
-        """Test that an error is raised if an unknown gradient function
-        is passed"""
-        a = np.array([0.1, 0.2], requires_grad=True)
-
-        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)[0]
-
-        with pytest.raises(ValueError, match="Unknown gradient function"):
-            res = qml.jacobian(cost)(a, device=dev)
-
     def test_incorrect_mode(self):
-        """Test that an error is raised if an gradient transform
+        """Test that an error is raised if a gradient transform
         is used with mode=forward"""
         a = np.array([0.1, 0.2], requires_grad=True)
 
@@ -98,7 +106,7 @@ class TestAutogradExecuteUnitTests:
                 qml.RX(a[1], wires=0)
                 qml.expval(qml.PauliZ(0))
 
-            return execute([tape], device, gradient_fn=param_shift, interface=None)[0]
+            return execute([tape], device, gradient_fn=param_shift, interface="None")[0]
 
         with pytest.raises(ValueError, match="Unknown interface"):
             cost(a, device=dev)
@@ -810,6 +818,90 @@ class TestHigherOrderDerivatives:
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
 
+class TestOverridingShots:
+    """Test overriding shots on execution"""
+
+    def test_changing_shots(self, mocker, tol):
+        """Test that changing shots works on execution"""
+        dev = qml.device("default.qubit", wires=2, shots=None)
+        a, b = np.array([0.543, -0.654], requires_grad=True)
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RY(a, wires=0)
+            qml.RX(b, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliY(1))
+
+        spy = mocker.spy(dev, "sample")
+
+        # execute with device default shots (None)
+        res = execute([tape], dev, gradient_fn=param_shift)
+        assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0)
+        spy.assert_not_called()
+
+        # execute with shots=100
+        res = execute([tape], dev, gradient_fn=param_shift, override_shots=100)
+        spy.assert_called()
+        assert spy.spy_return.shape == (100,)
+
+        # device state has been unaffected
+        assert dev.shots is None
+        spy = mocker.spy(dev, "sample")
+        res = execute([tape], dev, gradient_fn=param_shift)
+        assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0)
+        spy.assert_not_called()
+
+    def test_overriding_shots_with_same_value(self, mocker):
+        """Overriding shots with the same value as the device will have no effect"""
+        dev = qml.device("default.qubit", wires=2, shots=123)
+        a, b = np.array([0.543, -0.654], requires_grad=True)
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RY(a, wires=0)
+            qml.RX(b, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliY(1))
+
+        spy = mocker.Mock(wraps=qml.Device.shots.fset)
+        mock_property = qml.Device.shots.setter(spy)
+        mocker.patch.object(qml.Device, "shots", mock_property)
+
+        res = execute([tape], dev, gradient_fn=param_shift, override_shots=123)
+        # overriden shots is the same, no change
+        spy.assert_not_called()
+
+        res = execute([tape], dev, gradient_fn=param_shift, override_shots=100)
+        # overriden shots is not the same, shots were changed
+        spy.assert_called()
+
+        # shots were temporarily set to the overriden value
+        assert spy.call_args_list[0][0] == (dev, 100)
+        # shots were then returned to the built-in value
+        assert spy.call_args_list[1][0] == (dev, 123)
+
+    def test_gradient_integration(self, tol):
+        """Test that temporarily setting the shots works
+        for gradient computations"""
+        dev = qml.device("default.qubit", wires=2, shots=None)
+        a, b = np.array([0.543, -0.654], requires_grad=True)
+
+        def cost_fn(a, b, shots):
+            with qml.tape.JacobianTape() as tape:
+                qml.RY(a, wires=0)
+                qml.RX(b, wires=1)
+                qml.CNOT(wires=[0, 1])
+                qml.expval(qml.PauliY(1))
+
+            return execute([tape], dev, gradient_fn=param_shift, override_shots=shots)[0]
+
+        res = qml.jacobian(cost_fn)(a, b, shots=[10000, 10000, 10000])
+        assert dev.shots is None
+        assert len(res) == 3
+
+        expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]
+        assert np.allclose(np.mean(res, axis=0), expected, atol=0.1, rtol=0)
+
+
 execute_kwargs = [
     {"gradient_fn": param_shift},
     {"gradient_fn": finite_diff},
diff --git a/tests/interfaces/test_batch_tensorflow.py b/tests/interfaces/test_batch_tensorflow.py
index c8744b3f4..f85ab2a59 100644
--- a/tests/interfaces/test_batch_tensorflow.py
+++ b/tests/interfaces/test_batch_tensorflow.py
@@ -55,24 +55,6 @@ class TestTensorFlowExecuteUnitTests:
         for args in spy.call_args_list:
             assert args[1]["shift"] == np.pi / 4
 
-    def test_unknown_gradient_fn_error(self):
-        """Test that an error is raised if an unknown gradient function
-        is passed"""
-        a = tf.Variable([0.1, 0.2])
-
-        dev = qml.device("default.qubit", wires=1)
-
-        with tf.GradientTape() as t:
-            with qml.tape.JacobianTape() as tape:
-                qml.RY(a[0], wires=0)
-                qml.RX(a[1], wires=0)
-                qml.expval(qml.PauliZ(0))
-
-            res = execute([tape], dev, gradient_fn=lambda x: x, interface="tf")[0]
-
-        with pytest.raises(ValueError, match="Unknown gradient function"):
-            print(t.jacobian(res, a))
-
     def test_incorrect_mode(self):
         """Test that an error is raised if a gradient transform
         is used with mode=forward"""
diff --git a/tests/interfaces/test_batch_torch.py b/tests/interfaces/test_batch_torch.py
index 35fd68bbd..ba0e0ca41 100644
--- a/tests/interfaces/test_batch_torch.py
+++ b/tests/interfaces/test_batch_torch.py
@@ -53,23 +53,6 @@ class TestTorchExecuteUnitTests:
         for args in spy.call_args_list:
             assert args[1]["shift"] == np.pi / 4
 
-    def test_unknown_gradient_fn_error(self):
-        """Test that an error is raised if an unknown gradient function
-        is passed"""
-        a = torch.tensor([0.1, 0.2], requires_grad=True)
-
-        dev = qml.device("default.qubit", wires=1)
-
-        with qml.tape.JacobianTape() as tape:
-            qml.RY(a[0], wires=0)
-            qml.RX(a[1], wires=0)
-            qml.expval(qml.PauliZ(0))
-
-        res = execute([tape], dev, gradient_fn=lambda x: x, interface="torch")[0]
-
-        with pytest.raises(ValueError, match="Unknown gradient function"):
-            res.backward()
-
     def test_incorrect_mode(self):
         """Test that an error is raised if a gradient transform
         is used with mode=forward"""