Repository: PennyLaneAI/pennylane
Pull request #1860 (2021-11-05 01:38:34)

Update deprecated use of `qml.grad`/`qml.jacobian` in docstrings
Updates docstrings that use `qml.grad`/`qml.jacobian` in light of #1773.

The following files were already okay (or are being updated separately):

1. Containing  `qml.grad`:
* `interfaces/batch/__init__.py`
* `transforms/unitary_to_rot.py` (addressed in #1859)
* `transforms/qfunc_transforms.py`
* `transforms/metric_tensor.py`
* `transforms/get_unitary_matrix.py`
* `transforms/batch_transform.py`
* `gradients/__init__.py`

2. Containing  `qml.jacobian`:
* `math/is_independent.py`: https://github.com/PennyLaneAI/pennylane/pull/1862
* `interfaces/autograd.py`
* `interfaces/batch/__init__.py`
* `gradients` (multiple files)

Features explicitly using `qml.grad`/`qml.jacobian`:

* `transforms/classical_jacobian.py`: uses the argnum syntax internally and offers this option too
* `optimize/shot_adaptive.py`: good, uses the `argnum` syntax

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.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/1860?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> :exclamation: No coverage uploaded for pull request base (`v0.19.0-rc0@fbfae24`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#section-missing-base-commit).
> The diff coverage is `n/a`.

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

```diff
@@              Coverage Diff               @@
##             v0.19.0-rc0    #1860   +/-   ##
==============================================
  Coverage               ?   98.96%           
==============================================
  Files                  ?      218           
  Lines                  ?    16377           
  Branches               ?        0           
==============================================
  Hits                   ?    16208           
  Misses                 ?      169           
  Partials               ?        0           
```



------

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


Patch:
diff --git a/doc/code/qml_hf.rst b/doc/code/qml_hf.rst
index 33db7274b..c49d806d0 100644
--- a/doc/code/qml_hf.rst
+++ b/doc/code/qml_hf.rst
@@ -1,7 +1,161 @@
-qml.hf
-======
-
-.. currentmodule:: pennylane.hf
-
-.. automodapi:: pennylane.hf
-
+qml.hf
+======
+
+This module provides the functionality to perform differentiable Hartree-Fock (HF) calculations and
+construct molecular Hamiltonians that can be differentiated with respect to nuclear coordinates and
+basis set parameters.
+
+The differentiable HF solver allows computing exact gradients with respect to
+molecular geometry, basis set, and circuit parameters simultaneously using the techniques of
+automatic differentiation available in `Autograd <https://github.com/HIPS/autograd>`__. This makes
+the solver more versatile and robust compared to non-differentiable tools that mainly rely
+on numerical approaches for computing gradients, which can lead to inaccuracies and instability.
+Additionally, optimizing the basis set parameters allows reaching lower ground-state energies
+without increasing the size of the basis set. Overall, the solver allows users to execute end-to-end
+differentiable algorithms for quantum chemistry. 
+
+.. currentmodule:: pennylane.hf
+
+.. automodapi:: pennylane.hf
+    :no-inheritance-diagram:
+
+Using the differentiable HF solver
+----------------------------------
+
+The HF solver computes the integrals over basis functions, constructs the relevant matrices, and
+performs self-consistent-field iterations to obtain a set of optimized molecular orbital
+coefficients. These coefficients and the computed integrals over basis functions are used to
+construct the one- and two-body electron integrals in the molecular orbital basis, which can be
+used to generate differentiable second-quantized Hamiltonians in the fermionic and qubit basis.
+
+The following code shows the construction of the Hamiltonian for the hydrogen molecule where the
+geometry of the molecule and the basis set parameters are all differentiable.
+
+.. code-block:: python3
+
+    import pennylane as qml
+    from pennylane import numpy as np
+
+    symbols = ["H", "H"]
+    # This initial geometry is suboptimal and will be optimized by the algorithm
+    geometry = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 2.0]], requires_grad=True)
+
+    # The exponents and contraction coefficients of the Gaussian basis functions
+    alpha = np.array([[3.42525091, 0.62391373, 0.1688554],
+                      [3.42525091, 0.62391373, 0.1688554]], requires_grad = True)
+    coeff = np.array([[0.15432897, 0.53532814, 0.44463454],
+                      [0.15432897, 0.53532814, 0.44463454]], requires_grad = True)
+
+We create a molecule object with differentiable atomic coordinates and basis set parameters and then
+construct the Hamiltonian.
+
+.. code-block:: python3
+
+    mol = qml.hf.Molecule(symbols, geometry, alpha=alpha, coeff=coeff)
+    args_mol = [geometry, alpha, coeff] # initial values of the differentiable parameters
+
+    hamiltonian = qml.hf.generate_hamiltonian(mol)(*args_mol)
+
+>>> print(hamiltonian)
+  ((-0.3596823592263728+0j)) [I0]
++ ((-0.11496335836149135+0j)) [Z3]
++ ((-0.1149633583614913+0j)) [Z2]
++ ((0.1308241430373499+0j)) [Z0]
++ ((0.1308241430373499+0j)) [Z1]
++ ((0.10316898251626505+0j)) [Z0 Z2]
++ ((0.10316898251626505+0j)) [Z1 Z3]
++ ((0.1532995999427217+0j)) [Z0 Z3]
++ ((0.1532995999427217+0j)) [Z2 Z1]
++ ((0.1540549585580985+0j)) [Z0 Z1]
++ ((0.1609686663985837+0j)) [Z3 Z2]
++ ((-0.05013061742645664+0j)) [Y0 X2 X3 Y1]
++ ((-0.05013061742645664+0j)) [X0 Y2 Y3 X1]
++ ((0.05013061742645664+0j)) [Y0 X2 Y3 X1]
++ ((0.05013061742645664+0j)) [X0 Y2 X3 Y1]
+
+The generated Hamiltonian can be used in a circuit where the molecular geometry, the basis set
+parameters, and the circuit parameters are optimized simultaneously. Further information about
+molecular geometry optimization with PennyLane is provided in this
+`paper <https://arxiv.org/abs/2106.13840>`__ and this
+`demo <https://pennylane.ai/qml/demos/tutorial_mol_geo_opt.html>`__ .
+
+.. code-block:: python3
+
+    dev = qml.device("default.qubit", wires=4)
+    hf_state = np.array([1, 1, 0, 0])
+    params = [np.array([0.0], requires_grad=True)] # initial values of the circuit parameters
+
+    def generate_circuit(mol):
+        @qml.qnode(dev)
+        def circuit(*args):
+            qml.BasisState(hf_state, wires=[0, 1, 2, 3])
+            qml.DoubleExcitation(*args[0][0], wires=[0, 1, 2, 3])
+            return qml.expval(qml.hf.generate_hamiltonian(mol)(*args[1:]))
+        return circuit
+
+Now that the circuit is defined, we can create a geometry and parameter optimization loop:
+
+.. code-block:: python3
+
+    for n in range(21):
+
+        mol = qml.hf.Molecule(symbols, geometry, alpha=alpha, coeff=coeff)
+        args = [params, geometry, alpha, coeff] # initial values of the differentiable parameters
+
+        # compute gradients with respect to the circuit parameters and update the parameters
+        g_params = qml.grad(generate_circuit(mol), argnum = 0)(*args)
+        params = params - 0.25 * g_params[0]
+
+        # compute gradients with respect to the nuclear coordinates and update geometry
+        g_coor = qml.grad(generate_circuit(mol), argnum = 1)(*args)
+        geometry = geometry - 0.5 * g_coor
+
+        # compute gradients with respect to the Gaussian exponents and update the exponents
+        g_alpha = qml.grad(generate_circuit(mol), argnum = 2)(*args)
+        alpha = alpha - 0.25 * g_alpha
+
+        # compute gradients with respect to the Gaussian contraction coefficients and update them
+        g_coeff = qml.grad(generate_circuit(mol), argnum = 3)(*args)
+        coeff = coeff - 0.25 * g_coeff
+
+        if n%5 == 0:
+            print(f'Step: {n}, Energy: {generate_circuit(mol)(*args)}, Maximum Absolute Force: {abs(g_coor).max()}')
+
+
+Running this optimization, we get the following output in atomic units:
+
+.. code-block:: text
+
+    Step: 0, Energy: -1.0491709019853235, Maximum Force: 0.15801947189250276
+    Step: 5, Energy: -1.134986263549686, Maximum Force: 0.037660772858684355
+    Step: 10, Energy: -1.1399960673348044, Maximum Force: 0.005175326423590643
+    Step: 15, Energy: -1.1403213849556897, Maximum Force: 0.0004138322831406249
+    Step: 20, Energy: -1.1403680839521801, Maximum Force: 8.223301624310508e-06
+
+Note that the computed energy is lower than the Full-CI energy, -1.1373060483 Ha, computed without
+optimizing the basis set parameters.
+
+The components of the HF solver can also be differentiated individually. For instance, the overlap
+integral can be differentiated with respect to the basis set parameters as follows
+
+.. code-block:: python3
+
+    symbols = ["H", "H"]
+    geometry = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 2.0]], requires_grad=False)
+    alpha = np.array([[3.42525091, 0.62391373, 0.1688554],
+                    [3.42525091, 0.62391373, 0.1688554]], requires_grad = True)
+    coeff = np.array([[0.15432897, 0.53532814, 0.44463454],
+                    [0.15432897, 0.53532814, 0.44463454]], requires_grad = True)
+
+    mol = qml.hf.Molecule(symbols, geometry, alpha=alpha, coeff=coeff)
+    args = [alpha, coeff]
+
+    a = mol.basis_set[0]
+    b = mol.basis_set[1]
+
+    g_alpha = qml.grad(qml.hf.generate_overlap(a, b), argnum = 0)(*args)
+    g_coeff = qml.grad(qml.hf.generate_overlap(a, b), argnum = 1)(*args)
+
+>>> print(g_alpha)
+[[ 0.00169332 -0.14826928 -0.37296693]
+ [ 0.00169332 -0.14826928 -0.37296693]]
diff --git a/doc/introduction/chemistry.rst b/doc/introduction/chemistry.rst
index 1e02ad009..955b655cd 100644
--- a/doc/introduction/chemistry.rst
+++ b/doc/introduction/chemistry.rst
@@ -1,10 +1,10 @@
 Quantum Chemistry
 =================
 
