Repository: PennyLaneAI/pennylane
Pull request #1647 (2021-09-13 14:15:57)

Remove (mistakenly added) newlines in docstring
**Context:**
I managed to accidentally add newlines via my code review suggestions in #1591.

**Description of the Change:**
Removes all the mistakenly added newlines in the docstring.

**Benefits:**
Things look like they should, and the docstring will render correctly.

**Possible Drawbacks:**
None

**Related GitHub Issues:**
None


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.

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index b08b371fe..8e3a3b2e1 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -273,6 +273,7 @@
   [(#1549)](https://github.com/PennyLaneAI/pennylane/pull/1549)
   [(#1608)](https://github.com/PennyLaneAI/pennylane/pull/1608)
   [(#1618)](https://github.com/PennyLaneAI/pennylane/pull/1618)
+  [(#1637)](https://github.com/PennyLaneAI/pennylane/pull/1637)
 
   For example:
 
@@ -402,6 +403,11 @@ and requirements-ci.txt (unpinned). This latter would be used by the CI.
 
 <h3>Breaking changes</h3>
 
+* Specifying `shots=None` with `qml.sample` was previously deprecated.
+  From this release onwards, setting `shots=None` when sampling will
+  raise an error also for `default.qubit.jax`.
+  [(#1629)](https://github.com/PennyLaneAI/pennylane/pull/1629)
+
 * An error is raised during QNode creation when a user requests backpropagation on
   a device with finite-shots.
   [(#1588)](https://github.com/PennyLaneAI/pennylane/pull/1588)
@@ -458,7 +464,8 @@ This release contains contributions from (in alphabetical order):
 
 Vishnu Ajith, Akash Narayanan B, Thomas Bromley, Olivia Di Matteo, Sahaj Dhamija, Tanya Garg, Josh Izaac,
 Prateek Jain, Ankit Khandelwal, Christina Lee, Ian McLean, Johannes Jakob Meyer, Romain Moyard, Esteban Payares,
-Pratul Saini, Maria Schuld, Arshpreet Singh, Ingrid Strandberg, Slimane Thabet, David Wierichs, Vincent Wong.
+Pratul Saini, Maria Schuld, Arshpreet Singh, Ingrid Strandberg, Slimane Thabet, Antal Száva, David Wierichs,
+Vincent Wong.
 
 # Release 0.17.0 (current release)
 
diff --git a/pennylane/devices/default_qubit_jax.py b/pennylane/devices/default_qubit_jax.py
index 4c8e1a0ab..a9246cd89 100644
--- a/pennylane/devices/default_qubit_jax.py
+++ b/pennylane/devices/default_qubit_jax.py
@@ -14,9 +14,7 @@
 """This module contains an jax implementation of the :class:`~.DefaultQubit`
 reference plugin.
 """
-
-import warnings
-
+import pennylane as qml
 from pennylane.operation import DiagonalOperation
 from pennylane.devices import DefaultQubit
 from pennylane.devices import jax_ops
@@ -252,15 +250,13 @@ class DefaultQubitJax(DefaultQubit):
             List[int]: the sampled basis states
         """
         if self.shots is None:
-            warnings.warn(
-                "The number of shots has to be explicitly set on the jax device "
-                "when using sample-based measurements. Since no shots are specified, "
-                "a default of 1000 shots is used.\n"
-                "This warning will be replaced with an error in a future release.",
-                UserWarning,
+
+            raise qml.QuantumFunctionError(
+                "The number of shots has to be explicitly set on the device "
+                "when using sample-based measurements."
             )
 
-        shots = self.shots or 1000
+        shots = self.shots
 
         if self._prng_key is None:
             # Assuming op-by-op, so we'll just make one.
diff --git a/pennylane/gradients/gradient_transform.py b/pennylane/gradients/gradient_transform.py
index d70ef25f5..1846be0f5 100644
--- a/pennylane/gradients/gradient_transform.py
+++ b/pennylane/gradients/gradient_transform.py
@@ -141,21 +141,25 @@ class gradient_transform(qml.batch_transform):
         self.hybrid = hybrid
         super().__init__(transform_fn, expand_fn=expand_fn, differentiable=differentiable)
 
-    def qnode_execution_wrapper(self, qnode, targs, tkwargs):
+    def default_qnode_wrapper(self, qnode, targs, tkwargs):
         # Here, we overwrite the QNode execution wrapper in order
         # to take into account that classical processing may be present
         # inside the QNode.
         hybrid = tkwargs.pop("hybrid", self.hybrid)
-        _wrapper = super().qnode_execution_wrapper(qnode, targs, tkwargs)
+        _wrapper = super().default_qnode_wrapper(qnode, targs, tkwargs)
         cjac_fn = qml.transforms.classical_jacobian(qnode)
 
         def jacobian_wrapper(*args, **kwargs):
             qjac = _wrapper(*args, **kwargs)
-            cjac = cjac_fn(*args, **kwargs)
 
             if any(m.return_type is qml.operation.Probability for m in qnode.qtape.measurements):
                 qjac = qml.math.squeeze(qjac)
 
+            if not hybrid:
+                return qjac
+
+            cjac = cjac_fn(*args, **kwargs)
+
             if isinstance(cjac, tuple):
                 # Classical processing of multiple arguments is present. Return qjac @ cjac.
                 jacs = [
@@ -167,7 +171,7 @@ class gradient_transform(qml.batch_transform):
 
             is_square = cjac.shape == (1,) or (cjac.ndim == 2 and cjac.shape[0] == cjac.shape[1])
 
-            if not hybrid or (is_square and qml.math.allclose(cjac, qml.numpy.eye(cjac.shape[0]))):
+            if is_square and qml.math.allclose(cjac, qml.numpy.eye(cjac.shape[0])):
                 # Classical Jacobian is the identity. No classical processing
                 # is present inside the QNode.
                 return qjac
diff --git a/pennylane/gradients/vjp.py b/pennylane/gradients/vjp.py
index 0dfa57f68..fa886b687 100644
--- a/pennylane/gradients/vjp.py
+++ b/pennylane/gradients/vjp.py
@@ -318,7 +318,8 @@ def batch_vjp(tapes, dys, gradient_fn, reduction="append", gradient_kwargs=None)
             vjp_ = processing_fns[t_idx](res_t)
 
             if vjp_ is None:
-                vjps.append(None)
+                if reduction == "append":
+                    vjps.append(None)
                 continue
 
             if isinstance(reduction, str):
diff --git a/pennylane/interfaces/batch/autograd.py b/pennylane/interfaces/batch/autograd.py
index 9c3060a8c..97a56bb13 100644
--- a/pennylane/interfaces/batch/autograd.py
+++ b/pennylane/interfaces/batch/autograd.py
@@ -101,7 +101,7 @@ def _execute(
     if the nth-order derivative is requested. Do not set this argument unless you
     understand the consequences!
     """
-    with qml.tape.Unwrap(*tapes):
+    with qml.tape.Unwrap(*tapes, set_trainable=False):
         res, jacs = execute_fn(tapes, **gradient_kwargs)
 
     for i, r in enumerate(res):
@@ -158,7 +158,7 @@ def vjp(
         """Returns the vector-Jacobian product with given
         parameter values and output gradient dy"""
 
-        dy = dy[0]
+        dy = [qml.math.T(d) for d in dy[0]]
         jacs = ans[1]
 
         if jacs:
diff --git a/pennylane/interfaces/batch/tensorflow.py b/pennylane/interfaces/batch/tensorflow.py
index 3678a5cd1..1c4592dbf 100644
--- a/pennylane/interfaces/batch/tensorflow.py
+++ b/pennylane/interfaces/batch/tensorflow.py
@@ -65,9 +65,6 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
         list[list[tf.Tensor]]: A nested list of tape results. Each element in
         the returned list corresponds in order to the provided tapes.
     """
-    with qml.tape.Unwrap(*tapes):
-        # Forward pass: execute the tapes
-        res, jacs = execute_fn(tapes, **gradient_kwargs)
 
     parameters = []
     params_unwrapped = []
@@ -75,6 +72,8 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
     for i, tape in enumerate(tapes):
         # store the trainable parameters
         params = tape.get_parameters(trainable_only=False)
+        tape.trainable_params = qml.math.get_trainable_indices(params)
+
         parameters += [p for i, p in enumerate(params) if i in tape.trainable_params]
 
         # store all unwrapped parameters
@@ -82,6 +81,11 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
             [i.numpy() if isinstance(i, (tf.Variable, tf.Tensor)) else i for i in params]
         )
 
+    with qml.tape.Unwrap(*tapes, set_trainable=False):
+        # Forward pass: execute the tapes
+        res, jacs = execute_fn(tapes, **gradient_kwargs)
+
+    for i, tape in enumerate(tapes):
         # convert output to TensorFlow tensors
         r = np.hstack(res[i]) if res[i].dtype == np.dtype("object") else res[i]
         res[i] = tf.convert_to_tensor(r)
@@ -92,6 +96,8 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
             """Returns the vector-Jacobian product with given
             parameter values and output gradient dy"""
 
+            dy = [qml.math.T(d) for d in dy]
+
             if jacs:
                 # Jacobians were computed on the forward pass (mode="forward")
                 # No additional quantum evaluations needed; simply compute the VJPs directly.
diff --git a/pennylane/interfaces/batch/torch.py b/pennylane/interfaces/batch/torch.py
index 151cdd591..b81d6d16c 100644
--- a/pennylane/interfaces/batch/torch.py
+++ b/pennylane/interfaces/batch/torch.py
@@ -81,7 +81,7 @@ class ExecuteTapes(torch.autograd.Function):
         ctx.max_diff = kwargs["max_diff"]
         ctx._n = kwargs.get("_n", 1)
 
-        with qml.tape.Unwrap(*ctx.tapes):
+        with qml.tape.Unwrap(*ctx.tapes, set_trainable=False):
             res, ctx.jacs = ctx.execute_fn(ctx.tapes, **ctx.gradient_kwargs)
 
         # if any input tensor uses the GPU, the output should as well
@@ -93,11 +93,17 @@ class ExecuteTapes(torch.autograd.Function):
                 break
 
         for i, r in enumerate(res):
-            if r.dtype is np.dtype("object"):
+            if isinstance(r, np.ndarray) and r.dtype is np.dtype("object"):
                 # For backwards compatibility, we flatten ragged tape outputs
                 r = np.hstack(r)
 
-            res[i] = torch.as_tensor(r, device=ctx.torch_device)
+            if isinstance(r, (list, tuple)):
+                res[i] = [torch.as_tensor(t) for t in r]
+
+                if isinstance(r, tuple):
+                    res[i] = tuple(res[i])
+            else:
+                res[i] = torch.as_tensor(r, device=ctx.torch_device)
 
             if ctx.jacs:
                 ctx.jacs[i] = torch.as_tensor(ctx.jacs[i], device=ctx.torch_device)
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 9726d48de..d718042e1 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -115,32 +115,23 @@ class QNode:
             quantum structure will only be constructed on the *first* evaluation of the QNode,
             and is stored and re-used for further quantum evaluations. Only set this to False
             if it is known that the underlying quantum structure is **independent of QNode input**.
-
         max_expansion (int): The number of times the internal circuit should be expanded when
             executed on a device. Expansion occurs when an operation or measurement is not
             supported, and results in a gate decomposition. If any operations in the decomposition
             remain unsupported by the device, another expansion occurs.
-
         h (float): step size for the finite difference method
-
         order (int): The order of the finite difference method to use. ``1`` corresponds
             to forward finite differences, ``2`` to centered finite differences.
-
         shift (float): the size of the shift for two-term parameter-shift gradient computations
-
         adjoint_cache (bool): For TensorFlow and PyTorch interfaces and adjoint differentiation,
-
             this indicates whether to save the device state after the forward pass.  Doing so saves a
             forward execution. Device state automatically reused with autograd and JAX interfaces.
-
         argnum (int, list(int), None): Which argument(s) to compute the Jacobian
             with respect to. When there are fewer parameters specified than the
             total number of trainable parameters, the jacobian is being estimated. Note
             that this option is only applicable for the following differentiation methods:
             ``"parameter-shift"``, ``"finite-diff"`` and ``"reversible"``.
-
         kwargs: used to catch all unrecognized keyword arguments and provide a user warning
-
             about them
 
     **Example**
@@ -1141,32 +1132,23 @@ def qnode(
             quantum structure will only be constructed on the *first* evaluation of the QNode,
             and is stored and re-used for further quantum evaluations. Only set this to False
             if it is known that the underlying quantum structure is **independent of QNode input**.
-
         max_expansion (int): The number of times the internal circuit should be expanded when
             executed on a device. Expansion occurs when an operation or measurement is not
             supported, and results in a gate decomposition. If any operations in the decomposition
             remain unsupported by the device, another expansion occurs.
-
         h (float): step size for the finite difference method
-
         order (int): The order of the finite difference method to use. ``1`` corresponds
             to forward finite differences, ``2`` to centered finite differences.
-
         shift (float): the size of the shift for two-term parameter-shift gradient computations
-
         adjoint_cache (bool): For TensorFlow and PyTorch interfaces and adjoint differentiation,
-
             this indicates whether to save the device state after the forward pass.  Doing so saves a
             forward execution. Device state automatically reused with autograd and JAX interfaces.
-
         argnum (int, list(int), None): Which argument(s) to compute the Jacobian
             with respect to. When there are fewer parameters specified than the
             total number of trainable parameters, the jacobian is being estimated. Note
             that this option is only applicable for the following differentiation methods:
             ``"parameter-shift"``, ``"finite-diff"`` and ``"reversible"``.
-
         kwargs: used to catch all unrecognized keyword arguments and provide a user warning
-
             about them
 
     **Example**
diff --git a/pennylane/transforms/batch_transform.py b/pennylane/transforms/batch_transform.py
index ba4a47309..5c7999ce0 100644
--- a/pennylane/transforms/batch_transform.py
+++ b/pennylane/transforms/batch_transform.py
@@ -14,6 +14,7 @@
 """Contains tools and decorators for registering batch transforms."""
 # pylint: disable=too-few-public-methods
 import functools
+import types
 
 import pennylane as qml
 
@@ -158,9 +159,60 @@ class batch_transform:
         self.transform_fn = transform_fn
         self.expand_fn = expand_fn
         self.differentiable = differentiable
+        self.qnode_wrapper = self.default_qnode_wrapper
         functools.update_wrapper(self, transform_fn)
 
-    def qnode_execution_wrapper(self, qnode, targs, tkwargs):
+    def custom_qnode_wrapper(self, fn):
+        """Register a custom QNode execution wrapper function
+        for the batch transform.
+
+        **Example**
+
+        .. code-block:: python
+
+            def my_transform(tape, *targs, **tkwargs):
+                ...
+                return tapes, processing_fn
+
+            @my_transform.custom_qnode_wrapper
+            def my_custom_qnode_wrapper(self, qnode, targs, tkwargs):
+                def wrapper_fn(*args, **kwargs):
+                    # construct QNode
+                    qnode.construct(args, kwargs)
+                    # apply transform to QNode's tapes
+                    tapes, processing_fn = self.construct(qnode.qtape, *targs, **tkwargs)
+                    # execute tapes and return processed result
+                    ...
+                    return processing_fn(results)
+                return wrapper_fn
+
+        The custom QNode execution wrapper must have arguments
+        ``self`` (the batch transform object), ``qnode`` (the input QNode
+        to transform and execute), ``targs`` and ``tkwargs`` (the transform
+        arguments and keyword arguments respectively).
+
+        It should return a callable object that accepts the *same* arguments
+        as the QNode, and returns the transformed numerical result.
+
+        The default :meth:`~.default_qnode_wrapper` method may be called
+        if only pre- or post-processing dependent on QNode arguments is required:
+
+        .. code-block:: python
+
+            @my_transform.custom_qnode_wrapper
+            def my_custom_qnode_wrapper(self, qnode, targs, tkwargs):
+                transformed_qnode = self.default_qnode_wrapper(qnode)
+
+                def wrapper_fn(*args, **kwargs):
+                    args, kwargs = pre_process(args, kwargs)
+                    res = transformed_qnode(*args, **kwargs)
+                    ...
+                    return ...
+                return wrapper_fn
+        """
+        self.qnode_wrapper = types.MethodType(fn, self)
+
+    def default_qnode_wrapper(self, qnode, targs, tkwargs):
         """A wrapper method that takes a QNode and transform arguments,
         and returns a function that 'wraps' the QNode execution.
 
@@ -177,7 +229,7 @@ class batch_transform:
             interface = qnode.interface
 
             # TODO: extract gradient_fn from QNode
-            gradient_fn = qnode.diff_method
+            gradient_fn = getattr(qnode, "gradient_fn", qnode.diff_method)
 
             if interface is None or not self.differentiable:
                 gradient_fn = None
@@ -205,10 +257,10 @@ class batch_transform:
             # tapes, fn = some_transform(tape, *transform_args)
             return self.construct(qnode, *targs, **tkwargs)
 
-        if isinstance(qnode, qml.QNode):
+        if isinstance(qnode, (qml.QNode, qml.ExpvalCost)):
             # Input is a QNode:
             # result = some_transform(qnode, *transform_args)(*qnode_args)
-            wrapper = self.qnode_execution_wrapper(qnode, targs, tkwargs)
+            wrapper = self.qnode_wrapper(qnode, targs, tkwargs)
             wrapper = functools.wraps(qnode)(wrapper)
 
         else:
@@ -230,7 +282,7 @@ class batch_transform:
             targs = (qnode,) + targs
 
             def wrapper(qnode):
-                _wrapper = self.qnode_execution_wrapper(qnode, targs, tkwargs)
+                _wrapper = self.qnode_wrapper(qnode, targs, tkwargs)
                 _wrapper = functools.wraps(qnode)(_wrapper)
                 return _wrapper
 


Test patch:
diff --git a/tests/devices/test_default_qubit_jax.py b/tests/devices/test_default_qubit_jax.py
index 5593618cf..8db2f269e 100644
--- a/tests/devices/test_default_qubit_jax.py
+++ b/tests/devices/test_default_qubit_jax.py
@@ -186,22 +186,20 @@ class TestQNodeIntegration:
         assert not np.all(a == b)
 
     def test_sampling_analytic_mode(self):
-        """Test that when sampling with shots=None, dev uses 1000 shots and
-        raises deprecation warning.
-        """
+        """Test that when sampling with shots=None an error is raised."""
         dev = qml.device("default.qubit.jax", wires=1, shots=None)
 
         @qml.qnode(dev, interface="jax", diff_method=None)
         def circuit():
             return qml.sample(qml.PauliZ(wires=0))
 
-        with pytest.warns(
-            UserWarning, match="The number of shots has to be explicitly set on the jax device"
+        with pytest.raises(
+            qml.QuantumFunctionError,
+            match="The number of shots has to be explicitly set on the device "
+            "when using sample-based measurements.",
         ):
             res = circuit()
 
-        assert len(res) == 1000
-
     def test_gates_dont_crash(self):
         """Test for gates that weren't covered by other tests."""
         dev = qml.device("default.qubit.jax", wires=2, shots=1000)
diff --git a/tests/interfaces/test_batch_autograd.py b/tests/interfaces/test_batch_autograd.py
index 50f7cef69..aca6ec911 100644
--- a/tests/interfaces/test_batch_autograd.py
+++ b/tests/interfaces/test_batch_autograd.py
@@ -439,6 +439,38 @@ class TestAutogradExecuteIntegration:
         expected = [[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+    def test_tape_no_parameters(self, execute_kwargs, tol):
+        """Test that a tape with no parameters is correctly
+        ignored during the gradient computation"""
+        dev = qml.device("default.qubit", wires=1)
+
+        def cost(params):
+            with qml.tape.JacobianTape() as tape1:
+                qml.Hadamard(0)
+                qml.expval(qml.PauliX(0))
+
+            with qml.tape.JacobianTape() as tape2:
+                qml.RY(np.array(0.5, requires_grad=False), wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            with qml.tape.JacobianTape() as tape3:
+                qml.RY(params[0], wires=0)
+                qml.RX(params[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            return sum(execute([tape1, tape2, tape3], dev, **execute_kwargs))
+
+        params = np.array([0.1, 0.2], requires_grad=True)
+        x, y = params
+
+        res = cost(params)
+        expected = 1 + np.cos(0.5) + np.cos(x) * np.cos(y)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        grad = qml.grad(cost)(params)
+        expected = [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)]
+        assert np.allclose(grad, expected, atol=tol, rtol=0)
+
     def test_reusing_quantum_tape(self, execute_kwargs, tol):
         """Test re-using a quantum tape by passing new parameters"""
         a = np.array(0.1, requires_grad=True)
@@ -634,13 +666,14 @@ class TestAutogradExecuteIntegration:
 
         expected = np.array(
             [
-                [[-np.sin(x) / 2, 0], [np.sin(x) / 2, 0]],
+                [[-np.sin(x) / 2, 0], [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2]],
                 [
-                    [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2],
+                    [np.sin(x) / 2, 0],
                     [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
                 ],
             ]
         )
+
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
     def test_ragged_differentiation(self, execute_kwargs, tol):
diff --git a/tests/interfaces/test_batch_tensorflow.py b/tests/interfaces/test_batch_tensorflow.py
index 369dc48af..948d0a858 100644
--- a/tests/interfaces/test_batch_tensorflow.py
+++ b/tests/interfaces/test_batch_tensorflow.py
@@ -367,6 +367,36 @@ class TestTensorFlowExecuteIntegration:
         expected = [[-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]]
         assert np.allclose(expected, [agrad, bgrad], atol=tol, rtol=0)
 
+    def test_tape_no_parameters(self, execute_kwargs, tol):
+        """Test that a tape with no parameters is correctly
+        ignored during the gradient computation"""
+        dev = qml.device("default.qubit", wires=1)
+        params = tf.Variable([0.1, 0.2], dtype=tf.float64)
+        x, y = 1.0 * params
+
+        with tf.GradientTape() as t:
+            with qml.tape.JacobianTape() as tape1:
+                qml.Hadamard(0)
+                qml.expval(qml.PauliX(0))
+
+            with qml.tape.JacobianTape() as tape2:
+                qml.RY(0.5, wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            with qml.tape.JacobianTape() as tape3:
+                qml.RY(params[0], wires=0)
+                qml.RX(params[1], wires=0)
+                qml.expval(qml.PauliZ(0))
+
+            res = sum(execute([tape1, tape2, tape3], dev, **execute_kwargs))
+
+        expected = 1 + np.cos(0.5) + np.cos(x) * np.cos(y)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        grad = t.gradient(res, params)
+        expected = [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)]
+        assert np.allclose(grad, expected, atol=tol, rtol=0)
+
     def test_reusing_quantum_tape(self, execute_kwargs, tol):
         """Test re-using a quantum tape by passing new parameters"""
         a = tf.Variable(0.1, dtype=tf.float64)
@@ -600,12 +630,12 @@ class TestTensorFlowExecuteIntegration:
         expected = np.array(
             [
                 [
-                    [-tf.sin(x) / 2, -tf.sin(x) * tf.cos(y) / 2],
-                    [tf.sin(x) / 2, tf.cos(y) * tf.sin(x) / 2],
+                    [-tf.sin(x) / 2, tf.sin(x) / 2],
+                    [-tf.sin(x) * tf.cos(y) / 2, tf.cos(y) * tf.sin(x) / 2],
                 ],
                 [
-                    [0, -tf.cos(x) * tf.sin(y) / 2],
-                    [0, tf.cos(x) * tf.sin(y) / 2],
+                    [0, 0],
+                    [-tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2],
                 ],
             ]
         )
diff --git a/tests/interfaces/test_batch_torch.py b/tests/interfaces/test_batch_torch.py
index ba0e0ca41..2c9db85c4 100644
--- a/tests/interfaces/test_batch_torch.py
+++ b/tests/interfaces/test_batch_torch.py
@@ -439,6 +439,35 @@ class TestTorchExecuteIntegration:
         assert torch.allclose(a.grad, expected[0], atol=tol, rtol=0)
         assert torch.allclose(b.grad, expected[1], atol=tol, rtol=0)
 
+    def test_tape_no_parameters(self, torch_device, execute_kwargs, tol):
+        """Test that a tape with no parameters is correctly
+        ignored during the gradient computation"""
+        dev = qml.device("default.qubit", wires=1)
+        params = torch.tensor([0.1, 0.2], requires_grad=True, device=torch_device)
+        x, y = params.detach()
+
+        with qml.tape.JacobianTape() as tape1:
+            qml.Hadamard(0)
+            qml.expval(qml.PauliX(0))
+
+        with qml.tape.JacobianTape() as tape2:
+            qml.RY(0.5, wires=0)
+            qml.expval(qml.PauliZ(0))
+
+        with qml.tape.JacobianTape() as tape3:
+            qml.RY(params[0], wires=0)
+            qml.RX(params[1], wires=0)
+            qml.expval(qml.PauliZ(0))
+
+        res = sum(execute([tape1, tape2, tape3], dev, **execute_kwargs))
+        expected = 1 + np.cos(0.5) + np.cos(x) * np.cos(y)
+        assert np.allclose(res.detach(), expected, atol=tol, rtol=0)
+
+        res.backward()
+        grad = params.grad.detach()
+        expected = [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)]
+        assert np.allclose(grad, expected, atol=tol, rtol=0)
+
     def test_reusing_quantum_tape(self, torch_device, execute_kwargs, tol):
         """Test re-using a quantum tape by passing new parameters"""
         a = torch.tensor(0.1, requires_grad=True, device=torch_device)
@@ -760,6 +789,27 @@ class TestTorchExecuteIntegration:
         assert res.shape == (2, 10)
         assert isinstance(res, torch.Tensor)
 
+    def test_sampling_expval(self, torch_device, execute_kwargs):
+        """Test sampling works as expected if combined with expectation values"""
+        if execute_kwargs["gradient_fn"] == "device" and execute_kwargs["mode"] == "forward":
+            pytest.skip("Adjoint differentiation does not support samples")
+
+        dev = qml.device("default.qubit", wires=2, shots=10)
+
+        with qml.tape.JacobianTape() as tape:
+            qml.Hadamard(wires=[0])
+            qml.CNOT(wires=[0, 1])
+            qml.sample(qml.PauliZ(0))
+            qml.expval(qml.PauliX(1))
+
+        res = execute([tape], dev, **execute_kwargs)[0]
+
+        assert len(res) == 2
+        assert isinstance(res, tuple)
+        assert res[0].shape == (10,)
+        assert isinstance(res[0], torch.Tensor)
+        assert isinstance(res[1], torch.Tensor)
+
     def test_sampling_gradient_error(self, torch_device, execute_kwargs):
         """Test differentiating a tape with sampling results in an error"""
         if execute_kwargs["gradient_fn"] == "device" and execute_kwargs["mode"] == "forward":
diff --git a/tests/optimize/test_shot_adaptive.py b/tests/test_optimize_shot_adapative.py
similarity index 100%
rename from tests/optimize/test_shot_adaptive.py
rename to tests/test_optimize_shot_adapative.py
diff --git a/tests/transforms/test_batch_transform.py b/tests/transforms/test_batch_transform.py
index f9498e46a..01c14d506 100644
--- a/tests/transforms/test_batch_transform.py
+++ b/tests/transforms/test_batch_transform.py
@@ -235,6 +235,40 @@ class TestBatchTransform:
         expected = fn(dev.batch_execute(tapes))
         assert res == expected
 
+    def test_custom_qnode_wrapper(self, capsys):
+        """Test that the QNode execution wrapper can be overridden
+        if required."""
+        a = 0.654
+        x = 0.543
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.batch_transform
+        def my_transform(tape, a):
+            tape1 = tape.copy()
+            tape2 = tape.copy()
+            return [tape1, tape2], lambda res: a * qml.math.sum(res)
+
+        @my_transform.custom_qnode_wrapper
+        def qnode_wrapper(self, qnode, targs, tkwargs):
+            wrapper = self.default_qnode_wrapper(qnode, targs, tkwargs)
+            assert targs == (a,)
+            assert tkwargs == {}
+            print("custom wrapper called")
+            return wrapper
+
+        @my_transform(a)
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.Hadamard(wires=0)
+            qml.RX(x, wires=0)
+            return qml.expval(qml.PauliX(0))
+
+        circuit(x)
+
+        captured = capsys.readouterr()
+        assert captured.out == "custom wrapper called\n"
+
 
 @pytest.mark.parametrize("diff_method", ["parameter-shift", "backprop", "finite-diff"])
 class TestBatchTransformGradients: