Repository: PennyLaneAI/pennylane
Pull request #1637 (2021-09-10 07:42:21)

Bug fixes for batch execution
**Context:** As part of the WIP PR #1623, several bugs and issues relating to the new batch-execution pipeline have been identified. The corresponding bug fixes have been extracted into a separate PR, for ease of review.

**Description of the Change:**

- I discovered that Autograd and TF have a different convention for the ordering of the Jacobian! Without any changes, this results in `qml.jacobian(qnode)` returning the _transpose_ of the expected Jacobian when using Autograd and TF (PyTorch is unaffected). To fix this, the gradient-output vectors `dy` are transposed in the Autograd/TF interfaces. The corresponding tests have also been updated.

- A bug was present when batch executing a list of tapes that may contain tapes with _no_ parameters. To fix this, an `if` statement was added to the `vjp` function. Corresponding tests to check for this behaviour have also been added to all three interface tests.

- Inside the interfaces, `with qml.tape.Unwrap(*tapes)` appeared to not be correctly determining trainable parameters when used recursively (that is, it would work on the first call correctly, and then fail on subsequent nested calls). To fix this, all instances are now being called with `set_trainable=False`.

- The torch interface was not correctly casting output of tapes that had a mixture of sampling and expvals.

- Finally, a small change was made to `qml.batch_execute` to allow transforms to override the default QNode execution logic (if needed). This need arises in the new QNode.

**Benefits:** Various bug fixes and new tests to ensure these are caught in future.

**Possible Drawbacks:** n/a

**Related GitHub Issues:** n/a


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [.github/CHANGELOG.md](/PennyLaneAI/pennylane/blob/master/.github/CHANGELOG.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1637?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1637](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1637?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (a937d4d) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/e702624038ba11bbe32cb75ea62d99eb09e5bef2?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (e702624) will **decrease** coverage by `0.02%`.
> The diff coverage is `85.18%`.

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

```diff
@@            Coverage Diff             @@
##           master    #1637      +/-   ##
==========================================
- Coverage   99.14%   99.11%   -0.03%     
==========================================
  Files         196      196              
  Lines       14241    14253      +12     
==========================================
+ Hits        14119    14127       +8     
- Misses        122      126       +4     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1637?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/interfaces/batch/torch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1637/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvdG9yY2gucHk=) | `95.08% <57.14%> (-4.92%)` | :arrow_down: |
| [pennylane/transforms/batch\_transform.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1637/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvYmF0Y2hfdHJhbnNmb3JtLnB5) | `98.11% <88.88%> (-1.89%)` | :arrow_down: |
| [pennylane/gradients/gradient\_transform.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1637/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9ncmFkaWVudF90cmFuc2Zvcm0ucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/gradients/vjp.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1637/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy92anAucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/batch/autograd.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1637/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvYXV0b2dyYWQucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/batch/tensorflow.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1637/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvdGVuc29yZmxvdy5weQ==) | `100.00% <100.00%> (ø)` | |

------

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