-PennyLane provides a quantum chemistry module ``qchem`` to perform quantum simulations
-of the electronic structure of molecules. ``qchem`` contains tools to construct the
-electronic Hamiltonian of molecules, and uses PennyLane to implement the Variational Quantum
-Eigensolver (VQE) algorithm.
+PennyLane provides a differentiable Hartree-Fock solver module :mod:`~.hf` and a quantum chemistry
+module :mod:`~.qchem` to perform quantum simulations of the electronic structure of molecules. These
+modules contain tools to construct the electronic Hamiltonian of molecules that can be used to
+implement the Variational Quantum Eigensolver (VQE) algorithm.
 
 .. figure:: ../_static/sketch_pennylane_qchem.png
     :width: 80%
@@ -161,26 +161,30 @@ where a quantum computer is used to prepare the trial wave function of a molecul
 the expectation value of the *electronic Hamiltonian*, while a classical optimizer is used to
 find its ground state.
 
-We can use :class:`~.ExpvalCost` to automatically create the required PennyLane QNodes and define
-the cost function:
+PennyLane supports treating Hamiltonians just like any other observable, and the 
+expectation value of a Hamiltonian can be calculated using ``qml.expval``:
 
 .. code-block:: python
 
     import pennylane as qml
-    
+
     dev = qml.device('default.qubit', wires=4)
 
-    def circuit(params, wires):
-        qml.BasisState(np.array([1, 1, 0, 0]), wires=wires)
-        for i in wires:
+    hamiltonian = 2.0 * qml.PauliZ(0) @ qml.PauliZ(1)
+
+    @qml.qnode(dev)
+    def circuit(params):
+        qml.BasisState(np.array([1, 1, 0, 0]), wires=[0,1,2,3])
+        for i in range(4):
             qml.Rot(*params[i], wires=i)
         qml.CNOT(wires=[2, 3])
         qml.CNOT(wires=[2, 0])
         qml.CNOT(wires=[3, 1])
+        return qml.expval(hamiltonian)
 
-    cost = qml.ExpvalCost(circuit, hamiltonian, dev, interface="torch")
-    params = torch.rand([4, 3])
-    cost(params)
+    rng = np.random.default_rng(seed=42)
+    params = rng.random([4, 3])
+    circuit(params)
 
 The rotation angles can be optimized using the machine learning interface of choice
 until the energy difference between two consecutive iterations has converged to near zero.
diff --git a/doc/introduction/circuits.rst b/doc/introduction/circuits.rst
index 13cf9ff86..173141b35 100644
--- a/doc/introduction/circuits.rst
+++ b/doc/introduction/circuits.rst
@@ -148,9 +148,9 @@ For example:
 Executing this, we will get an output of size ``(3, 2)``:
 
 >>> circuit(0.5)
-[[0.33333333 1.        ]
-[0.2        1.        ]
-[0.012      0.868     ]]
+tensor([[ 1.   ,  1.   ],
+        [ 0.2  ,  1.   ],
+        [-0.022,  0.876]], requires_grad=True)
 
 Custom wire labels
 ^^^^^^^^^^^^^^^^^^
