Repository: PennyLaneAI/pennylane
Pull request #2111 (2022-01-19 08:47:01)

Moves the application of batch transforms out of the QNode and into `qml.execute`
**Context:** Devices can provide device-level batch transformations that might be required to be applied _before_ `qml.execute(tapes)` is called, since they require differentiability.

Currently, the QNode manually performs this logic:

https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/qnode.py#L558-L570

Currently, the default `Device.batch_transform` simply applies `qml.transforms.hamiltonian_expand` if the device doesn't support Hamiltonians, or finite-shots are used.

However, this causes an issue with batch transforms. Batch transforms work by generating multiple tapes, and then manually calling `qml.execute()`, which means that `Device.batch_transform` is skipped - leading to a weird situation where standard QNodes might work with `expval(H)`, but a batch transformed QNode won't (e.g., `qml.gradients.param_shift(qnode)(args)` will then fail if finite-shots are used).

**Description of the Change:**

- The logic for applying `device.batch_transform` is moved from the QNode to inside `qml.execute`. This is much more scalable, and more in-scope.

- However, QNodes currently only ever execute one tape; `qml.execute` can handle multiple tapes. As a result, we must provide a way to support mapping a single batch transform over multiple tapes. To do so, `qml.transforms.map_batch_transform` has been added as a small utility function.

- Tests have been added.

**Benefits:** Batch transforms now no longer ignore device batch transforms.

**Possible Drawbacks:** For some reason, there is a dimensional issue I can't quite solve. To get around this, I had to manually 'squeeze' tuples and lists of length 1.

