Repository: PennyLaneAI/pennylane
Pull request #6741 (2024-12-27 16:49:36)

add qml.math.grad and qml.math.jacobian for differentiating any interface
**Context:**

We often have to rewrite tests four different ways to test differentiability for each different interface.  This is rather time-intensive and annoying.  By adding and using a `qml.math.grad`, we can write one test that applies to all interfaces.

**Description of the Change:**

Adds a `qml.math.grad` and `qml.math.jacobian` with a jax-like interface that can differentiate with respect to any argument inputs.

**Benefits:**

**Possible Drawbacks:**

We have several different options for interface, but I decided to do my best to mirror the jax interface.

**Related GitHub Issues:**


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6741?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: Patch coverage is `97.61905%` with `1 line` in your changes missing coverage. Please review.
> Project coverage is 99.60%. Comparing base [(`7d1b0e5`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/7d1b0e5af4d907b2f65b563b5d23f8a137fcd00b?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`45c6e0c`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/45c6e0c3cfcb4b4b22587de36ffa61bd49e2f4e2?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 5 commits behind head on master.

| [Files with missing lines](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6741?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/math/grad.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6741?src=pr&el=tree&filepath=pennylane%2Fmath%2Fgrad.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvZ3JhZC5weQ==) | 96.66% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6741?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #6741      +/-   ##
==========================================
- Coverage   99.60%   99.60%   -0.01%     
==========================================
  Files         475      476       +1     
  Lines       45121    45157      +36     
==========================================
+ Hits        44944    44979      +35     
- Misses        177      178       +1     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6741?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

JerryChen97: Also, the similar error handling branch in `grad.py` seems missing tests coverage.

Patch:
diff --git a/doc/introduction/interfaces/numpy.rst b/doc/introduction/interfaces/numpy.rst
index 1b16a2053..859ca1343 100644
--- a/doc/introduction/interfaces/numpy.rst
+++ b/doc/introduction/interfaces/numpy.rst
@@ -3,8 +3,6 @@
 NumPy interface
 ===============
 
-.. note:: This interface is the default interface supported by PennyLane's :class:`QNode <pennylane.QNode>`.
-
 
 Using the NumPy interface
 -------------------------
@@ -275,10 +273,10 @@ we would get an error message. This is because the `gradient <https://en.wikiped
 only defined for scalar functions, i.e., functions which return a single value. In the case where the QNode
 returns multiple expectation values, the correct differential operator to use is
 the `Jacobian matrix <https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant>`_.
-This can be accessed in PennyLane as :func:`~.jacobian`.
+This can be accessed in PennyLane as :func:`~pennylane.jacobian`.
 
 As the ``circuit5`` returns a tuple of numpy arrays instead of a single numpy array, the results need
-to be stacked into a single array before use with :func:`~.jacobian`.
+to be stacked into a single array before use with :func:`~pennylane.jacobian`.
 
 >>> j1 = qml.jacobian(lambda x: np.stack(circuit5(x)))
 >>> j1(params)
@@ -286,7 +284,7 @@ array([[ 0.        , -0.98006658],
        [-0.98006658,  0.        ]])
 
 
-The output of :func:`~.jacobian` is a two-dimensional vector, with the first/second element being
+The output of :func:`~pennylane.jacobian` is a two-dimensional vector, with the first/second element being
 the partial derivative of the first/second expectation value with respect to the input parameter.
 
 
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index f035d4c55..3680c3dd6 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -340,6 +340,10 @@ such as `shots`, `rng` and `prng_key`.
 
 <h4>Other Improvements</h4>
 
+* `qml.math.grad` and `qml.math.jacobian` added to differentiate a function with inputs of any
+  interface in a jax-like manner.
+  [(#6741)](https://github.com/PennyLaneAI/pennylane/pull/6741)
+
 * `qml.GroverOperator` now has a `work_wires` property.
   [(#6738)](https://github.com/PennyLaneAI/pennylane/pull/6738)
 
diff --git a/pennylane/_grad.py b/pennylane/_grad.py
index 7e75e65b9..868c03fce 100644
--- a/pennylane/_grad.py
+++ b/pennylane/_grad.py
@@ -251,6 +251,21 @@ class grad:
         return grad_value, ans
 
 
+def _error_if_not_array(f):
+    """A function decorator that raises an error if the function output is not an autograd, pennylane, or numpy array."""
+
+    @wraps(f)
+    def new_f(*args, **kwargs):
+        output = f(*args, **kwargs)
+        if output.__class__.__module__.split(".")[0] not in {"autograd", "pennylane", "numpy"}:
+            raise ValueError(
+                f"autograd can only differentiate with respect to arrays, not {type(output)}. Ensure the output class is an autograd array."
+            )
+        return output
+
+    return new_f
+
+
 def jacobian(func, argnum=None, method=None, h=None):
     """Returns the Jacobian as a callable function of vector-valued (functions of) QNodes.
     This function is compatible with Autograd and :func:`~.qjit`.
@@ -514,7 +529,7 @@ def jacobian(func, argnum=None, method=None, h=None):
                 "If this is unintended, please add trainable parameters via the "
                 "'requires_grad' attribute or 'argnum' keyword."
             )
-        jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
+        jac = tuple(_jacobian(_error_if_not_array(func), arg)(*args, **kwargs) for arg in _argnum)
 
         return jac[0] if unpack else jac
 
diff --git a/pennylane/gradients/classical_jacobian.py b/pennylane/gradients/classical_jacobian.py
index 0434bf572..20d255217 100644
--- a/pennylane/gradients/classical_jacobian.py
+++ b/pennylane/gradients/classical_jacobian.py
@@ -113,7 +113,7 @@ def classical_jacobian(qnode, argnum=None, expand_fn=None, trainable_only=True):
          - ``tuple(array)``
 
     [1] If there only is one trainable QNode argument, the tuple is unpacked to a
-    single ``array``, as is the case for :func:`.jacobian`.
+    single ``array``, as is the case for :func:`pennylane.jacobian`.
 
     [2] For JAX, ``argnum=None`` defaults to ``argnum=0`` in contrast to all other
     interfaces. This means that only the classical Jacobian with respect to the first
@@ -158,7 +158,7 @@ def classical_jacobian(qnode, argnum=None, expand_fn=None, trainable_only=True):
         if qnode.interface == "autograd":
             jac = qml.jacobian(classical_preprocessing, argnum=wrapper_argnum)(*args, **kwargs)
 
-        if qnode.interface == "torch":
+        elif qnode.interface == "torch":
             import torch
 
             def _jacobian(*args, **kwargs):  # pylint: disable=unused-argument
@@ -177,7 +177,7 @@ def classical_jacobian(qnode, argnum=None, expand_fn=None, trainable_only=True):
 
             jac = _jacobian(*args, **kwargs)
 
-        if qnode.interface in ["jax", "jax-jit"]:
+        elif qnode.interface in ["jax", "jax-jit"]:
             import jax
 
             argnum = 0 if wrapper_argnum is None else wrapper_argnum
@@ -187,7 +187,7 @@ def classical_jacobian(qnode, argnum=None, expand_fn=None, trainable_only=True):
 
             jac = _jacobian(*args, **kwargs)
 
-        if qnode.interface == "tf":
+        elif qnode.interface == "tf":
             import tensorflow as tf
 
             def _jacobian(*args, **kwargs):
@@ -206,6 +206,9 @@ def classical_jacobian(qnode, argnum=None, expand_fn=None, trainable_only=True):
 
             jac = _jacobian(*args, **kwargs)
 
+        else:
+            raise ValueError(f"Undifferentiable interface {qnode.interface}.")
+
         if old_interface == "auto":
             qnode.interface = "auto"
 
diff --git a/pennylane/math/__init__.py b/pennylane/math/__init__.py
index 6107c4009..6d46896b0 100644
--- a/pennylane/math/__init__.py
+++ b/pennylane/math/__init__.py
@@ -101,6 +101,7 @@ from .interface_utils import (
     get_interface,
     Interface,
 )
+from .grad import grad, jacobian
 
 sum = ar.numpy.sum
 toarray = ar.numpy.to_numpy
@@ -168,10 +169,12 @@ __all__ = [
     "get_canonical_interface_name",
     "get_deep_interface",
     "get_trainable_indices",
+    "grad",
     "in_backprop",
     "is_abstract",
     "is_independent",
     "iscomplex",
+    "jacobian",
     "marginal_prob",
     "max_entropy",
     "min_entropy",
diff --git a/pennylane/math/grad.py b/pennylane/math/grad.py
new file mode 100644
index 000000000..3e4cd53ce
--- /dev/null
+++ b/pennylane/math/grad.py
@@ -0,0 +1,246 @@
+# Copyright 2024 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 submodule defines grad and jacobian for differentiating circuits in an interface-independent way.
+"""
+
+from typing import Callable, Sequence, Union
+
+from pennylane._grad import grad as _autograd_grad
+from pennylane._grad import jacobian as _autograd_jacobian
+
+from .interface_utils import get_interface
+
+
+# pylint: disable=import-outside-toplevel
+def grad(f: Callable, argnums: Union[Sequence[int], int] = 0) -> Callable:
+    """Compute the gradient in a jax-like manner for any interface.
+
+    Args:
+        f (Callable): a function with a single 0-D scalar output
+        argnums (Sequence[int] | int ) = 0 : which arguments to differentiate
+
+    Returns:
+        Callable: a function with the same signature as ``f`` that returns the gradient.
+
+    .. seealso:: :func:`pennylane.math.jacobian`
+
+    Note that this function follows the same design as jax. By default, the function will return the gradient
+    of the first argument, whether or not other arguments are trainable.
+
+    >>> import jax, torch, tensorflow as tf
+    >>> def f(x, y):
+    ...     return  x * y
+    >>> qml.math.grad(f)(qml.numpy.array(2.0), qml.numpy.array(3.0))
+    tensor(3., requires_grad=True)
+    >>> qml.math.grad(f)(jax.numpy.array(2.0), jax.numpy.array(3.0))
+    Array(3., dtype=float32, weak_type=True)
+    >>> qml.math.grad(f)(torch.tensor(2.0, requires_grad=True), torch.tensor(3.0, requires_grad=True))
+    tensor(3.)
+    >>> qml.math.grad(f)(tf.Variable(2.0), tf.Variable(3.0))
+    <tf.Tensor: shape=(), dtype=float32, numpy=3.0>
+
+    ``argnums`` can be provided to differentiate multiple arguments.
+
+    >>> qml.math.grad(f, argnums=(0,1))(torch.tensor(2.0, requires_grad=True), torch.tensor(3.0, requires_grad=True))
+    (tensor(3.), tensor(2.))
+
+    Note that the selected arguments *must* be of an appropriately trainable datatype, or an error may occur.
+
+    >>> qml.math.grad(f)(torch.tensor(1.0), torch.tensor(2.))
+    RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
+
+    """
+
+    argnums_integer = False
+    if isinstance(argnums, int):
+        argnums = (argnums,)
+        argnums_integer = True
+
+    def compute_grad(*args, **kwargs):
+        interface = get_interface(*args)
+
+        if interface == "autograd":
+            g = _autograd_grad(f, argnum=argnums)(*args, **kwargs)
+            return g[0] if argnums_integer else g
+
+        if interface == "jax":
+            import jax
+
+            g = jax.grad(f, argnums=argnums)(*args, **kwargs)
+            return g[0] if argnums_integer else g
+
+        if interface == "torch":
+            y = f(*args, **kwargs)
+            y.backward()
+            g = tuple(args[i].grad for i in argnums)
+            return g[0] if argnums_integer else g
+
+        if interface == "tensorflow":
+            import tensorflow as tf
+
+            with tf.GradientTape() as tape:
+                y = f(*args, **kwargs)
+
+            g = tape.gradient(y, tuple(args[i] for i in argnums))
+            return g[0] if argnums_integer else g
+
+        raise ValueError(f"Interface {interface} is not differentiable.")
+
+    return compute_grad
+
+
+# pylint: disable=import-outside-toplevel
+def _torch_jac(f, argnums, args, kwargs):
+    """Calculate a jacobian via torch."""
+    from torch.autograd.functional import jacobian as _torch_jac
+
+    argnums_torch = (argnums,) if isinstance(argnums, int) else argnums
+    trainable_args = tuple(args[i] for i in argnums_torch)
+
+    # keep track of output type to know how to unpack
+    output_type_cache = []
+
+    def partial_f(*_trainables):
+        full_args = list(args)
+        for argnum, value in zip(argnums_torch, _trainables, strict=True):
+            full_args[argnum] = value
+        result = f(*full_args, **kwargs)
+        output_type_cache.append(type(result))
+        return result
+
+    jac = _torch_jac(partial_f, trainable_args)
+    if output_type_cache[-1] is tuple:
+        return tuple(j[0] for j in jac) if isinstance(argnums, int) else jac
+    # else array
+    return jac[0] if isinstance(argnums, int) else jac
+
+
+# pylint: disable=import-outside-toplevel
+def _tensorflow_jac(f, argnums, args, kwargs):
+    """Calculate a jacobian via tensorflow"""
+    import tensorflow as tf
+
+    with tf.GradientTape() as tape:
+        y = f(*args, **kwargs)
+
+    if get_interface(y) != "tensorflow":
+        raise ValueError(
+            f"qml.math.jacobian does not work with tensorflow and non-tensor outputs. Got {y} of type {type(y)}."
+        )
+
+    argnums_integer = False
+    if isinstance(argnums, int):
+        argnums_tf = (argnums,)
+        argnums_integer = True
+    else:
+        argnums_tf = argnums
+
+    g = tape.jacobian(y, tuple(args[i] for i in argnums_tf))
+    return g[0] if argnums_integer else g
+
+
+# pylint: disable=import-outside-toplevel
+def jacobian(f: Callable, argnums: Union[Sequence[int], int] = 0) -> Callable:
+    """Compute the Jacobian in a jax-like manner for any interface.
+
+    Args:
+        f (Callable): a function with a vector valued output
+        argnums (Sequence[int] | int ) = 0 : which arguments to differentiate
+
+    Returns:
+        Callable: a function with the same signature as ``f`` that returns the jacobian
+
+    .. seealso:: :func:`pennylane.math.grad`
+
+    Note that this function follows the same design as jax. By default, the function will return the gradient
+    of the first argument, whether or not other arguments are trainable.
+
+    >>> import jax, torch, tensorflow as tf
+    >>> def f(x, y):
+    ...     return  x * y
+    >>> qml.math.jacobian(f)(qml.numpy.array([2.0, 3.0]), qml.numpy.array(3.0))
+    array([[3., 0.],
+              [0., 3.]])
+    >>> qml.math.jacobian(f)(jax.numpy.array([2.0, 3.0]), jax.numpy.array(3.0))
+    Array([[3., 0.],
+               [0., 3.]], dtype=float32)
+    >>> x_torch = torch.tensor([2.0, 3.0], requires_grad=True)
+    >>> y_torch = torch.tensor(3.0, requires_grad=True)
+    >>> qml.math.jacobian(f)(x_torch, y_torch)
+    tensor([[3., 0.],
+                [0., 3.]])
+    >>> qml.math.jacobian(f)(tf.Variable([2.0, 3.0]), tf.Variable(3.0))
+    <tf.Tensor: shape=(2, 2), dtype=float32, numpy=
+    array([[3., 0.],
+              [0., 3.]], dtype=float32)>
+
+    ``argnums`` can be provided to differentiate multiple arguments.
+
+    >>> qml.math.jacobian(f, argnums=(0,1))(x_torch, y_torch)
+    (tensor([[3., 0.],
+            [0., 3.]]),
+    tensor([2., 3.]))
+
+    While jax can handle taking jacobians of outputs with any pytree shape:
+
+    >>> def pytree_f(x):
+    ...     return {"a": 2*x, "b": 3*x}
+    >>> qml.math.jacobian(pytree_f)(jax.numpy.array(2.0))
+    {'a': Array(2., dtype=float32, weak_type=True),
+    'b': Array(3., dtype=float32, weak_type=True)}
+
+    Torch can only differentiate arrays and tuples:
+
+    >>> def tuple_f(x):
+    ...     return x**2, x**3
+    >>> qml.math.jacobian(tuple_f)(torch.tensor(2.0))
+    (tensor(4.), tensor(12.))
+    >>> qml.math.jacobian(pytree_f)(torch.tensor(2.0))
+    TypeError: The outputs of the user-provided function given to jacobian must be
+    either a Tensor or a tuple of Tensors but the given outputs of the user-provided
+    function has type <class 'dict'>.
+
+
+    But tensorflow and autograd can only handle array-valued outputs:
+
+    >>> qml.math.jacobian(tuple_f)(qml.numpy.array(2.0))
+    ValueError: autograd can only differentiate with respect to arrays, not <class 'tuple'>
+    >>> qml.math.jacobian(tuple_f)(tf.Variable(2.0))
+    ValueError: qml.math.jacobian does not work with tensorflow and non-tensor outputs.
+    Got (<tf.Tensor: shape=(), dtype=float32, numpy=4.0>,
+    <tf.Tensor: shape=(), dtype=float32, numpy=8.0>) of type <class 'tuple'>.
+
+    """
+
+    def compute_jacobian(*args, **kwargs):
+        interface = get_interface(*args)
+
+        if interface == "autograd":
+            return _autograd_jacobian(f, argnum=argnums)(*args, **kwargs)
+
+        if interface == "jax":
+            import jax
+
+            return jax.jacobian(f, argnums=argnums)(*args, **kwargs)
+
+        if interface == "torch":
+            return _torch_jac(f, argnums, args, kwargs)
+
+        if interface == "tensorflow":
+            return _tensorflow_jac(f, argnums, args, kwargs)
+
+        raise ValueError(f"Interface {interface} is not differentiable.")
+
+    return compute_jacobian
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index 64a74e9f2..6d9902d41 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -301,10 +301,12 @@ ar.register_function(
 )
 
 
-def _tf_convert_to_tensor(x, **kwargs):
+def _tf_convert_to_tensor(x, requires_grad=False, **kwargs):
     if isinstance(x, _i("tf").Tensor) and "dtype" in kwargs:
-        return _i("tf").cast(x, **kwargs)
-    return _i("tf").convert_to_tensor(x, **kwargs)
+        out = _i("tf").cast(x, **kwargs)
+    else:
+        out = _i("tf").convert_to_tensor(x, **kwargs)
+    return _i("tf").Variable(out) if requires_grad else out
 
 
 ar.register_function("tensorflow", "asarray", _tf_convert_to_tensor)
@@ -541,7 +543,7 @@ def _to_numpy_torch(x):
 ar.register_function("torch", "to_numpy", _to_numpy_torch)
 
 
-def _asarray_torch(x, dtype=None, **kwargs):
+def _asarray_torch(x, dtype=None, requires_grad=False, **kwargs):
     import torch
 
     dtype_map = {
@@ -556,9 +558,9 @@ def _asarray_torch(x, dtype=None, **kwargs):
         np.complex128: torch.complex128,
         "float64": torch.float64,
     }
-
-    if dtype in dtype_map:
-        return torch.as_tensor(x, dtype=dtype_map[dtype], **kwargs)
+    dtype = dtype_map.get(dtype, dtype)
+    if requires_grad:
+        return torch.tensor(x, dtype=dtype, **kwargs, requires_grad=True)
 
     return torch.as_tensor(x, dtype=dtype, **kwargs)
 
@@ -814,6 +816,14 @@ ar.register_function("jax", "block_diag", lambda x: _i("jax").scipy.linalg.block
 ar.register_function("jax", "gather", lambda x, indices: x[np.array(indices)])
 
 
+# pylint: disable=unused-argument
+def _asarray_jax(x, dtype=None, requires_grad=False, **kwargs):
+    return _i("jax").numpy.array(x, dtype=dtype, **kwargs)
+
+
+ar.register_function("jax", "asarray", _asarray_jax)
+
+
 def _ndim_jax(x):
     import jax.numpy as jnp
 


Test patch:
diff --git a/tests/math/test_grad.py b/tests/math/test_grad.py
new file mode 100644
index 000000000..6e47a9772
--- /dev/null
+++ b/tests/math/test_grad.py
@@ -0,0 +1,193 @@
+# Copyright 2024 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.
+"""
+Test the qml.math.grad and qml.math.jacobian functions.
+"""
+
+import numpy as np
+import pytest
+
+from pennylane import math
+
+pytestmark = pytest.mark.all_interfaces
+
+
+@pytest.mark.parametrize("grad_fn", (math.grad, math.jacobian))
+def test_no_interface_error_numpy(grad_fn):
+    """Test that an error is raised for an unknown interface."""
+
+    with pytest.raises(ValueError, match="Interface numpy is not differentiable"):
+        grad_fn(lambda x: x**2)(np.array(2.0))
+
+
+@pytest.mark.parametrize("interface", ("autograd", "jax", "tensorflow", "torch"))
+class TestGrad:
+    """Tests for qml.math.grad"""
+
+    def test_differentiate_first_arg(self, interface):
+        """Test that we just differentiate the first argument by default."""
+
+        def f(x, y):
+            return x * y
+
+        x = math.asarray(2.0, like=interface, requires_grad=True)
+        y = math.asarray(3.0, like=interface, requires_grad=True)
+
+        g = math.grad(f)(x, y)
+        if interface != "autograd":
+            assert math.get_interface(g) == interface
+        assert math.allclose(g, 3.0)
+
+    def test_multiple_argnums(self, interface):
+        """Test that we can differentiate multiple arguments."""
+
+        def g(x, y):
+            return 2 * x + 3 * y
+
+        x = math.asarray(0.5, like=interface, requires_grad=True)
+        y = math.asarray(2.5, like=interface, requires_grad=True)
+
+        g1, g2 = math.grad(g, argnums=(0, 1))(x, y)
+        if interface != "autograd":
+            assert math.get_interface(g1) == interface
+            assert math.get_interface(g2) == interface
+
+        assert math.allclose(g1, 2.0)
+        assert math.allclose(g2, 3.0)
+
+    def test_keyword_arguments(self, interface):
+        """Test that keyword arguments are considered."""
+
+        def f(x, *, constant):
+            return constant * x
+
+        x = math.asarray(2.0, like=interface, requires_grad=True)
+
+        g = math.grad(f)(x, constant=2.0)
+        assert math.allclose(g, 2.0)
+
+
+@pytest.mark.parametrize("interface", ("autograd", "jax", "tensorflow", "torch"))
+class TestJacobian:
+    """Tests for the math.jacobian function."""
+
+    def test_jac_first_arg(self, interface):
+        """Test taking the jacobian of the first argument."""
+
+        def f(x, y):
+            return x * y
+
+        x = math.asarray([2.0, 3.0], like=interface, requires_grad=True)
+        y = math.asarray(3.0, like=interface, requires_grad=True)
+
+        g = math.jacobian(f)(x, y)
+        print(g)
+        if interface != "autograd":
+            assert math.get_interface(g) == interface
+        expected = math.asarray([[3.0, 0.0], [0.0, 3.0]])
+        assert math.allclose(g, expected)
+
+    def test_multiple_argnums(self, interface):
+        """Test that we can differentiate multiple arguments."""
+
+        def g(x, y):
+            return 2 * x + 3 * y
+
+        x = math.asarray([0.5, 1.2], like=interface, requires_grad=True)
+        y = math.asarray([2.5, 4.8], like=interface, requires_grad=True)
+
+        g1, g2 = math.jacobian(g, argnums=(0, 1))(x, y)
+        if interface != "autograd":
+            assert math.get_interface(g1) == interface
+            assert math.get_interface(g2) == interface
+
+        assert math.allclose(g1, 2 * math.eye(2))
+        assert math.allclose(g2, 3 * math.eye(2))
+
+    def test_keyword_arguments(self, interface):
+        """Test that keyword arguments are considered."""
+
+        def f(x, *, constant):
+            return constant * x
+
+        x = math.asarray([2.0, 3.0], like=interface, requires_grad=True)
+
+        g = math.jacobian(f)(x, constant=2.0)
+        assert math.allclose(g, 2 * math.eye(2))
+
+
+class TestJacobianPytreeOutput:
+    """Test jacobians of non-array outputs."""
+
+    def test_jacobian_autograd_error(self):
+        """Test that an informative error if the output of a function isnt an array."""
+
+        def f(x):
+            return (x,)
+
+        with pytest.raises(
+            ValueError, match="autograd can only differentiate with respect to arrays,"
+        ):
+            math.jacobian(f)(math.asarray(2.0, like="autograd", requires_grad=True))
+
+    def test_jacobian_tf_error(self):
+        """Test that an informative error is raised if the output of a function isnt an array with tensorflow."""
+
+        def f(x):
+            return (x,)
+
+        with pytest.raises(
+            ValueError, match="qml.math.jacobian does not work with tensorflow and non-tensor"
+        ):
+            math.jacobian(f)(math.asarray(2.0, like="tensorflow", requires_grad=True))
+
+    @pytest.mark.parametrize("interface", ("torch", "jax"))
+    def test_tuple_output_scalar_argnum(self, interface):
+        """Test the shape outputted for a tuple valued function."""
+
+        def f(x):
+            return (5 * x, x**2)
+
+        jac = math.jacobian(f)(math.asarray(2.0, like=interface, requires_grad=True))
+        assert isinstance(jac, tuple)
+        assert len(jac) == 2
+        for j in jac:
+            assert math.get_interface(j) == interface
+
+        assert math.allclose(jac[0], 5)
+        assert math.allclose(jac[1], 4)
+
+    @pytest.mark.parametrize("interface", ("torch", "jax"))
+    def test_tuple_output_tuple_argnum(self, interface):
+        """Test the shape outputted for a tuple valued function with multiple traianble arguments."""
+
+        def f(x, y):
+            return (2 * y**2, 3 * x**2)
+
+        x = math.asarray(1.5, like=interface, requires_grad=True)
+        y = math.asarray(2.5, like=interface, requires_grad=True)
+
+        jac = math.jacobian(f, argnums=(0, 1))(x, y)
+
+        assert len(jac) == 2
+        assert len(jac[0]) == 2
+        assert len(jac[1]) == 2
+        dy0_dx0_expected = 0
+        assert math.allclose(jac[0][0], dy0_dx0_expected)
+        dy0_dx1_expected = 4 * y
+        assert math.allclose(jac[0][1], dy0_dx1_expected)
+        dy1_dx0_expected = 6 * x
+        assert math.allclose(jac[1][0], dy1_dx0_expected)
+        dy1_dx1_expected = 0
+        assert math.allclose(jac[1][1], dy1_dx1_expected)
diff --git a/tests/test_autograd_grad.py b/tests/test_autograd_grad.py
new file mode 100644
index 000000000..8851a6a8e
--- /dev/null
+++ b/tests/test_autograd_grad.py
@@ -0,0 +1,29 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Unit tests for qml.grad and qml.jacobian
+"""
+import pytest
+
+import pennylane as qml
+
+
+def test_informative_error_on_bad_shape():
+    """Test that an informative error is raised if taking the jacobian of a non-array."""
+
+    def f(x):
+        return (2 * x,)
+
+    with pytest.raises(ValueError, match="autograd can only differentiate with"):
+        qml.jacobian(f)(qml.numpy.array(2.0))
diff --git a/tests/transforms/test_classical_jacobian.py b/tests/transforms/test_classical_jacobian.py
index d2a4c560f..f5bb8996f 100644
--- a/tests/transforms/test_classical_jacobian.py
+++ b/tests/transforms/test_classical_jacobian.py
@@ -145,6 +145,18 @@ def test_autograd_without_argnum(circuit, args, expected_jac, diff_method, inter
 interfaces = ["tf"]
 
 
+def test_error_undefined_interface():
+    """Test that an error is raised in the qnode interface is not differentiable."""
+
+    @qml.qnode(qml.device("default.qubit"), interface=None)
+    def circuit(_var):
+        qml.RX(_var, 0)
+        return qml.expval(qml.Z(0))
+
+    with pytest.raises(ValueError, match="Undifferentiable interface numpy"):
+        classical_jacobian(circuit)(np.array(0.5))
+
+
 @pytest.mark.tf
 @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, all_args, class_jacs))