@@ -166,7 +166,7 @@ This is done by passing an iterable of wire labels to the ``wires`` argument:
 
 .. code-block:: python
 
-    dev = qml.device('default.qubit', wires=['wire1', 'wire2'], shots=1000)
+    dev = qml.device('default.qubit', wires=['wire1', 'wire2'])
 
 In the quantum function you can now use your own labels to address wires:
 
@@ -180,13 +180,13 @@ In the quantum function you can now use your own labels to address wires:
 
 Allowed wire labels can be of the following types:
 
-* *strings* like ``wires=['a', 'd', 'b', ...]`` or ``wires=['ancilla', 'q1', 'q2', ...]``,
+* *strings* like ``wires=['a', 'd', 'b', ...]`` or ``wires=['auxiliary', 'q1', 'q2', ...]``,
 
 * *integers* like ``wires=[0, 4, 7]`` or even ``wires=[-1, 0, 4]``
 
 * *floats* and other *numbers* like ``wires=[1., 2., 4.]``
 
-* *mixed types* like ``wires=['ancilla', -1, 0, 'q3']``
+* *mixed types* like ``wires=['auxiliary', -1, 0, 'q3']``
 
 .. note::
 
@@ -212,14 +212,14 @@ The QNode can be used to compute the result of a quantum circuit as if it was a
 function. It takes the same arguments as the original quantum function:
 
 >>> circuit(np.pi/4, 0.7)
-0.7648421872844883
+tensor(0.764, requires_grad=True)
 
 To view the quantum circuit given specific parameter values, we can use the :func:`~.pennylane.draw`
 transform:
 
 >>> print(qml.draw(circuit)(np.pi/4, 0.7))
-0: ──RZ(0.785)──╭C───────────┤
-1: ─────────────╰X──RY(0.7)──┤ ⟨Z⟩
+ wire1: ──RZ(0.785)──╭C───────────┤     
+ wire2: ─────────────╰X──RY(0.7)──┤ ⟨Z⟩ 
 
 .. _intro_vcirc_decorator:
 
@@ -328,7 +328,6 @@ pennylane.collections.qnode_collection.QNodeCollection
 >>> params = [0.54, 0.12]
 >>> qnodes(params)
 array([-0.02854835  0.99280864])
-
 Functions are available to process QNode collections, including :func:`~.pennylane.collections.dot`,
 :func:`~.pennylane.collections.sum`, and :func:`~.pennylane.collections.apply`:
 
diff --git a/doc/introduction/interfaces/numpy.rst b/doc/introduction/interfaces/numpy.rst
index 3aa63f062..4993d925e 100644
--- a/doc/introduction/interfaces/numpy.rst
+++ b/doc/introduction/interfaces/numpy.rst
@@ -108,8 +108,8 @@ with respect to both QNode parameters ``phi`` and ``theta``:
 
 .. code-block:: python
 
-    phi = np.array([0.5, 0.1])
-    theta = 0.2
+    phi = np.array([0.5, 0.1], requires_grad=True)
+    theta = np.array(0.2, requires_grad=True)
     dcircuit = qml.grad(circuit3)
 
 Evaluating this gradient function at specific parameter values:
@@ -125,12 +125,16 @@ How does PennyLane know which arguments of a quantum function to differentiate,
 For example, you may want to pass arguments to a QNode but *not* have
 PennyLane consider them when computing gradients.
 
-As a basic rule, **all positional arguments provided to the QNode are assumed to be differentiable
-by default**. To accomplish this, arguments are internally turned into arrays of the PennyLane NumPy module,
+All positional arguments provided to the QNode are assumed to be differentiable
+by default, but this behaviour is deprecated. In a future release, only arguments explicitly
+marked as trainable or marked using the `argnum` keyword argument will be treated as trainable. Not explicitly marking arguments currently raises a 
+deprecation warning.
+
+For now, arguments are internally turned into arrays from the PennyLane NumPy module,
 which have a special flag ``requires_grad`` specifying whether they are trainable or not:
 
 >>> from pennylane import numpy as np
->>> np.array([0.1, 0.2])
+>>> np.array([0.1, 0.2], requires_grad=True)
 tensor([0.1, 0.2], requires_grad=True)
 
 If you would like to provide explicit non-differentiable arguments to the
@@ -169,8 +173,8 @@ and two non-differentiable arguments ``data`` and ``wires``:
 
 For ``data``, which is a PennyLane NumPy array, we can simply specify ``requires_grad=False``:
 
->>> np.random.seed(42)  # make the results reproducable
->>> data = np.random.random([2**3], requires_grad=False)
+>>> rng = np.random.default_rng(seed=42)  # make the results reproducable
+>>> data = rng.random([2**3], requires_grad=False)
 
 But ``wires`` is a list in this example, and if we turn it into a PennyLane NumPy array we would have to
 create a device that understands custom wire labels of this type.
@@ -180,14 +184,14 @@ QNode using keyword argument syntax:
 >>> wires = [2, 0, 1]
 >>> weights = np.array([0.1, 0.2, 0.3])
 >>> circuit(weights, data, wires=wires)
-0.4124409353413991
+tensor(-0.03404945, requires_grad=True)
 
 When we compute the derivative, arguments with ``requires_grad=False`` as well as arguments passed as
 keyword arguments are ignored by :func:`~.grad`:
 
 >>> grad_fn = qml.grad(circuit)
 >>> grad_fn(weights, data, wires=wires)
-[-4.1382126e-02  0.0000000e+00 -6.9388939e-18]
+array([ 3.41633993e-03,  8.23993651e-18, -6.93889390e-18])
 
 .. note::
 
@@ -235,7 +239,7 @@ lead to a final expectation value of 0.5:
     opt = qml.GradientDescentOptimizer(stepsize=0.4)
 
     steps = 100
-    params = np.array([0.011, 0.012, 0.05])
+    params = np.array([0.011, 0.012, 0.05], requires_grad=True)
 
     for i in range(steps):
         # update the circuit parameters
@@ -244,9 +248,9 @@ lead to a final expectation value of 0.5:
 The final weights and circuit value are:
 
 >>> params
-array([ 0.19846757,  0.012     ,  1.03559806])
+tensor([0.19846757, 0.012     , 1.03559806], requires_grad=True)
 >>> circuit4(params)
-0.5
+tensor(0.5, requires_grad=True)
 
 For more details on the NumPy optimizers, check out the tutorials, as well as the
 :mod:`pennylane.optimize` documentation.