**Related GitHub Issues:** n/a


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/2111?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2111](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2111?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (d79529a) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/64806dc30fbf3beaf33da5462465de9c1cf99b6c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (64806dc) will **decrease** coverage by `2.21%`.
> The diff coverage is `100.00%`.

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

```diff
@@            Coverage Diff             @@
##           master    #2111      +/-   ##
==========================================
- Coverage   99.19%   96.98%   -2.22%     
==========================================
  Files         228      228              
  Lines       17532    17558      +26     
==========================================
- Hits        17391    17028     -363     
- Misses        141      530     +389     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2111?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19kZXZpY2UucHk=) | `97.97% <100.00%> (ø)` | |
| [pennylane/interfaces/batch/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvX19pbml0X18ucHk=) | `99.11% <100.00%> (-0.89%)` | :arrow_down: |
| [pennylane/qnode.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fub2RlLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/transforms/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/transforms/batch\_transform.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvYmF0Y2hfdHJhbnNmb3JtLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4LnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4LnB5) | `0.00% <0.00%> (-97.44%)` | :arrow_down: |
| [pennylane/devices/default\_qubit\_jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9qYXgucHk=) | `0.00% <0.00%> (-96.63%)` | :arrow_down: |
| [pennylane/qnn/cost.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi9jb3N0LnB5) | `42.85% <0.00%> (-57.15%)` | :arrow_down: |
| [pennylane/math/multi\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2111/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvbXVsdGlfZGlzcGF0Y2gucHk=) | `76.68% <0.00%> (-23.32%)` | :arrow_down: |
| ... and [18 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2111/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2111?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/2111?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [64806dc...d79529a](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2111?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/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 157818680..6ce75d755 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -426,6 +426,10 @@
 
 <h3>Bug fixes</h3>
 
+* Fixes a bug where batch transformed QNodes would fail to apply batch transforms
+  provided by the underlying device.
+  [(#2111)](https://github.com/PennyLaneAI/pennylane/pull/2111)
+
 * An error is raised during QNode creation if backpropagation is requested on a device with
   finite-shots specified.
   [(#2114)](https://github.com/PennyLaneAI/pennylane/pull/2114)
diff --git a/pennylane/interfaces/batch/__init__.py b/pennylane/interfaces/batch/__init__.py
index 82d8b65b8..5d28a9f75 100644
--- a/pennylane/interfaces/batch/__init__.py
+++ b/pennylane/interfaces/batch/__init__.py
@@ -203,6 +203,7 @@ def execute(
     override_shots=False,
     expand_fn="device",
     max_expansion=10,
+    device_batch_transform=True,
 ):
     """Execute a batch of tapes on a device in an autodifferentiable-compatible manner.
 
@@ -239,6 +240,10 @@ def execute(
             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.
+        device_batch_transform (bool): Whether to apply any batch transforms defined by the device
+            (within :meth:`Device.batch_transform`) to each tape to be executed. The default behaviour
+            of the device batch transform is to expand out Hamiltonian measurements into
+            constituent terms if not supported on the device.
 
     Returns:
         list[list[float]]: A nested list of tape results. Each element in
@@ -300,6 +305,11 @@ def execute(
     """
     gradient_kwargs = gradient_kwargs or {}
 
+    if device_batch_transform:
+        tapes, batch_fn = qml.transforms.map_batch_transform(device.batch_transform, tapes)
+    else:
+        batch_fn = lambda res: res
+
     if isinstance(cache, bool) and cache:
         # cache=True: create a LRUCache object
         cache = LRUCache(maxsize=cachesize, getsizeof=lambda x: qml.math.shape(x)[0])
@@ -315,10 +325,12 @@ def execute(
                 tapes
             )
 
-        return res
+        return batch_fn(res)
 
     if gradient_fn == "backprop" or interface is None:
-        return cache_execute(batch_execute, cache, return_tuple=False, expand_fn=expand_fn)(tapes)
+        return batch_fn(
+            cache_execute(batch_execute, cache, return_tuple=False, expand_fn=expand_fn)(tapes)
+        )
 
     # the default execution function is batch_execute
     execute_fn = cache_execute(batch_execute, cache, expand_fn=expand_fn)
@@ -391,4 +403,4 @@ def execute(
         tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_diff=max_diff, mode=_mode
     )
 
-    return res
+    return batch_fn(res)
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 1e17da4fa..57adb1724 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -20,6 +20,8 @@ import functools
 import inspect
 import warnings
 
+import autograd
+
 import pennylane as qml
 from pennylane import Device
 from pennylane.interfaces.batch import set_shots, SUPPORTED_INTERFACES
@@ -554,11 +556,8 @@ class QNode:
         # construct the tape
         self.construct(args, kwargs)
 
-        # preprocess the tapes by applying any device-specific transforms
-        tapes, processing_fn = self.device.batch_transform(self.tape)
-
         res = qml.execute(
-            tapes,
+            [self.tape],
             device=self.device,
             gradient_fn=self.gradient_fn,
             interface=self.interface,
@@ -567,7 +566,19 @@ class QNode:
             **self.execute_kwargs,
         )
 
-        res = processing_fn(res)
+        if autograd.isinstance(res, (tuple, list)) and len(res) == 1:
+            # If a device batch transform was applied, we need to 'unpack'
+            # the returned tuple/list to a float.
+            #
+            # Note that we use autograd.isinstance, because on the backwards pass
+            # with Autograd, lists and tuples are converted to autograd.box.SequenceBox.
+            # autograd.isinstance is a 'safer' isinstance check that supports
+            # autograd backwards passes.
+            #
+            # TODO: find a more explicit way of determining that a batch transform
+            # was applied.
+
+            res = res[0]
 
         if override_shots is not False:
             # restore the initialization gradient function
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index 3b7817589..0f57ea8f6 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -107,6 +107,7 @@ to help build custom QNode, quantum function, and tape transforms:
     ~batch_transform
     ~qfunc_transform
     ~transforms.make_tape
+    ~transforms.map_batch_transform
     ~transforms.create_expand_fn
     ~transforms.create_decomp_expand_fn
     ~transforms.expand_invalid_trainable
@@ -115,7 +116,7 @@ to help build custom QNode, quantum function, and tape transforms:
     ~transforms.expand_nonunitary_gen
 """
 # Import the decorators first to prevent circular imports when used in other transforms
-from .batch_transform import batch_transform
+from .batch_transform import batch_transform, map_batch_transform
 from .qfunc_transforms import make_tape, single_tape_transform, qfunc_transform
 from .adjoint import adjoint
 from .batch_params import batch_params
diff --git a/pennylane/transforms/batch_transform.py b/pennylane/transforms/batch_transform.py
index 5e3f9d48f..7d282160a 100644
--- a/pennylane/transforms/batch_transform.py
+++ b/pennylane/transforms/batch_transform.py
@@ -424,3 +424,70 @@ class batch_transform:
 
     def _tape_wrapper(self, *targs, **tkwargs):
         return lambda tape: self.construct(tape, *targs, **tkwargs)
+
+
+def map_batch_transform(transform, tapes):
+    """Map a batch transform over multiple tapes.
+
+    Args:
+        transform (.batch_transform): the batch transform
+            to be mapped
+        tapes (Sequence[QuantumTape]): The sequence of tapes the batch
+            transform should be applied to. Each tape in the sequence
+            is transformed by the batch transform.
+
+    **Example**
+
+    Consider the following tapes:
+
+    .. code-block:: python
+
+        H = qml.PauliZ(0) @ qml.PauliZ(1) - qml.PauliX(0)
+
+        with qml.tape.QuantumTape() as tape1:
+            qml.RX(0.5, wires=0)
+            qml.RY(0.1, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(H)
+
+        with qml.tape.QuantumTape() as tape2:
+            qml.Hadamard(wires=0)
+            qml.CRX(0.5, wires=[0, 1])
+            qml.CNOT(wires=[0, 1])
+            qml.expval(H + 0.5 * qml.PauliY(0))
+
+    We can use ``map_batch_transform`` to map a single
+    batch transform across both of the these tapes in such a way
+    that allows us to submit a single job for execution:
+
+    >>> tapes, fn = map_batch_transform(qml.transforms.hamiltonian_expand, [tape1, tape2])
+    >>> dev = qml.device("default.qubit", wires=2)
+    >>> fn(qml.execute(tapes, dev, qml.gradients.param_shift))
+    [0.9950041652780257, 0.8150893013179248]
+    """
+    execution_tapes = []
+    batch_fns = []
+    tape_counts = []
+
+    for t in tapes:
+        # Preprocess the tapes by applying batch transforms
+        # to each tape, and storing corresponding tapes
+        # for execution, processing functions, and list of tape lengths.
+        new_tapes, fn = transform(t)
+        execution_tapes.extend(new_tapes)
+        batch_fns.append(fn)
+        tape_counts.append(len(new_tapes))
+
+    def processing_fn(res):
+        count = 0
+        final_results = []
+
+        for idx, s in enumerate(tape_counts):
+            # apply any batch transform post-processing
+            new_res = batch_fns[idx](res[count : count + s])
+            final_results.append(new_res)
+            count += s
+
+        return final_results
+
+    return execution_tapes, processing_fn


Test patch:
diff --git a/tests/interfaces/test_batch_autograd.py b/tests/interfaces/test_batch_autograd.py
index a8ad8083c..e71a4489e 100644
--- a/tests/interfaces/test_batch_autograd.py
+++ b/tests/interfaces/test_batch_autograd.py
@@ -167,6 +167,36 @@ class TestAutogradExecuteUnitTests:
         spy_gradients.assert_called()
 
 
+class TestBatchTransformExecution:
+    """Tests to ensure batch transforms can be correctly executed
+    via qml.execute and map_batch_transform"""
+
+    def test_no_batch_transform(self, mocker):
+        """Test that batch transforms can be disabled and enabled"""
+        dev = qml.device("default.qubit", wires=2, shots=100000)
+
+        H = qml.PauliZ(0) @ qml.PauliZ(1) - qml.PauliX(0)
+        x = 0.6
+        y = 0.2
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(x, wires=0)
+            qml.RY(y, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(H)
+
+        spy = mocker.spy(dev, "batch_transform")
+
+        with pytest.raises(AssertionError, match="Hamiltonian must be used with shots=None"):
+            qml.execute([tape], dev, None, device_batch_transform=False)
+
+        spy.assert_not_called()
+
+        res = qml.execute([tape], dev, None, device_batch_transform=True)
+        spy.assert_called()
+        assert np.allclose(res[0], np.cos(y), atol=0.1)
+
+
 class TestCaching:
     """Test for caching behaviour"""
 
diff --git a/tests/transforms/test_batch_transform.py b/tests/transforms/test_batch_transform.py
index 1fb9e4610..24edbcae1 100644
--- a/tests/transforms/test_batch_transform.py
+++ b/tests/transforms/test_batch_transform.py
@@ -605,3 +605,99 @@ class TestBatchTransformGradients:
         grad = jax.grad(cost, argnums=[0, 1])(x, weights)
         expected = qml.grad(self.expval)(np.array(x), np.array(weights))
         assert all(np.allclose(g, e) for g, e in zip(grad, expected))
+
+    def test_batch_transforms_qnode(self, diff_method, mocker):
+        """Test that batch transforms can be applied to a QNode
+        without affecting device batch transforms"""
+        if diff_method == "backprop":
+            pytest.skip("Test only supports finite shots")
+
+        dev = qml.device("default.qubit", wires=2, shots=100000)
+
+        H = qml.PauliZ(0) @ qml.PauliZ(1) - qml.PauliX(0)
+        weights = np.array([0.5, 0.3], requires_grad=True)
+
+        @qml.gradients.param_shift
+        @qml.qnode(dev, diff_method=diff_method)
+        def circuit(weights):
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(H)
+
+        spy = mocker.spy(dev, "batch_transform")
+
+        res = circuit(weights)
+        spy.assert_called()
+        assert np.allclose(res, [0, -np.sin(weights[1])], atol=0.1)
+
+
+class TestMapBatchTransform:
+    """Tests for the map_batch_transform function"""
+
+    def test_result(self, mocker):
+        """Test that it correctly applies the transform to be mapped"""
+        dev = qml.device("default.qubit", wires=2)
+        H = qml.PauliZ(0) @ qml.PauliZ(1) - qml.PauliX(0)
+        x = 0.6
+        y = 0.7
+
+        with qml.tape.JacobianTape() as tape1:
+            qml.RX(x, wires=0)
+            qml.RY(y, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(H)
+
+        with qml.tape.JacobianTape() as tape2:
+            qml.Hadamard(wires=0)
+            qml.CRX(x, wires=[0, 1])
+            qml.CNOT(wires=[0, 1])
+            qml.expval(H + 0.5 * qml.PauliY(0))
+
+        spy = mocker.spy(qml.transforms, "hamiltonian_expand")
+        tapes, fn = qml.transforms.map_batch_transform(
+            qml.transforms.hamiltonian_expand, [tape1, tape2]
+        )
+
+        spy.assert_called()
+        assert len(tapes) == 5
+
+        res = qml.execute(tapes, dev, qml.gradients.param_shift, device_batch_transform=False)
+        expected = [np.cos(y), 0.5 + 0.5 * np.cos(x) - 0.5 * np.sin(x / 2)]
+
+        assert np.allclose(fn(res), expected)
+
+    def test_differentiation(self):
+        """Test that an execution using map_batch_transform can be differentiated"""
+        dev = qml.device("default.qubit", wires=2)
+        H = qml.PauliZ(0) @ qml.PauliZ(1) - qml.PauliX(0)
+
+        weights = np.array([0.6, 0.8], requires_grad=True)
+
+        def cost(weights):
+            with qml.tape.JacobianTape() as tape1:
+                qml.RX(weights[0], wires=0)
+                qml.RY(weights[1], wires=1)
+                qml.CNOT(wires=[0, 1])
+                qml.expval(H)
+
+            with qml.tape.JacobianTape() as tape2:
+                qml.Hadamard(wires=0)
+                qml.CRX(weights[0], wires=[0, 1])
+                qml.CNOT(wires=[0, 1])
+                qml.expval(H + 0.5 * qml.PauliY(0))
+
+            tapes, fn = qml.transforms.map_batch_transform(
+                qml.transforms.hamiltonian_expand, [tape1, tape2]
+            )
+            res = qml.execute(tapes, dev, qml.gradients.param_shift, device_batch_transform=False)
+            return np.sum(fn(res))
+
+        res = cost(weights)
+        x, y = weights
+        expected = np.cos(y) + 0.5 + 0.5 * np.cos(x) - 0.5 * np.sin(x / 2)
+        assert np.allclose(res, expected)
+
+        res = qml.grad(cost)(weights)
+        expected = [-0.5 * np.sin(x) - 0.25 * np.cos(x / 2), -np.sin(y)]
+        assert np.allclose(res, expected)