Repository: PennyLaneAI/pennylane
Pull request #5212 (2024-02-15 02:36:54)

Replacing `map_batch_transform` in source code
**Context:** In the PennyLane source code, the `map_batch_transform` method is currently used in several locations. However, a new functionality has recently been added to `TransformDispatcher` so that transforms can be dispatched onto a batch of tapes (see #5163).

**Description of the Change:** We replace the current usage of `map_batch_transform` in the source code, taking advantage of the new method `_batch_transform` implemented in the `TransformDispatcher` class.

**Benefits:** Implementing transforms over a batch of tapes without resorting to `map_batch_transform` is much more straightforward.

**Possible Drawbacks:** The new method in `TransformDispatcher` replacing `map_batch_transform` is called only if the transform is applied on a sequence consisting of `qml.tape.QuantumScript` instances. Finally, compared to `map_batch_transform`, the new method returns a `tuple` as a first variable (instead of a `list`).

**Related GitHub Issues:** None

[sc-56689]

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://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5212?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`256b1cc`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/256b1ccefe2420d9951ed276b16b230502d6f33a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.68% compared to head [(`3ba0b7a`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5212?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.68%.


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


```diff
@@            Coverage Diff             @@
##           master    #5212      +/-   ##
==========================================
- Coverage   99.68%   99.68%   -0.01%     
==========================================
  Files         396      396              
  Lines       36640    36367     -273     
==========================================
- Hits        36526    36252     -274     
- Misses        114      115       +1     
```



</details>

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


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 0bf8c407c..e6585fb31 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -373,6 +373,10 @@
 
 <h3>Deprecations 👋</h3>
 
+* Replacing `map_batch_transform` in the source code with the method `_batch_transform` 
+  implemented in `TransformDispatcher`.
+  [(#5212)](https://github.com/PennyLaneAI/pennylane/pull/5212)
+
 * `TransformDispatcher` can now dispatch onto a batch of tapes, so that it is easier to compose transforms
   when working in the tape paradigm.
   [(#5163)](https://github.com/PennyLaneAI/pennylane/pull/5163)
diff --git a/pennylane/_device.py b/pennylane/_device.py
index aaab968cf..3553e2b8c 100644
--- a/pennylane/_device.py
+++ b/pennylane/_device.py
@@ -803,9 +803,7 @@ class Device(abc.ABC):
             return circuits, hamiltonian_fn
 
         # Expand each of the broadcasted Hamiltonian-expanded circuits
-        expanded_tapes, expanded_fn = qml.transforms.map_batch_transform(
-            qml.transforms.broadcast_expand, circuits
-        )
+        expanded_tapes, expanded_fn = qml.transforms.broadcast_expand(circuits)
 
         # Chain the postprocessing functions of the broadcasted-tape expansions and the Hamiltonian
         # expansion. Note that the application order is reversed compared to the expansion order,
diff --git a/pennylane/transforms/mitigate.py b/pennylane/transforms/mitigate.py
index 9b2c1d847..51af815e8 100644
--- a/pennylane/transforms/mitigate.py
+++ b/pennylane/transforms/mitigate.py
@@ -525,6 +525,10 @@ def mitigate_with_zne(
 
     def processing_fn(results):
         """Maps from input tape executions to an error-mitigated estimate"""
+
+        # content of `results` must be modified in this post-processing function
+        results = list(results)
+
         for i, tape in enumerate(out_tapes):
             # stack the results if there are multiple measurements
             # this will not create ragged arrays since only expval measurements are allowed
diff --git a/pennylane/workflow/execution.py b/pennylane/workflow/execution.py
index 9884f3ed3..b369059dc 100644
--- a/pennylane/workflow/execution.py
+++ b/pennylane/workflow/execution.py
@@ -196,8 +196,10 @@ def _batch_transform(
     """
     # TODO: Remove once old device are removed
     if device_batch_transform:
-        dev_batch_transform = set_shots(device, override_shots)(device.batch_transform)
-        return *qml.transforms.map_batch_transform(dev_batch_transform, tapes), config
+        dev_batch_transform = qml.transform(
+            set_shots(device, override_shots)(device.batch_transform)
+        )
+        return *dev_batch_transform(tapes), config
 
     def null_post_processing_fn(results):
         """A null post processing function used because the user requested not to use the device batch transform."""
diff --git a/pennylane/workflow/jacobian_products.py b/pennylane/workflow/jacobian_products.py
index 91bf08e70..e902ed2ad 100644
--- a/pennylane/workflow/jacobian_products.py
+++ b/pennylane/workflow/jacobian_products.py
@@ -15,7 +15,6 @@
 Defines classes that take the vjps, jvps, and jacobians of circuits.
 """
 import abc
-from functools import partial
 import inspect
 import logging
 from typing import Tuple, Callable, Optional, Union
@@ -310,10 +309,7 @@ class TransformJacobianProducts(JacobianProductCalculator):
 
         num_result_tapes = len(tapes)
 
-        partial_gradient_fn = partial(self._gradient_transform, **self._gradient_kwargs)
-        jac_tapes, jac_postprocessing = qml.transforms.map_batch_transform(
-            partial_gradient_fn, tapes
-        )
+        jac_tapes, jac_postprocessing = self._gradient_transform(tapes, **self._gradient_kwargs)
 
         full_batch = tapes + tuple(jac_tapes)
         full_results = self._inner_execute(full_batch)
@@ -327,10 +323,7 @@ class TransformJacobianProducts(JacobianProductCalculator):
             logger.debug("compute_jacobian called with %s", tapes)
         if tapes in self._cache:
             return self._cache[tapes]
-        partial_gradient_fn = partial(self._gradient_transform, **self._gradient_kwargs)
-        jac_tapes, batch_post_processing = qml.transforms.map_batch_transform(
-            partial_gradient_fn, tapes
-        )
+        jac_tapes, batch_post_processing = self._gradient_transform(tapes, **self._gradient_kwargs)
         results = self._inner_execute(jac_tapes)
         jacs = tuple(batch_post_processing(results))
         self._cache[tapes] = jacs


Test patch:
diff --git a/tests/interfaces/test_autograd.py b/tests/interfaces/test_autograd.py
index cd2ef932d..f914c902a 100644
--- a/tests/interfaces/test_autograd.py
+++ b/tests/interfaces/test_autograd.py
@@ -219,7 +219,7 @@ class TestBatchTransformExecution:
         H = 2.0 * qml.PauliZ(0)
         qscript = qml.tape.QuantumScript(measurements=[qml.expval(H)])
         res = qml.execute([qscript], dev, interface=None, override_shots=10)
-        assert res == [2.0]
+        assert res == (2.0,)
 
 
 class TestCaching:
diff --git a/tests/interfaces/test_jax.py b/tests/interfaces/test_jax.py
index ad762b9a7..44801ce4c 100644
--- a/tests/interfaces/test_jax.py
+++ b/tests/interfaces/test_jax.py
@@ -18,6 +18,7 @@ import numpy as np
 
 import pennylane as qml
 from pennylane.gradients import param_shift
+from pennylane.typing import TensorLike
 from pennylane import execute
 
 pytestmark = pytest.mark.jax
@@ -576,7 +577,7 @@ class TestJaxExecuteIntegration:
             return execute(tapes=[tape1, tape2], device=dev, **execute_kwargs)
 
         res = cost_fn(params)
-        assert isinstance(res, list)
+        assert isinstance(res, TensorLike)
         assert all(isinstance(r, jax.numpy.ndarray) for r in res)
         assert all(r.shape == () for r in res)
 
@@ -856,7 +857,7 @@ class TestVectorValued:
             x, y, dev, interface="jax-python", ek=execute_kwargs
         )
 
-        assert isinstance(res, list)
+        assert isinstance(res, TensorLike)
         assert len(res) == 2
 
         for r, exp_shape in zip(res, [(), (2,)]):
diff --git a/tests/interfaces/test_jax_jit.py b/tests/interfaces/test_jax_jit.py
index 495fb24fa..509b537f1 100644
--- a/tests/interfaces/test_jax_jit.py
+++ b/tests/interfaces/test_jax_jit.py
@@ -19,6 +19,7 @@ import numpy as np
 
 import pennylane as qml
 from pennylane.gradients import param_shift
+from pennylane.typing import TensorLike
 from pennylane import execute
 
 pytestmark = pytest.mark.jax
@@ -580,7 +581,7 @@ class TestJaxExecuteIntegration:
             return execute(tapes=[tape1, tape2], device=dev, **execute_kwargs)
 
         res = jax.jit(cost_fn)(params)
-        assert isinstance(res, list)
+        assert isinstance(res, TensorLike)
         assert all(isinstance(r, jax.numpy.ndarray) for r in res)
         assert all(r.shape == () for r in res)
 
diff --git a/tests/interfaces/test_torch.py b/tests/interfaces/test_torch.py
index 4c519360d..71e2145d0 100644
--- a/tests/interfaces/test_torch.py
+++ b/tests/interfaces/test_torch.py
@@ -18,6 +18,7 @@ import pytest
 
 import pennylane as qml
 from pennylane.gradients import finite_diff, param_shift
+from pennylane.typing import TensorLike
 from pennylane import execute
 
 pytestmark = pytest.mark.torch
@@ -421,7 +422,7 @@ class TestTorchExecuteIntegration:
 
         res = execute([tape1, tape2], dev, **execute_kwargs)
 
-        assert isinstance(res, list)
+        assert isinstance(res, TensorLike)
         assert len(res) == 2
         assert res[0].shape == ()
         assert res[1].shape == ()