@@ -274,8 +278,9 @@ How does automatic differentiation work in the case where the QNode returns mult
 If we were to naively try computing the gradient of ``circuit5`` using the :func:`~.grad` function,
 
 >>> g1 = qml.grad(circuit5)
->>> params = np.array([np.pi/2, 0.2])
+>>> params = np.array([np.pi/2, 0.2], requires_grad=True)
 >>> g1(params)
+TypeError: Grad only applies to real scalar-output functions. Try jacobian, elementwise_grad or holomorphic_grad.
 
 we would get an error message. This is because the `gradient <https://en.wikipedia.org/wiki/Gradient>`_ is
 only defined for scalar functions, i.e., functions which return a single value. In the case where the QNode
@@ -365,15 +370,17 @@ to the ``scipy.minimize`` function:
     def cost(x):
         return np.abs(circuit(x) - 0.5) ** 2
 
-    params = np.array([0.011, 0.012, 0.05])
+    params = np.array([0.011, 0.012, 0.05], requires_grad=True)
 
     minimize(cost, params, method='BFGS')
 
 Some of the SciPy minimization methods require information about the gradient
 of the cost function via the ``jac`` keyword argument. This is easy to include; we
-can simply create a function that computes the gradient using ``qml.grad``:
+can simply create a function that computes the gradient using ``qml.grad``.  Since
+``minimize`` does not use our wrapped version of numpy, we need to explicitly 
+specify which arguments are trainable via the ``argnum`` keyword.
 
