Repository: PennyLaneAI/pennylane
Pull request #4675 (2023-10-13 17:23:30)

Defining custom decompositions for DQ2
**Context:**
The `custom_decomps` and `decomp_depth` arguments when calling `qml.device("default.qubit")` did not work, and as such were not at feature parity with DQL. The `set_decomposition` context also didn't work with the new implementation of `default.qubit`

**Description of the Change:**
Kwargs `decomposer` and `max_expansion` are added to the `preprocessing.decompose` transform. These can be specified to create custom behaviour for decompositions.

A `custom_decomp_preprocessing` function is created, analogous to `create_decomp_expand_fn` for DQL. It defines a custom `decomposer` for `device.decompose` and modifies device `stopping_conditions` based on the `custom_decomps` dictionary provided. It also modifies `device.decompose.max_expansion` if a `decomp_depth` is specified.

There's some special treatment for controlled operators where we override the `op.base.compute_decomposition` method for that instance of the operator if there is a custom decomposition defined for the base.

The tests are all copy-pasted from the old device API tests, with slight modifications where relevant for the new API, and a couple of additional tests that seemed relevant.

**Benefits:**
Feature parity with DQL wrt custom decompositions. It works.

**Possible Drawbacks:**
It feels messy and not very nice (but it does the thing)


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/4675?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified lines are covered by tests :white_check_mark:
> Comparison is base [(`4c3ab72`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/4c3ab729c0715c0a78a0e9b3454a650f0e7e7096?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64% compared to head [(`00b68cc`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4675?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64%.
> Report is 1 commits behind head on master.

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


```diff
@@           Coverage Diff           @@
##           master    #4675   +/-   ##
=======================================
  Coverage   99.64%   99.64%           
=======================================
  Files         376      376           
  Lines       33828    33876   +48     
=======================================
+ Hits        33707    33755   +48     
  Misses        121      121           
```


| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4675?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_\_init\_\_.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4675?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/devices/preprocess.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4675?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcHJlcHJvY2Vzcy5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/drawer/draw.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4675?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RyYXdlci9kcmF3LnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/transforms/tape\_expand.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4675?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvdGFwZV9leHBhbmQucHk=) | `100.00% <100.00%> (ø)` | |

... and [1 file with indirect coverage changes](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4675/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4675?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 60f265ff6..b93593862 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -257,6 +257,13 @@
   of the computed qubit operator, if imaginary components are smaller than a threshold. 
   [(#4639)](https://github.com/PennyLaneAI/pennylane/pull/4639)
 
+* Updated `qml.device`, `devices.preprocessing` and the `tape_expand.set_decomposition` context 
+  manager to bring `DefaultQubit2` to feature parity with `default.qubit.legacy` with regards to 
+  using custom decompositions. The `DefaultQubit2` device can now be included in a `set_decomposition` 
+  context or initialized with a `custom_decomps` dictionary, as well as a custom `max_depth` for 
+  decomposition.
+  [(#4675)](https://github.com/PennyLaneAI/pennylane/pull/4675)
+
 
 <h3>Breaking changes 💔</h3>
 
diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index 252542704..0d869bf99 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -163,6 +163,7 @@ def refresh_devices():
 plugin_devices = _get_device_entrypoints()
 
 
+# pylint: disable=protected-access
 def device(name, *args, **kwargs):
     r"""device(name, wires=1, *args, **kwargs)
     Load a :class:`~.Device` and return the instance.
@@ -348,10 +349,18 @@ def device(name, *args, **kwargs):
         # Once the device is constructed, we set its custom expansion function if
         # any custom decompositions were specified.
         if custom_decomps is not None:
-            custom_decomp_expand_fn = pennylane.transforms.create_decomp_expand_fn(
-                custom_decomps, dev, decomp_depth=decomp_depth
-            )
-            dev.custom_expand(custom_decomp_expand_fn)
+            if isinstance(dev, pennylane.Device):
+                custom_decomp_expand_fn = pennylane.transforms.create_decomp_expand_fn(
+                    custom_decomps, dev, decomp_depth=decomp_depth
+                )
+                dev.custom_expand(custom_decomp_expand_fn)
+            else:
+                custom_decomp_preprocess = (
+                    pennylane.transforms.tape_expand._create_decomp_preprocessing(
+                        custom_decomps, dev, decomp_depth=decomp_depth
+                    )
+                )
+                dev.preprocess = custom_decomp_preprocess
 
         return dev
 
diff --git a/pennylane/devices/preprocess.py b/pennylane/devices/preprocess.py
index b2ee41641..7444e96bf 100644
--- a/pennylane/devices/preprocess.py
+++ b/pennylane/devices/preprocess.py
@@ -14,7 +14,8 @@
 
 """This module contains functions for preprocessing `QuantumTape` objects to ensure
 that they are supported for execution by a device."""
-# pylint: disable=protected-access
+# pylint: disable=protected-access, too-many-arguments
+
 import os
 from typing import Generator, Callable, Union, Sequence, Optional
 from copy import copy
@@ -45,21 +46,35 @@ def null_postprocessing(results):
 def _operator_decomposition_gen(
     op: qml.operation.Operator,
     acceptance_function: Callable[[qml.operation.Operator], bool],
+    decomposer: Callable[[qml.operation.Operator], Sequence[qml.operation.Operator]],
+    max_expansion: Optional[int] = None,
+    current_depth=0,
     name: str = "device",
 ) -> Generator[qml.operation.Operator, None, None]:
     """A generator that yields the next operation that is accepted."""
-    if acceptance_function(op):
+    max_depth_reached = False
+    if max_expansion is not None and max_expansion <= current_depth:
+        max_depth_reached = True
+    if acceptance_function(op) or max_depth_reached:
         yield op
     else:
         try:
-            decomp = op.decomposition()
+            decomp = decomposer(op)
+            current_depth += 1
         except qml.operation.DecompositionUndefinedError as e:
             raise DeviceError(
                 f"Operator {op} not supported on {name} and does not provide a decomposition."
             ) from e
 
         for sub_op in decomp:
-            yield from _operator_decomposition_gen(sub_op, acceptance_function, name)
+            yield from _operator_decomposition_gen(
+                sub_op,
+                acceptance_function,
+                decomposer=decomposer,
+                max_expansion=max_expansion,
+                current_depth=current_depth,
+                name=name,
+            )
 
 
 #######################
@@ -214,6 +229,10 @@ def decompose(
     tape: qml.tape.QuantumTape,
     stopping_condition: Callable[[qml.operation.Operator], bool],
     skip_initial_state_prep: bool = True,
+    decomposer: Optional[
+        Callable[[qml.operation.Operator], Sequence[qml.operation.Operator]]
+    ] = None,
+    max_expansion: Union[int, None] = None,
     name: str = "device",
 ) -> (Sequence[qml.tape.QuantumTape], Callable):
     """Decompose operations until the stopping condition is met.
@@ -224,6 +243,10 @@ def decompose(
             should be decomposed. If an operator cannot be decomposed and is not accepted by ``stopping_condition``,
             a ``DecompositionUndefinedError`` will be raised.
         skip_initial_state_prep=True (bool): If ``True``, the first operator will not be decomposed if it inherits from :class:`~.StatePrepBase`.
+        decomposer (Callable): an optional callable that takes an operator and implements the relevant decomposition.
+            If None, defaults to using a callable returning ``op.decomposition()`` for any :class:`~.Operator` .
+        max_expansion (int): The maximum depth of the expansion.
+
 
     Returns:
         pennylane.QNode or qfunc or Tuple[List[.QuantumTape], Callable]: If a QNode is passed,
@@ -274,6 +297,10 @@ def decompose(
     RZ(1.5707963267948966, wires=[1])]
 
     """
+    if decomposer is None:
+
+        def decomposer(op):
+            return op.decomposition()
 
     if not all(stopping_condition(op) for op in tape.operations):
         try:
@@ -285,7 +312,13 @@ def decompose(
             new_ops = [
                 final_op
                 for op in tape.operations[bool(prep_op) :]
-                for final_op in _operator_decomposition_gen(op, stopping_condition, name)
+                for final_op in _operator_decomposition_gen(
+                    op,
+                    stopping_condition,
+                    decomposer=decomposer,
+                    max_expansion=max_expansion,
+                    name=name,
+                )
             ]
         except RecursionError as e:
             raise DeviceError(
diff --git a/pennylane/drawer/draw.py b/pennylane/drawer/draw.py
index ce3939e28..b07721289 100644
--- a/pennylane/drawer/draw.py
+++ b/pennylane/drawer/draw.py
@@ -253,7 +253,9 @@ def _draw_qnode(
 ):
     @wraps(qnode)
     def wrapper(*args, **kwargs):
-        if expansion_strategy == "device" and isinstance(qnode.device, qml.devices.Device):
+        if isinstance(qnode.device, qml.devices.Device) and (
+            expansion_strategy == "device" or getattr(qnode, "expansion_strategy", None) == "device"
+        ):
             qnode.construct(args, kwargs)
             program, _ = qnode.device.preprocess()
             tapes = program([qnode.tape])
diff --git a/pennylane/transforms/tape_expand.py b/pennylane/transforms/tape_expand.py
index fc96b1be8..8b418191f 100644
--- a/pennylane/transforms/tape_expand.py
+++ b/pennylane/transforms/tape_expand.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 """This module contains tape expansion functions and stopping criteria to
 generate such functions from."""
-# pylint: disable=unused-argument,invalid-unary-operand-type, unsupported-binary-operation
+# pylint: disable=unused-argument,invalid-unary-operand-type, unsupported-binary-operation, no-member
 import contextlib
 
 import pennylane as qml
@@ -278,7 +278,7 @@ def _custom_decomp_context(custom_decomps):
             yield
 
         finally:
-            obj.compute_decomposition = original_decomp_method
+            obj.compute_decomposition = staticmethod(original_decomp_method)
 
     # Loop through the decomposition dictionary and create all the contexts
     try:
@@ -350,6 +350,80 @@ def create_decomp_expand_fn(custom_decomps, dev, decomp_depth=10):
     return custom_decomp_expand
 
 
+def _create_decomp_preprocessing(custom_decomps, dev, decomp_depth=10):
+    """Creates a custom preprocessing method for a device that applies
+    a set of specified custom decompositions.
+
+    Args:
+        custom_decomps (Dict[Union(str, qml.operation.Operation), Callable]): Custom
+            decompositions to be applied by the device at runtime.
+        dev (pennylane.devices.Device): A quantum device.
+        decomp_depth: The maximum depth of the expansion.
+
+    Returns:
+        Callable: A custom preprocessing method that a device can call to expand
+        its tapes.
+
+    **Example**
+
+    Suppose we would like a custom expansion function that decomposes all CNOTs
+    into CZs. We first define a decomposition function:
+
+    .. code-block:: python
+
+        def custom_cnot(wires):
+            return [
+                qml.Hadamard(wires=wires[1]),
+                qml.CZ(wires=[wires[0], wires[1]]),
+                qml.Hadamard(wires=wires[1])
+            ]
+
+    We then create the custom function (passing a device, in order to pick up any
+    additional stopping criteria the expansion should have), and then register the
+    result as a custom function of the device:
+
+    >>> custom_decomps = {qml.CNOT : custom_cnot}
+    >>> new_preprocessing = _create_decomp_preprocessing(custom_decomps, dev)
+    >>> dev.preprocess = new_preprocessing
+    """
+
+    def decomposer(op):
+        if isinstance(op, qml.ops.Controlled) and type(op.base) in custom_decomps:
+            op.base.compute_decomposition = custom_decomps[type(op.base)]
+            return op.decomposition()
+        if op.name in custom_decomps:
+            return custom_decomps[op.name](*op.data, wires=op.wires, **op.hyperparameters)
+        if type(op) in custom_decomps:
+            return custom_decomps[type(op)](*op.data, wires=op.wires, **op.hyperparameters)
+        return op.decomposition()
+
+    original_preprocess = dev.preprocess
+
+    # pylint: disable=cell-var-from-loop
+    def new_preprocess(execution_config=qml.devices.DefaultExecutionConfig):
+        program, config = original_preprocess(execution_config)
+
+        for container in program:
+            if container.transform == qml.devices.preprocess.decompose.transform:
+                container.kwargs["decomposer"] = decomposer
+                container.kwargs["max_expansion"] = decomp_depth
+
+                original_stopping_condition = container.kwargs["stopping_condition"]
+
+                def stopping_condition(obj):
+                    if obj.name in custom_decomps or type(obj) in custom_decomps:
+                        return False
+                    return original_stopping_condition(obj)
+
+                container.kwargs["stopping_condition"] = stopping_condition
+
+                break
+
+        return program, config
+
+    return new_preprocess
+
+
 @contextlib.contextmanager
 def set_decomposition(custom_decomps, dev, decomp_depth=10):
     """Context manager for setting custom decompositions.
@@ -398,18 +472,32 @@ def set_decomposition(custom_decomps, dev, decomp_depth=10):
     1: ──H─╰Z──H─┤
 
     """
-    original_custom_expand_fn = dev.custom_expand_fn
+    if isinstance(dev, qml.Device):
+        original_custom_expand_fn = dev.custom_expand_fn
 
-    # Create a new expansion function; stop at things that do not have
-    # custom decompositions, or that satisfy the regular device stopping criteria
-    new_custom_expand_fn = qml.transforms.create_decomp_expand_fn(
-        custom_decomps, dev, decomp_depth=decomp_depth
-    )
+        # Create a new expansion function; stop at things that do not have
+        # custom decompositions, or that satisfy the regular device stopping criteria
+        new_custom_expand_fn = create_decomp_expand_fn(
+            custom_decomps, dev, decomp_depth=decomp_depth
+        )
 
-    # Set the custom expand function within this context only
-    try:
-        dev.custom_expand(new_custom_expand_fn)
-        yield
+        # Set the custom expand function within this context only
+        try:
+            dev.custom_expand(new_custom_expand_fn)
+            yield
 
-    finally:
-        dev.custom_expand_fn = original_custom_expand_fn
+        finally:
+            dev.custom_expand_fn = original_custom_expand_fn
+
+    else:
+        original_preprocess = dev.preprocess
+        new_preprocess = _create_decomp_preprocessing(
+            custom_decomps, dev, decomp_depth=decomp_depth
+        )
+
+        try:
+            dev.preprocess = new_preprocess
+            yield
+
+        finally:
+            dev.preprocess = original_preprocess


Test patch:
diff --git a/tests/devices/test_preprocess.py b/tests/devices/test_preprocess.py
index 0a4d58fff..006db9f4a 100644
--- a/tests/devices/test_preprocess.py
+++ b/tests/devices/test_preprocess.py
@@ -59,6 +59,10 @@ class NoMatNoDecompOp(Operation):
 class TestPrivateHelpers:
     """Test the private helpers for preprocessing."""
 
+    @staticmethod
+    def decomposer(op):
+        return op.decomposition()
+
     @pytest.mark.parametrize("op", (qml.PauliX(0), qml.RX(1.2, wires=0), qml.QFT(wires=range(3))))
     def test_operator_decomposition_gen_accepted_operator(self, op):
         """Test the _operator_decomposition_gen function on an operator that is accepted."""
@@ -66,7 +70,7 @@ class TestPrivateHelpers:
         def stopping_condition(op):
             return op.has_matrix
 
-        casted_to_list = list(_operator_decomposition_gen(op, stopping_condition))
+        casted_to_list = list(_operator_decomposition_gen(op, stopping_condition, self.decomposer))
         assert len(casted_to_list) == 1
         assert casted_to_list[0] is op
 
@@ -78,7 +82,7 @@ class TestPrivateHelpers:
             return op.has_matrix
 
         op = NoMatOp("a")
-        casted_to_list = list(_operator_decomposition_gen(op, stopping_condition))
+        casted_to_list = list(_operator_decomposition_gen(op, stopping_condition, self.decomposer))
         assert len(casted_to_list) == 2
         assert qml.equal(casted_to_list[0], qml.PauliX("a"))
         assert qml.equal(casted_to_list[1], qml.PauliY("a"))
@@ -98,7 +102,7 @@ class TestPrivateHelpers:
                 return [NoMatOp(self.wires), qml.S(self.wires), qml.adjoint(NoMatOp(self.wires))]
 
         op = RaggedDecompositionOp("a")
-        final_decomp = list(_operator_decomposition_gen(op, stopping_condition))
+        final_decomp = list(_operator_decomposition_gen(op, stopping_condition, self.decomposer))
         assert len(final_decomp) == 5
         assert qml.equal(final_decomp[0], qml.PauliX("a"))
         assert qml.equal(final_decomp[1], qml.PauliY("a"))
@@ -113,7 +117,11 @@ class TestPrivateHelpers:
             DeviceError,
             match=r"not supported on abc and does",
         ):
-            tuple(_operator_decomposition_gen(op, lambda op: op.has_matrix, name="abc"))
+            tuple(
+                _operator_decomposition_gen(
+                    op, lambda op: op.has_matrix, self.decomposer, name="abc"
+                )
+            )
 
 
 def test_no_sampling():
diff --git a/tests/transforms/test_tape_expand.py b/tests/transforms/test_tape_expand.py
index 82f5ecdfb..e8eceb892 100644
--- a/tests/transforms/test_tape_expand.py
+++ b/tests/transforms/test_tape_expand.py
@@ -16,8 +16,10 @@ Unit tests for tape expansion stopping criteria and expansion functions.
 """
 # pylint: disable=too-few-public-methods, invalid-unary-operand-type, no-member,
 # pylint: disable=arguments-differ, arguments-renamed,
+
 import pytest
 import numpy as np
+
 import pennylane as qml
 from pennylane.wires import Wires
 
@@ -422,9 +424,31 @@ def custom_basic_entangler_layers(weights, wires, **kwargs):
 
 
 class TestCreateCustomDecompExpandFn:
-    """Tests for the gradient expand function"""
+    """Tests for the custom_decomps argument for devices"""
+
+    @pytest.mark.parametrize("device_name", ["default.qubit", "default.qubit.legacy"])
+    def test_string_and_operator_allowed(self, device_name):
+        """Test that the custom_decomps dictionary accepts both strings and operator classes as keys."""
+
+        custom_decomps = {"Hadamard": custom_hadamard, qml.CNOT: custom_cnot}
+        decomp_dev = qml.device(device_name, wires=2, custom_decomps=custom_decomps)
+
+        @qml.qnode(decomp_dev, expansion_strategy="device")
+        def circuit():
+            qml.Hadamard(wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0))
+
+        _ = circuit()
+        decomp_ops = circuit.tape.operations
+
+        assert len(decomp_ops) == 7
+        for op in decomp_ops:
+            assert op.name != "Hadamard"
+            assert op.name != "CNOT"
 
-    def test_no_custom_decomp(self):
+    @pytest.mark.parametrize("device_name", ["default.qubit", "default.qubit.legacy"])
+    def test_no_custom_decomp(self, device_name):
         """Test that sending an empty dictionary results in no decompositions."""
 
         def circuit():
@@ -432,8 +456,8 @@ class TestCreateCustomDecompExpandFn:
             qml.CNOT(wires=[0, 1])
             return qml.expval(qml.PauliZ(0))
 
-        original_dev = qml.device("default.qubit.legacy", wires=3)
-        decomp_dev = qml.device("default.qubit.legacy", wires=3, custom_decomps={})
+        original_dev = qml.device(device_name, wires=3)
+        decomp_dev = qml.device(device_name, wires=3, custom_decomps={})
 
         original_qnode = qml.QNode(circuit, original_dev, expansion_strategy="device")
         decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
@@ -449,7 +473,8 @@ class TestCreateCustomDecompExpandFn:
             )
         ]
 
-    def test_no_custom_decomp_template(self):
+    @pytest.mark.parametrize("device_name", ["default.qubit", "default.qubit.legacy"])
+    def test_no_custom_decomp_template(self, device_name):
         """Test that sending an empty dictionary results in no decomposition
         when a template is involved, except the decomposition expected from the device."""
 
@@ -457,8 +482,8 @@ class TestCreateCustomDecompExpandFn:
             qml.BasicEntanglerLayers([[0.1, 0.2]], wires=[0, 1])
             return qml.expval(qml.PauliZ(0))
 
-        original_dev = qml.device("default.qubit.legacy", wires=3)
-        decomp_dev = qml.device("default.qubit.legacy", wires=3, custom_decomps={})
+        original_dev = qml.device(device_name, wires=3)
+        decomp_dev = qml.device(device_name, wires=3, custom_decomps={})
 
         original_qnode = qml.QNode(circuit, original_dev, expansion_strategy="device")
         decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
@@ -474,20 +499,23 @@ class TestCreateCustomDecompExpandFn:
             )
         ]
 
-    @pytest.mark.parametrize("device_name", ["default.qubit.legacy", "lightning.qubit"])
+    @pytest.mark.parametrize(
+        "device_name", ["default.qubit", "default.qubit.legacy", "lightning.qubit"]
+    )
     def test_one_custom_decomp(self, device_name):
         """Test that specifying a single custom decomposition works as expected."""
 
+        custom_decomps = {"Hadamard": custom_hadamard}
+        decomp_dev = qml.device(device_name, wires=2, custom_decomps=custom_decomps)
+
+        @qml.qnode(decomp_dev, expansion_strategy="device")
         def circuit():
             qml.Hadamard(wires=0)
             qml.CNOT(wires=[0, 1])
             return qml.expval(qml.PauliZ(0))
 
-        custom_decomps = {"Hadamard": custom_hadamard}
-        decomp_dev = qml.device(device_name, wires=2, custom_decomps=custom_decomps)
-        decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
-        _ = decomp_qnode()
-        decomp_ops = decomp_qnode.qtape.operations
+        _ = circuit()
+        decomp_ops = circuit.tape.operations
 
         assert len(decomp_ops) == 3
 
@@ -499,27 +527,28 @@ class TestCreateCustomDecompExpandFn:
 
         assert decomp_ops[2].name == "CNOT"
 
-    def test_no_decomp_with_depth_zero(self):
+    @pytest.mark.parametrize("device_name", ["default.qubit", "default.qubit.legacy"])
+    def test_no_decomp_with_depth_zero(self, device_name):
         """Test that specifying a single custom decomposition works as expected."""
 
+        custom_decomps = {"Hadamard": custom_hadamard, "CNOT": custom_cnot}
+        decomp_dev = qml.device(device_name, wires=2, custom_decomps=custom_decomps, decomp_depth=0)
+
+        @qml.qnode(decomp_dev, expansion_strategy="device")
         def circuit():
             qml.Hadamard(wires=0)
             qml.CNOT(wires=[0, 1])
             return qml.expval(qml.PauliZ(0))
 
-        custom_decomps = {"Hadamard": custom_hadamard, "CNOT": custom_cnot}
-        decomp_dev = qml.device(
-            "default.qubit.legacy", wires=2, custom_decomps=custom_decomps, decomp_depth=0
-        )
-        decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
-        _ = decomp_qnode()
-        decomp_ops = decomp_qnode.qtape.operations
+        _ = circuit()
+        decomp_ops = circuit.tape.operations
 
         assert len(decomp_ops) == 2
         assert decomp_ops[0].name == "Hadamard"
         assert decomp_ops[1].name == "CNOT"
 
-    def test_one_custom_decomp_gradient(self):
+    @pytest.mark.parametrize("device_name", ["default.qubit", "default.qubit.legacy"])
+    def test_one_custom_decomp_gradient(self, device_name):
         """Test that gradients are still correctly computed after a decomposition
         that performs transpilation."""
 
@@ -529,8 +558,8 @@ class TestCreateCustomDecompExpandFn:
             qml.Hadamard(wires=0)
             return qml.expval(qml.PauliZ(0))
 
-        original_dev = qml.device("default.qubit.legacy", wires=3)
-        decomp_dev = qml.device("default.qubit.legacy", wires=3, custom_decomps={"Rot": custom_rot})
+        original_dev = qml.device(device_name, wires=3)
+        decomp_dev = qml.device(device_name, wires=3, custom_decomps={"Rot": custom_rot})
 
         original_qnode = qml.QNode(circuit, original_dev, expansion_strategy="device")
         decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
@@ -548,20 +577,22 @@ class TestCreateCustomDecompExpandFn:
         expected_ops = ["Hadamard", "RZ", "PauliX", "RY", "PauliX", "RZ", "Hadamard"]
         assert all(op.name == name for op, name in zip(decomp_qnode.qtape.operations, expected_ops))
 
-    def test_nested_custom_decomp(self):
+    @pytest.mark.parametrize("device_name", ["default.qubit", "default.qubit.legacy"])
+    def test_nested_custom_decomp(self, device_name):
         """Test that specifying two custom decompositions that have interdependence
         works as expected."""
 
+        custom_decomps = {"Hadamard": custom_hadamard, qml.CNOT: custom_cnot}
+        decomp_dev = qml.device(device_name, wires=2, custom_decomps=custom_decomps)
+
+        @qml.qnode(decomp_dev, expansion_strategy="device")
         def circuit():
             qml.Hadamard(wires=0)
             qml.CNOT(wires=[0, 1])
             return qml.expval(qml.PauliZ(0))
 
-        custom_decomps = {"Hadamard": custom_hadamard, qml.CNOT: custom_cnot}
-        decomp_dev = qml.device("default.qubit.legacy", wires=2, custom_decomps=custom_decomps)
-        decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
-        _ = decomp_qnode()
-        decomp_ops = decomp_qnode.qtape.operations
+        _ = circuit()
+        decomp_ops = circuit.tape.operations
 
         assert len(decomp_ops) == 7
 
@@ -585,21 +616,23 @@ class TestCreateCustomDecompExpandFn:
 
         assert decomp_ops[4].name == "CZ"
 
-    def test_nested_custom_decomp_with_template(self):
+    @pytest.mark.parametrize("device_name", ["default.qubit", "default.qubit.legacy"])
+    def test_nested_custom_decomp_with_template(self, device_name):
         """Test that specifying two custom decompositions that have interdependence
         works as expected even when there is a template."""
 
+        custom_decomps = {"Hadamard": custom_hadamard, qml.CNOT: custom_cnot}
+        decomp_dev = qml.device(device_name, wires=2, custom_decomps=custom_decomps)
+
+        @qml.qnode(decomp_dev, expansion_strategy="device")
         def circuit():
             # -RX(0.1)-C- -> -RX(0.1)---C--- -> -RX(0.1)-----------------C----------------
             # -RX(0.2)-X- -> -RX(0.2)-H-Z-H- -> -RX(0.2)-RZ(pi)-RY(pi/2)-Z-RY(pi/2)-RZ(pi)-
             qml.BasicEntanglerLayers([[0.1, 0.2]], wires=[0, 1])
             return qml.expval(qml.PauliZ(0))
 
-        custom_decomps = {"Hadamard": custom_hadamard, qml.CNOT: custom_cnot}
-        decomp_dev = qml.device("default.qubit.legacy", wires=2, custom_decomps=custom_decomps)
-        decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
-        _ = decomp_qnode()
-        decomp_ops = decomp_qnode.qtape.operations
+        _ = circuit()
+        decomp_ops = circuit.tape.operations
 
         assert len(decomp_ops) == 7
 
@@ -630,20 +663,22 @@ class TestCreateCustomDecompExpandFn:
         assert np.isclose(decomp_ops[6].parameters[0], np.pi / 2)
         assert decomp_ops[6].wires == Wires(1)
 
-    def test_custom_decomp_template_to_template(self):
+    @pytest.mark.parametrize("device_name", ["default.qubit", "default.qubit.legacy"])
+    def test_custom_decomp_template_to_template(self, device_name):
         """Test that decomposing a template into another template and some
         gates yields the correct results."""
 
+        # BasicEntanglerLayers custom decomposition involves AngleEmbedding
+        custom_decomps = {"BasicEntanglerLayers": custom_basic_entangler_layers, "RX": custom_rx}
+        decomp_dev = qml.device(device_name, wires=2, custom_decomps=custom_decomps)
+
+        @qml.qnode(decomp_dev, expansion_strategy="device")
         def circuit():
             qml.BasicEntanglerLayers([[0.1, 0.2]], wires=[0, 1])
             return qml.expval(qml.PauliZ(0))
 
-        # BasicEntanglerLayers custom decomposition involves AngleEmbedding
-        custom_decomps = {"BasicEntanglerLayers": custom_basic_entangler_layers, "RX": custom_rx}
-        decomp_dev = qml.device("default.qubit.legacy", wires=2, custom_decomps=custom_decomps)
-        decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
-        _ = decomp_qnode()
-        decomp_ops = decomp_qnode.qtape.operations
+        _ = circuit()
+        decomp_ops = circuit.tape.operations
 
         assert len(decomp_ops) == 5
 
@@ -664,24 +699,35 @@ class TestCreateCustomDecompExpandFn:
         assert decomp_ops[4].name == "CNOT"
         assert decomp_ops[4].wires == Wires([0, 1])
 
-    def test_custom_decomp_different_depth(self):
+    @pytest.mark.parametrize("device_name", ["default.qubit", "default.qubit.legacy"])
+    def test_custom_decomp_different_depth(self, device_name):
         """Test that alternative expansion depths can be specified."""
 
-        def circuit():
-            qml.BasicEntanglerLayers([[0.1, 0.2]], wires=[0, 1])
-            return qml.expval(qml.PauliZ(0))
-
         # BasicEntanglerLayers custom decomposition involves AngleEmbedding. If
         # expansion depth is 2, the AngleEmbedding will still be decomposed into
         # RX (since it's not a supported operation on the device), but the RX will
         # not be further decomposed even though the custom decomposition is specified.
         custom_decomps = {"BasicEntanglerLayers": custom_basic_entangler_layers, "RX": custom_rx}
-        decomp_dev = qml.device(
-            "default.qubit.legacy", wires=2, custom_decomps=custom_decomps, decomp_depth=2
+
+        decomp_dev_2 = qml.device(
+            device_name, wires=2, custom_decomps=custom_decomps, decomp_depth=2
         )
-        decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
-        _ = decomp_qnode()
-        decomp_ops = decomp_qnode.qtape.operations
+
+        decomp_dev_3 = qml.device(
+            device_name, wires=2, custom_decomps=custom_decomps, decomp_depth=3
+        )
+
+        def circuit():
+            qml.BasicEntanglerLayers([[0.1, 0.2]], wires=[0, 1])
+            return qml.expval(qml.PauliZ(0))
+
+        circuit2 = qml.QNode(circuit, decomp_dev_2, expansion_strategy="device")
+        circuit3 = qml.QNode(circuit, decomp_dev_3, expansion_strategy="device")
+
+        _ = circuit2()
+        _ = circuit3()
+
+        decomp_ops = circuit2.tape.operations
 
         assert len(decomp_ops) == 3
 
@@ -696,20 +742,25 @@ class TestCreateCustomDecompExpandFn:
         assert decomp_ops[2].name == "CNOT"
         assert decomp_ops[2].wires == Wires([0, 1])
 
-    def test_custom_decomp_with_adjoint(self):
+        decomp_ops = circuit3.tape.operations
+        assert len(decomp_ops) == 5
+
+    @pytest.mark.parametrize("device_name", ["default.qubit", "default.qubit.legacy"])
+    def test_custom_decomp_with_adjoint(self, device_name):
         """Test that applying an adjoint in the circuit results in the adjoint
         undergoing the custom decomposition."""
 
+        custom_decomps = {qml.RX: custom_rx}
+        decomp_dev = qml.device(device_name, wires="a", custom_decomps=custom_decomps)
+
+        @qml.qnode(decomp_dev, expansion_strategy="device")
         def circuit():
             # Adjoint is RX(-0.2), so expect RY(-0.2) H
             qml.adjoint(qml.RX, lazy=False)(0.2, wires="a")
             return qml.expval(qml.PauliZ("a"))
 
-        custom_decomps = {qml.RX: custom_rx}
-        decomp_dev = qml.device("default.qubit.legacy", wires="a", custom_decomps=custom_decomps)
-        decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
-        _ = decomp_qnode()
-        decomp_ops = decomp_qnode.qtape.operations
+        _ = circuit()
+        decomp_ops = circuit.tape.operations
 
         assert len(decomp_ops) == 2
 
@@ -720,7 +771,8 @@ class TestCreateCustomDecompExpandFn:
         assert decomp_ops[1].name == "Hadamard"
         assert decomp_ops[1].wires == Wires("a")
 
-    def test_custom_decomp_with_control(self):
+    @pytest.mark.parametrize("device_name", ["default.qubit", "default.qubit.legacy"])
+    def test_custom_decomp_with_control(self, device_name):
         """Test that decomposing a controlled version of a gate uses the custom decomposition
         for the base gate."""
 
@@ -728,11 +780,13 @@ class TestCreateCustomDecompExpandFn:
             num_wires = 1
 
             @staticmethod
-            def compute_decomposition(wire):
-                return [qml.S(wire)]
+            def compute_decomposition(wires):
+                return [qml.S(wires)]
+
+        original_decomp = CustomOp(0).decomposition()
 
         custom_decomps = {CustomOp: lambda wires: [qml.T(wires), qml.T(wires)]}
-        decomp_dev = qml.device("default.qubit.legacy", wires=2, custom_decomps=custom_decomps)
+        decomp_dev = qml.device(device_name, wires=2, custom_decomps=custom_decomps)
 
         @qml.qnode(decomp_dev, expansion_strategy="device")
         def circuit():
@@ -748,7 +802,10 @@ class TestCreateCustomDecompExpandFn:
             assert isinstance(op, qml.ops.op_math.Controlled)
             assert qml.equal(op.base, qml.T(0))
 
-    def test_custom_decomp_in_separate_context(self):
+        # check that new instances of the operator are not affected by the modifications made to get the decomposition
+        assert [op1 == op2 for op1, op2 in zip(CustomOp(0).decomposition(), original_decomp)]
+
+    def test_custom_decomp_in_separate_context_legacy(self):
         """Test that the set_decomposition context manager works."""
 
         dev = qml.device("default.qubit.legacy", wires=2)
@@ -784,7 +841,55 @@ class TestCreateCustomDecompExpandFn:
         assert circuit.qtape.operations[0].name == "CNOT"
         assert dev.custom_expand_fn is None
 
+    def test_custom_decomp_in_separate_context(self, mocker):
+        """Test that the set_decomposition context manager works for the new device API."""
+
+        # pylint:disable=protected-access
+
+        dev = qml.device("default.qubit", wires=2)
+        spy = mocker.spy(dev, "execute")
+
+        @qml.qnode(dev, expansion_strategy="device")
+        def circuit():
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(wires=0))
+
+        # Initial test
+        _ = circuit()
+
+        tape = spy.call_args_list[0][0][0][0]
+        assert len(tape.operations) == 1
+        assert tape.operations[0].name == "CNOT"
+
+        assert dev.preprocess()[0][2].transform.__name__ == "decompose"
+        assert dev.preprocess()[0][2].kwargs.get("decomposer", None) is None
+
+        # Test within the context manager
+        with qml.transforms.set_decomposition({qml.CNOT: custom_cnot}, dev):
+            _ = circuit()
+
+            assert dev.preprocess()[0][2].transform.__name__ == "decompose"
+            assert dev.preprocess()[0][2].kwargs.get("decomposer", None) is not None
+
+        tape = spy.call_args_list[1][0][0][0]
+        ops_in_context = tape.operations
+        assert len(ops_in_context) == 3
+        assert ops_in_context[0].name == "Hadamard"
+        assert ops_in_context[1].name == "CZ"
+        assert ops_in_context[2].name == "Hadamard"
+
+        # Check that afterwards, the device has gone back to normal
+        _ = circuit()
+
+        tape = spy.call_args_list[2][0][0][0]
+        ops_in_context = tape.operations
+        assert len(tape.operations) == 1
+        assert tape.operations[0].name == "CNOT"
+        assert dev.preprocess()[0][2].transform.__name__ == "decompose"
+        assert dev.preprocess()[0][2].kwargs.get("decomposer", None) is None
+
     # pylint: disable=cell-var-from-loop
+
     def test_custom_decomp_used_twice(self):
         """Test that creating a custom decomposition includes overwriting the
         correct method under the hood and produces expected results."""