->>> minimize(cost, params, method='BFGS', jac=qml.grad(cost))
+>>> minimize(cost, params, method='BFGS', jac=qml.grad(cost, argnum=0))
       fun: 6.3491130264451484e-18
  hess_inv: array([[ 1.85642354e+00, -8.84954187e-22,  3.89539943e+00],
        [-8.84954187e-22,  1.00000000e+00, -4.02571211e-21],
diff --git a/doc/introduction/templates.rst b/doc/introduction/templates.rst
index 2a6b5921c..7451affa2 100644
--- a/doc/introduction/templates.rst
+++ b/doc/introduction/templates.rst
@@ -316,7 +316,7 @@ The shape can for example be used to construct random weights at the beginning o
     weights = np.random.random(size=shape)
 
 >>> circuit(weights)
-0.7258859204630561
+tensor(0.72588592, requires_grad=True)
 
 If a template takes more than one weight tensor, the ``shape`` method returns a list of shape tuples.
 
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 3a30829d7..9fde1b5d6 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -625,6 +625,15 @@
 
 <h3>Improvements</h3>
 
+* The `qml.AmplitudeEmbedding` template now supports just-in-time compilation when
+  used with either `@jax.jit` or `@tf.function`, resulting in improved performance.
+  [(#1845)](https://github.com/PennyLaneAI/pennylane/pull/1845)
+
+* A new utility function `qml.math.is_abstract(tensor)` has been added. This function
+  returns `True` if the tensor is *abstract*; that is, it has no value or shape.
+  This can occur if within a function that has been just-in-time compiled.
+  [(#1845)](https://github.com/PennyLaneAI/pennylane/pull/1845)
+
 * Refactored the `expand_fn` functionality in the Device class to avoid any
   edge cases leading to failures with plugins.
   [(#1840)](https://github.com/PennyLaneAI/pennylane/pull/1840)
diff --git a/doc/xanadu_theme/layout.html b/doc/xanadu_theme/layout.html
index a28078a47..bf25be232 100644
--- a/doc/xanadu_theme/layout.html
+++ b/doc/xanadu_theme/layout.html
@@ -308,5 +308,17 @@ if (downloadNote.length >= 1) {
 }
 </script>
 <script type="text/javascript">jQuery.noConflict(true);</script>
+<script type="text/javascript">
+  let divs = document.getElementsByClassName('field-odd');
+
+  for (let x = 0; x < divs.length; x++) {
+      let div = divs[x];
+      let content = div.innerHTML.trim();
+
+      if (content == 'orphan') {
+          div.style.display = 'none';
+      }
+  }
+</script>
         {% include "footer.html" %}
 {%- endblock %}
diff --git a/pennylane/_grad.py b/pennylane/_grad.py
index 89da4b80b..bb9765857 100644
--- a/pennylane/_grad.py
+++ b/pennylane/_grad.py
@@ -20,6 +20,7 @@ from functools import partial
 import numpy as onp
 from autograd import jacobian as _jacobian
 from autograd.core import make_vjp as _make_vjp
+from autograd.numpy.numpy_boxes import ArrayBox
 from autograd.extend import vspace
 from autograd.wrap_util import unary_to_nary
 
@@ -85,8 +86,11 @@ class grad:
         argnum = []
 
         for idx, arg in enumerate(args):
+
             trainable = getattr(arg, "requires_grad", None)
-            if trainable is None:
+            array_box = isinstance(arg, ArrayBox)
+
+            if trainable is None and not array_box:
 
                 warnings.warn(
                     "Starting with PennyLane v0.20.0, when using Autograd, inputs "
@@ -95,6 +99,8 @@ class grad:
                     "identified.",
                     UserWarning,
                 )
+
+            if trainable is None:
                 trainable = True
 
             if trainable:
@@ -184,8 +190,11 @@ def jacobian(func, argnum=None):
         argnum = []
 
         for idx, arg in enumerate(args):
+
             trainable = getattr(arg, "requires_grad", None)
-            if trainable is None:
+            array_box = isinstance(arg, ArrayBox)
+
+            if trainable is None and not array_box:
 
                 warnings.warn(
                     "Starting with PennyLane v0.20.0, when using Autograd, inputs "
@@ -194,6 +203,8 @@ def jacobian(func, argnum=None):
                     "identified.",
                     UserWarning,
                 )
+
+            if trainable is None:
                 trainable = True
 
             if trainable:
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index fe1e23217..9db32366b 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -630,18 +630,12 @@ class DefaultQubit(QubitDevice):
         state = self._asarray(state, dtype=self.C_DTYPE)
         n_state_vector = state.shape[0]
 
-        if state.ndim != 1 or n_state_vector != 2 ** len(device_wires):
+        if len(qml.math.shape(state)) != 1 or n_state_vector != 2 ** len(device_wires):
             raise ValueError("State vector must be of length 2**wires.")
 
-        norm_error_message = "Sum of amplitudes-squared does not equal one."
-        if qml.math.get_interface(state) != "jax":
+        if not qml.math.is_abstract(state):
             if not qml.math.allclose(qml.math.linalg.norm(state, ord=2), 1.0, atol=tolerance):
-                raise ValueError(norm_error_message)
-        else:
-            # Case for jax without jit, full_lower is an attribute for abstract tracers
-            if not hasattr(qml.math.linalg.norm(state, ord=2), "full_lower"):
-                if not qml.math.allclose(qml.math.linalg.norm(state, ord=2), 1.0, atol=tolerance):
-                    raise ValueError(norm_error_message)
+                raise ValueError("Sum of amplitudes-squared does not equal one.")
 
         if len(device_wires) == self.num_wires and sorted(device_wires) == device_wires:
             # Initialize the entire wires with the state
diff --git a/pennylane/devices/default_qubit_autograd.py b/pennylane/devices/default_qubit_autograd.py
index 92518e417..260ca58c8 100644
--- a/pennylane/devices/default_qubit_autograd.py
+++ b/pennylane/devices/default_qubit_autograd.py
@@ -50,7 +50,7 @@ class DefaultQubitAutograd(DefaultQubit):
     ...     qml.RX(x[1], wires=0)
     ...     qml.Rot(x[0], x[1], x[2], wires=0)
     ...     return qml.expval(qml.PauliZ(0))
-    >>> weights = np.array([0.2, 0.5, 0.1])
+    >>> weights = np.array([0.2, 0.5, 0.1], requires_grad=True)
     >>> grad_fn = qml.grad(circuit)
     >>> print(grad_fn(weights))
     array([-2.2526717e-01 -1.0086454e+00  1.3877788e-17])
diff --git a/pennylane/math/__init__.py b/pennylane/math/__init__.py
index b3fb9d462..558e42f02 100644
--- a/pennylane/math/__init__.py
+++ b/pennylane/math/__init__.py
@@ -55,6 +55,7 @@ from .utils import (
     allequal,
     cast,
     cast_like,
+    is_abstract,
     convert_like,
     get_interface,
     requires_grad,
@@ -84,6 +85,7 @@ __all__ = [
     "allequal",
     "cast",
     "cast_like",
+    "is_abstract",
     "convert_like",
     "get_interface",
     "requires_grad",
diff --git a/pennylane/math/is_independent.py b/pennylane/math/is_independent.py
index ee2d6321f..5a2bea1e6 100644
--- a/pennylane/math/is_independent.py
+++ b/pennylane/math/is_independent.py
@@ -23,6 +23,7 @@ a function is independent of its arguments for the interfaces
 import warnings
 
 import numpy as np
+from pennylane import numpy as pnp
 
 from autograd.tracer import isbox, new_box, trace_stack
 from autograd.core import VJPNode
@@ -199,6 +200,10 @@ def _get_random_args(args, interface, num, seed, bounds):
             tuple(np.random.random(np.shape(arg)) * width + bounds[0] for arg in args)
             for _ in range(num)
         ]
+        if interface == "autograd":
+
+            # Mark the arguments as trainable with Autograd
+            rnd_args = pnp.array(rnd_args, requires_grad=True)
 
     return rnd_args
 
@@ -296,7 +301,7 @@ def is_independent(
     The analytic and numeric tests used are as follows.
 
     - The analytic test performed depends on the provided ``interface``,
-       both in its method and its degree of reliability.
+      both in its method and its degree of reliability.
 
     - For the numeric test, the function is evaluated at a series of random positions,
       and the outputs numerically compared to verify that the output
@@ -328,8 +333,8 @@ def is_independent(
 
     .. code-block:: pycon
 
-        >>> x = np.array([0.2, 9.1, -3.2])
-        >>> weights = np.array([1.1, -0.7, 1.8])
+        >>> x = np.array([0.2, 9.1, -3.2], requires_grad=True)
+        >>> weights = np.array([1.1, -0.7, 1.8], requires_grad=True)
         >>> qml.math.is_independent(lin, "autograd", (x,), {"weights": weights})
         False
 
diff --git a/pennylane/math/quantum.py b/pennylane/math/quantum.py
index 69edd2459..2bc65bfe0 100644
--- a/pennylane/math/quantum.py
+++ b/pennylane/math/quantum.py
@@ -67,7 +67,7 @@ def cov_matrix(prob, obs, wires=None, diag_approx=False):
     We can now compute the covariance matrix:
 
     >>> shape = qml.templates.StronglyEntanglingLayers.shape(n_layers=2, n_wires=3)
-    >>> weights = np.random.random(shape)
+    >>> weights = np.random.random(shape, requires_grad=True)
     >>> cov = qml.math.cov_matrix(circuit(weights), obs_list)
     >>> cov
     array([[0.98707611, 0.03665537],
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index 0a21c0952..46a3bd2cc 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -283,7 +283,6 @@ ar.register_function("torch", "diag", lambda x, k=0: _i("torch").diag(x, diagona
 ar.register_function("torch", "expand_dims", lambda x, axis: _i("torch").unsqueeze(x, dim=axis))
 ar.register_function("torch", "shape", lambda x: tuple(x.shape))
 ar.register_function("torch", "gather", lambda x, indices: x[indices])
-ar.register_function("torch", "gather", lambda x, indices: x[indices])
 
 try:
     if semantic_version.match(">=1.10", _i("torch").__version__):
diff --git a/pennylane/math/utils.py b/pennylane/math/utils.py
index f2b973971..b9398e924 100644
--- a/pennylane/math/utils.py
+++ b/pennylane/math/utils.py
@@ -200,6 +200,119 @@ def get_interface(tensor):
     return res
 
 
+def is_abstract(tensor, like=None):
+    """Returns True if the tensor is considered abstract.
+
+    Abstract arrays have no internal value, and are used primarily when
+    tracing Python functions, for example, in order to perform just-in-time
+    (JIT) compilation.
+
+    Abstract tensors most commonly occur within a function that has been
+    decorated using ``@tf.function`` or ``@jax.jit``.
+
+    .. note::
+
+        Currently Autograd tensors and Torch tensors will always return ``False``.
+        This is because:
+
+        - Autograd does not provide JIT compilation, and
+
+        - ``@torch.jit.script`` is not currently compatible with QNodes.
+
+    Args:
+        tensor (tensor_like): input tensor
+        like (str): The name of the interface. Will be determined automatically
+            if not provided.
+
+    Returns:
+        bool: whether the tensor is abstract or not
+
+    **Example**
+
+    Consider the following JAX function:
+
+    .. code-block:: python
+
+        import jax
+        from jax import numpy as jnp
+
+        def function(x):
+            print("Value:", x)
+            print("Abstract:", qml.math.is_abstract(x))
+            return jnp.sum(x ** 2)
+
+    When we execute it, we see that the tensor is not abstract; it has known value:
+
+    >>> x = jnp.array([0.5, 0.1])
+    >>> function(x)
+    Value: [0.5, 0.1]
+    Abstract: False
+    DeviceArray(0.26, dtype=float32)
+
+    However, if we use the ``@jax.jit`` decorator, the tensor will now be abstract:
+
+    >>> x = jnp.array([0.5, 0.1])
+    >>> jax.jit(function)(x)
+    Value: Traced<ShapedArray(float32[2])>with<DynamicJaxprTrace(level=0/1)>
+    Abstract: True
+    DeviceArray(0.26, dtype=float32)
+
+    Note that JAX uses an abstract *shaped* array, so although we won't be able to
+    include conditionals within our function that depend on the value of the tensor,
+    we *can* include conditionals that depend on the shape of the tensor.
+
+    Similarly, consider the following TensorFlow function:
+
+    .. code-block:: python
+
+        import tensorflow as tf
+
+        def function(x):
+            print("Value:", x)
+            print("Abstract:", qml.math.is_abstract(x))
+            return tf.reduce_sum(x ** 2)
+
+    >>> x = tf.Variable([0.5, 0.1])
+    >>> function(x)
+    Value: <tf.Variable 'Variable:0' shape=(2,) dtype=float32, numpy=array([0.5, 0.1], dtype=float32)>
+    Abstract: False
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.26>
+
+    If we apply the ``@tf.function`` decorator, the tensor will now be abstract:
+
+    >>> tf.function(function)(x)
+    Value: <tf.Variable 'Variable:0' shape=(2,) dtype=float32>
+    Abstract: True
+    <tf.Tensor: shape=(), dtype=float32, numpy=0.26>
+    """
+    interface = like or get_interface(tensor)
+
+    if interface == "jax":
+        import jax
+        from jax.interpreters.partial_eval import DynamicJaxprTracer
+
+        if isinstance(tensor, (jax.ad.JVPTracer, jax.interpreters.batching.BatchTracer)):
+            # Tracer objects will be used when computing gradients or applying transforms.
+            # If the value of the tracer is known, it will contain a ConcreteArray.
+            # Otherwise, it will be abstract.
+            return not isinstance(tensor.aval, jax.core.ConcreteArray)
+
+        return isinstance(tensor, DynamicJaxprTracer)
+
+    if interface == "tensorflow":
+        import tensorflow as tf
+        from tensorflow.python.framework.ops import EagerTensor
+
+        return not isinstance(tf.convert_to_tensor(tensor), EagerTensor)
+
+    # Autograd does not have a JIT
+
+    # QNodes do not currently support TorchScript:
+    #   NotSupportedError: Compiled functions can't take variable number of arguments or
+    #   use keyword-only arguments with defaults.
+    return False
+
+
 def requires_grad(tensor, interface=None):
     """Returns True if the tensor is considered trainable.
 
diff --git a/pennylane/templates/embeddings/amplitude.py b/pennylane/templates/embeddings/amplitude.py
index e0200cf85..d58303939 100644
--- a/pennylane/templates/embeddings/amplitude.py
+++ b/pennylane/templates/embeddings/amplitude.py
@@ -190,9 +190,13 @@ class AmplitudeEmbedding(Operation):
             # normalize
             norm = qml.math.sum(qml.math.abs(feature_set) ** 2)
 
-            if not qml.math.allclose(norm, 1.0, atol=TOLERANCE):
+            if qml.math.is_abstract(norm):
                 if normalize or pad_with:
-                    feature_set = feature_set / np.sqrt(norm)
+                    feature_set = feature_set / qml.math.sqrt(norm)
+
+            elif not qml.math.allclose(norm, 1.0, atol=TOLERANCE):
+                if normalize or pad_with:
+                    feature_set = feature_set / qml.math.sqrt(norm)
                 else:
                     raise ValueError(
                         f"Features must be a vector of norm 1.0; got norm {norm}."
diff --git a/pennylane/tracker.py b/pennylane/tracker.py
index 27115d4ba..e006585be 100644
--- a/pennylane/tracker.py
+++ b/pennylane/tracker.py
@@ -64,7 +64,7 @@ class Tracker:
             qml.RX(x, wires=0)
             return qml.expval(qml.PauliZ(0))
 
-        x = np.array(0.1)
+        x = np.array(0.1, requires_grad=True)
 
         with qml.Tracker(dev) as tracker:
             qml.grad(circuit)(x)
@@ -102,8 +102,9 @@ class Tracker:
         >>> def shots_info(totals, history, latest):
         ...     if 'shots' in latest.keys():
         ...         print("Total shots: ", totals['shots'])
+        >>> x = np.array(0.1, requires_grad=True)
         >>> with qml.Tracker(circuit.device, callback=shots_info) as tracker:
-        ...     qml.grad(circuit)(0.1)
+        ...     qml.grad(circuit)(x)
         Total shots:  100
         Total shots:  200
         Total shots:  300
diff --git a/pennylane/transforms/batch_params.py b/pennylane/transforms/batch_params.py
index 64563d1ed..877b06b3c 100644
--- a/pennylane/transforms/batch_params.py
+++ b/pennylane/transforms/batch_params.py
@@ -71,21 +71,22 @@ def batch_params(tape, all_operations=False):
 
     >>> batch_size = 3
     >>> x = np.linspace(0.1, 0.5, batch_size)
-    >>> weights = np.random.random((batch_size, 10, 3, 3))
+    >>> rng = np.random.default_rng(seed=1234)
+    >>> weights = rng.random((batch_size, 10, 3, 3), requires_grad=True)
 
     If we evaluate the QNode with these inputs, we will get an output
     of shape ``(batch_size,)``:
 
     >>> circuit(x, weights)
-    [-0.30773348  0.23135516  0.13086565]
+    tensor([ 0.00800498,  0.2735391 , -0.24395442], requires_grad=True)
 
     QNodes with a batch dimension remain fully differentiable:
 
     >>> cost_fn = lambda x, weights: np.sum(circuit(x, weights))
     >>> cost_fn(x, weights)
-    -0.8581269507766536
+    tensor(0.03758966, requires_grad=True)
     >>> qml.grad(cost_fn)(x, weights)[0]
-    [ 0.23235464  0.00928953 -0.30083487]
+    array([-0.30262974,  0.06320878,  0.00811555])
 
     If we pass the ``all_operations`` argument, we can specify that
     *all* operation parameters in the transformed QNode, regardless of whether they
@@ -93,7 +94,7 @@ def batch_params(tape, all_operations=False):
 
     .. code-block:: python
 
-        @functools.partial(qml.batch_params, all_operations=True)
+        @qml.batch_params(all_operations=True)
         @qml.beta.qnode(dev)
         def circuit(x, weights):
             qml.RX(x, wires=0)
@@ -104,9 +105,9 @@ def batch_params(tape, all_operations=False):
     >>> cost_fn = lambda x, weights: np.sum(circuit(x, weights))
     >>> weights.requires_grad = False
     >>> cost_fn(x, weights)
-    0.5497108163237583
+    tensor(0.03758966, requires_grad=True)
     >>> qml.grad(cost_fn)(x, weights)[0]
-    0.43792281188363347
+    -0.30262974103192636
     """
     params = tape.get_parameters(trainable_only=not all_operations)
     output_tapes = []
@@ -133,4 +134,4 @@ def batch_params(tape, all_operations=False):
         new_tape.set_parameters(p, trainable_only=not all_operations)
         output_tapes.append(new_tape)
 
-    return output_tapes, qml.math.stack
+    return output_tapes, lambda x: qml.math.squeeze(qml.math.stack(x))
diff --git a/pennylane/transforms/batch_transform.py b/pennylane/transforms/batch_transform.py
index 8f0b8463a..742882ee8 100644
--- a/pennylane/transforms/batch_transform.py
+++ b/pennylane/transforms/batch_transform.py
@@ -150,7 +150,8 @@ class batch_transform:
 
     Batch tape transforms are fully differentiable:
 
-    >>> gradient = qml.grad(circuit)(-0.5)
+    >>> x = np.array(-0.5, requires_grad=True)
+    >>> gradient = qml.grad(circuit)(x)
     >>> print(gradient)
     2.5800122591960153
 
diff --git a/pennylane/transforms/unitary_to_rot.py b/pennylane/transforms/unitary_to_rot.py
index 270ff5f1f..a03a8eff1 100644
--- a/pennylane/transforms/unitary_to_rot.py
+++ b/pennylane/transforms/unitary_to_rot.py
@@ -100,8 +100,9 @@ def unitary_to_rot(tape):
             transformed_qnode = qml.QNode(transformed_qfunc, dev)
 
         >>> g = qml.grad(transformed_qnode)
-        >>> g([0.2, 0.3])
-        [array(-0.24454992), array(-0.23913254)]
+        >>> params = np.array([0.2, 0.3], requires_grad=True)
+        >>> g(params)
+        array([ 0.00296633, -0.29392145])
 
         However, the following example will **not** be differentiable:
 


Test patch:
diff --git a/tests/interfaces/test_batch_autograd.py b/tests/interfaces/test_batch_autograd.py
index f3c011644..4a2ae90a9 100644
--- a/tests/interfaces/test_batch_autograd.py
+++ b/tests/interfaces/test_batch_autograd.py
@@ -744,7 +744,7 @@ class TestHigherOrderDerivatives:
             np.array([-2.0, 0], requires_grad=True),
         ],
     )
-    def test_parameter_shift_hessian(self, params, tol):
+    def test_parameter_shift_hessian(self, params, tol, recwarn):
         """Tests that the output of the parameter-shift transform
         can be differentiated using autograd, yielding second derivatives."""
         dev = qml.device("default.qubit.autograd", wires=2)
@@ -786,7 +786,11 @@ class TestHigherOrderDerivatives:
         )
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-    def test_adjoint_hessian(self, tol):
+        # Check that no deprecation warnigns are emitted due to trainable
+        # parameters with the Autograd interface
+        assert len(recwarn) == 0
+
+    def test_adjoint_hessian(self, tol, recwarn):
         """Since the adjoint hessian is not a differentiable transform,
         higher-order derivatives are not supported."""
         dev = qml.device("default.qubit.autograd", wires=2)
@@ -811,6 +815,10 @@ class TestHigherOrderDerivatives:
 
         assert np.allclose(res, np.zeros([2, 2]), atol=tol, rtol=0)
 
+        # Check that no additional deprecation warnigns are emitted due to trainable
+        # parameters with the Autograd interface
+        assert len(recwarn) == 0
+
     def test_max_diff(self, tol):
         """Test that setting the max_diff parameter blocks higher-order
         derivatives"""
diff --git a/tests/math/test_is_abstract.py b/tests/math/test_is_abstract.py
new file mode 100644
index 000000000..c7433634b
--- /dev/null
+++ b/tests/math/test_is_abstract.py
@@ -0,0 +1,159 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Unit tests for the :func:`pennylane.math.is_abstract` function.
+"""
+import functools
+
+import numpy as np
+import pytest
+
+import pennylane as qml
+
+
+class TestTensorFlow:
+    """Test that tensorflow integrates with is_abstract"""
+
+    def test_eager(self):
+        """Test that no tensors are abstract when in eager mode"""
+        tf = pytest.importorskip("tensorflow")
+
+        def cost(x, w):
+            y = x ** 2
+            z = tf.ones([2, 2])
+
+            assert tf.executing_eagerly()
+
+            assert not qml.math.is_abstract(w)
+            assert not qml.math.is_abstract(x)
+            assert not qml.math.is_abstract(y)
+            assert not qml.math.is_abstract(z)
+
+            return tf.reduce_sum(y)
+
+        x = tf.Variable([0.5, 0.1])
+        w = tf.constant(0.1)
+
+        assert not qml.math.is_abstract(w)
+        assert not qml.math.is_abstract(x)
+
+        with tf.GradientTape() as tape:
+            res = cost(x, w)
+
+        assert res == 0.26
+
+        grad = tape.gradient(res, x)
+        assert np.allclose(grad, 2 * x)
+
+    @pytest.mark.parametrize("jit_compile", [True, False])
+    def test_jit(self, jit_compile):
+        """Test that all tensors are abstract when in autograd mode"""
+        tf = pytest.importorskip("tensorflow")
+
+        @tf.function(jit_compile=jit_compile)
+        def cost(x, w):
+            y = x ** 2
+            z = tf.ones([2, 2])
+
+            assert not tf.executing_eagerly()
+
+            assert qml.math.is_abstract(w)
+            assert qml.math.is_abstract(x)
+            assert qml.math.is_abstract(y)
+            assert qml.math.is_abstract(z)
+
+            return tf.reduce_sum(y)
+
+        x = tf.Variable([0.5, 0.1])
+        w = tf.constant(0.1)
+
+        assert not qml.math.is_abstract(w)
+        assert not qml.math.is_abstract(x)
+
+        with tf.GradientTape() as tape:
+            res = cost(x, w)
+
+        assert res == 0.26
+
+        grad = tape.gradient(res, x)
+        assert np.allclose(grad, 2 * x)
+
+
+class TestJAX:
+    """Test that JAX integrates with is_abstract"""
+
+    def test_eager(self):
+        """Test that no tensors are abstract when in eager mode"""
+        jax = pytest.importorskip("jax")
+        jnp = jax.numpy
+
+        def cost(x, w):
+            y = x ** 2
+            z = jnp.ones([2, 2])
+
+            assert not qml.math.is_abstract(w)
+            assert not qml.math.is_abstract(x)
+            assert not qml.math.is_abstract(y)
+            assert not qml.math.is_abstract(z)
+
+            return jnp.sum(y)
+
+        x = jnp.array([0.5, 0.1])
+        w = jnp.array(0.1)
+
+        assert not qml.math.is_abstract(w)
+        assert not qml.math.is_abstract(x)
+
+        res = cost(x, w)
+        assert res == 0.26
+
+        grad = jax.grad(cost, argnums=0)(x, w)
+        assert np.allclose(grad, 2 * x)
+
+    def test_jit(self):
+        """Test that all tensors are abstract when in JIT mode.
+        Note that `jax.grad` has slightly different behaviour, and will
+        avoid making abstract tensors for non-differentiable arguments."""
+        jax = pytest.importorskip("jax")
+        jnp = jax.numpy
+
+        @functools.partial(jax.jit, static_argnums=[2])
+        def cost(x, w, w_is_abstract=True):
+            y = x ** 2
+            z = jnp.ones([2, 2])
+
+            assert qml.math.is_abstract(w) == w_is_abstract
+
+            assert qml.math.is_abstract(x)
+            assert qml.math.is_abstract(y)
+            assert qml.math.is_abstract(z)
+
+            return jnp.sum(y)
+
+        x = jnp.array([0.5, 0.1])
+        w = jnp.array(0.1)
+
+        assert not qml.math.is_abstract(w)
+        assert not qml.math.is_abstract(x)
+
+        res = cost(x, w)
+        assert res == 0.26
+
+        # since we only specify argnums=0, w will not be abstract
+        grad = jax.grad(cost, argnums=0)(x, w, w_is_abstract=False)
+        assert np.allclose(grad, 2 * x)
+
+        # Otherwise, w will be abstract
+        grad = jax.grad(cost, argnums=[0, 1])(x, w, w_is_abstract=True)
+        assert np.allclose(grad[0], 2 * x)
diff --git a/tests/math/test_is_independent.py b/tests/math/test_is_independent.py
index 87e8bc38c..6953b584d 100644
--- a/tests/math/test_is_independent.py
+++ b/tests/math/test_is_independent.py
@@ -17,6 +17,7 @@ Unit tests for the :func:`pennylane.math.is_independent` function.
 import pytest
 
 import numpy as np
+from pennylane import numpy as pnp
 
 import pennylane as qml
 from pennylane.math import is_independent
@@ -205,6 +206,33 @@ class TestIsIndependentAutograd:
         assert not is_independent(f, self.interface, args, {"kw": True})
         assert is_independent(jac, self.interface, args, {"kw": True})
 
+    def test_no_trainable_params_deprecation_grad(self, recwarn):
+        """Tests that no deprecation warning arises when using qml.grad with
+        qml.math.is_independent.
+        """
+
+        def lin(x):
+            return pnp.sum(x)
+
+        x = pnp.array([0.2, 9.1, -3.2], requires_grad=True)
+        jac = qml.grad(lin)
+        assert qml.math.is_independent(jac, "autograd", (x,), {})
+
+        assert len(recwarn) == 0
+
+    def test_no_trainable_params_deprecation_jac(self, recwarn):
+        """Tests that no deprecation arises when using qml.jacobian with
+        qml.math.is_independent."""
+
+        def lin(x, weights=None):
+            return np.dot(x, weights)
+
+        x = pnp.array([0.2, 9.1, -3.2], requires_grad=True)
+        weights = pnp.array([1.1, -0.7, 1.8], requires_grad=True)
+        jac = qml.jacobian(lin)
+        assert qml.math.is_independent(jac, "autograd", (x,), {"weights": weights})
+        assert len(recwarn) == 0
+
 
 if have_jax:
 
diff --git a/tests/templates/test_embeddings/test_amplitude.py b/tests/templates/test_embeddings/test_amplitude.py
index 9f47df53b..846ed9271 100644
--- a/tests/templates/test_embeddings/test_amplitude.py
+++ b/tests/templates/test_embeddings/test_amplitude.py
@@ -291,6 +291,24 @@ class TestInterfaces:
 
         assert qml.math.allclose(res, res2, atol=tol, rtol=0)
 
+    def test_jax_jit(self, tol):
+        """Tests jax tensors when using JIT."""
+
+        jax = pytest.importorskip("jax")
+        import jax.numpy as jnp
+
+        features = jnp.array([1 / 2, 0, 1 / 2, 0, 1 / 2, 1 / 2, 0, 0])
+
+        dev = qml.device("default.qubit", wires=3)
+
+        circuit = jax.jit(qml.QNode(circuit_template, dev, interface="jax"))
+        circuit2 = jax.jit(qml.QNode(circuit_decomposed, dev, interface="jax"))
+
+        res = circuit(features)
+        res2 = circuit2(features)
+
+        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+
     def test_tf(self, tol):
         """Tests tf tensors."""
 
@@ -308,6 +326,23 @@ class TestInterfaces:
 
         assert qml.math.allclose(res, res2, atol=tol, rtol=0)
 
+    def test_tf_jit(self, tol):
+        """Tests tf tensors when using JIT."""
+
+        tf = pytest.importorskip("tensorflow")
+
+        features = tf.Variable([1 / 2, 0, 1 / 2, 0, 1 / 2, 1 / 2, 0, 0])
+
+        dev = qml.device("default.qubit", wires=3)
+
+        circuit = tf.function(jit_compile=True)(qml.QNode(circuit_template, dev, interface="tf"))
+        circuit2 = tf.function(jit_compile=True)(qml.QNode(circuit_decomposed, dev, interface="tf"))
+
+        res = circuit(features)
+        res2 = circuit2(features)
+
+        assert qml.math.allclose(res, res2, atol=tol, rtol=0)
+
     def test_torch(self, tol):
         """Tests torch tensors."""
 
diff --git a/tests/transforms/test_batch_params.py b/tests/transforms/test_batch_params.py
index a40844cca..91f1ad2ba 100644
--- a/tests/transforms/test_batch_params.py
+++ b/tests/transforms/test_batch_params.py
@@ -41,7 +41,7 @@ def test_simple_circuit(mocker):
 
     spy = mocker.spy(circuit.device, "batch_execute")
     res = circuit(data, x, weights)
-    assert res.shape == (batch_size, 1, 4)
+    assert res.shape == (batch_size, 4)
     assert len(spy.call_args[0][0]) == batch_size
 
 
@@ -224,7 +224,7 @@ def test_all_operations(mocker):
 
     spy = mocker.spy(circuit.device, "batch_execute")
     res = circuit(x, weights)
-    assert res.shape == (batch_size, 1, 4)
+    assert res.shape == (batch_size, 4)
     assert len(spy.call_args[0][0]) == batch_size