Repository: PennyLaneAI/pennylane
Pull request #1900 (2021-11-16 15:49:03)

[OpRefactor] Add custom decomposition context manager to device
**Context:** Operator overhaul

**Description of the Change:** Adds custom decompositions at the device level by creating a `custom_expand_fn` and context manager. Also adds a context manager to enable custom decompositions to be temporarily used by devices. Subsumes PR #1872.  

Example: suppose we are running on a trapped-ion machine, and would like to implement a decomposition of the `CNOT` gate into `IsingXX`. We can do so by defining a decomposition (which has the same signature as that of the normal operation), and passing the device a dictionary detailing the custom decompositions.

```python
def ion_trap_cnot(wires):
    return [
        qml.RY(np.pi/2, wires=wires[0]),
        qml.IsingXX(np.pi/2, wires=wires),
        qml.RX(-np.pi/2, wires=wires[0]),
        qml.RY(-np.pi/2, wires=wires[0]),
        qml.RY(-np.pi/2, wires=wires[1])
    ]

dev = qml.device('default.qubit', wires=2, custom_decomps={"CNOT" : ion_trap_cnot})
```

Now, we can see the decomposition get applied by drawing a QNode using the `device` expansion strategy.

```pycon
>>> @qml.qnode(dev)
... def run_cnot():
...     qml.CNOT(wires=[0, 1])
...     return qml.expval(qml.PauliX(wires=1))
>>> print(qml.draw(run_cnot, expansion_strategy="device")())
 0: ──RY(1.57)──╭IsingXX(1.57)──RX(-1.57)──RY(-1.57)──┤     
 1: ────────────╰IsingXX(1.57)──RY(-1.57)─────────────┤ ⟨X⟩ 
```

**Benefits:** Provides an intuitive means of implementing an arbitrary number of custom decompositions at runtime. There is minimal modification to the code of the device, and no modification of any operator code, since everything is done via context managers. 

**Possible Drawbacks:** Currently, this works by over-riding the `decompose` method of the specified operations within a given context. However this may not work for templates, which implement `expand`.

**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/1900?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1900](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1900?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (dfdd2f1) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/590704b9ce58301b0cb18ea99f8ffc8c02188e1a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (590704b) will **decrease** coverage by `0.17%`.
> The diff coverage is `26.82%`.

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

```diff
@@            Coverage Diff             @@
##           master    #1900      +/-   ##
==========================================
- Coverage   98.98%   98.80%   -0.18%     
==========================================
  Files         219      219              
  Lines       16512    16548      +36     
==========================================
+ Hits        16344    16350       +6     
- Misses        168      198      +30     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1900?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_qubit\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1900/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19xdWJpdF9kZXZpY2UucHk=) | `90.05% <18.91%> (-8.64%)` | :arrow_down: |
| [pennylane/\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1900/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19kZXZpY2UucHk=) | `98.29% <100.00%> (+<0.01%)` | :arrow_up: |
| [pennylane/devices/default\_qubit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1900/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdC5weQ==) | `100.00% <100.00%> (ø)` | |

------

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

glassnotes: > Is this outdated, since you have some tests that work with templates?

Actually what I was thinking about here is if someone wanted to implement a custom decomposition *for* a template. 

> A thought: do we have to test gradients/differentiation? Could anything go wrong?

If one is not smart about the decomposition (e.g., the new decomposition removes the parameter, or changes the operation that gets implemented like an incorrect circuit identity), then the gradients would not match up. But that is on the user, I think, to define decompositions that make sense!

> Re: the discussion on the device signatures, I think it might make sense to add the call to dev.custom_expand inside the device loader qml.device, rather than the Device classes/subclasses (see #1900 (comment)). This will ensure that it will work with all devices, without us needing to update the signatures of any of the devices.

Brilliant!

>  Finally, it would be nice to add a small utility decomposition_context (see #1900 (comment)), as atm it seems quite difficult to work out how to use the create_decomposition_custom_expand function!

Sure, will look into this!

glassnotes: @josh146 @mariaschuld  Think I've addressed everything except the following comment:

> Finally, it would be nice to add a small utility decomposition_context (see #1900 (comment)), as atm it seems quite difficult to work out how to use the create_decomposition_custom_expand function!

I still have to think about this, because my hope here was to eliminate the need for the user to work with the contexts at all, and have everything handled by the device (well, now, the device loader).
glassnotes: [sc-7225]

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index b8493ace9..f05b43d14 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,6 +4,75 @@
 
 <h3>New features since last release</h3>
 
+* Custom decompositions can now be applied to operations at the device level.
+  [(#1900)](https://github.com/PennyLaneAI/pennylane/pull/1900)
+
+  For example, suppose we would like to implement the following QNode:
+
+  ```python
+  def circuit(weights):
+      qml.BasicEntanglerLayers(weights, wires=[0, 1, 2])
+      return qml.expval(qml.PauliZ(0))
+
+  original_dev = qml.device("default.qubit", wires=3)
+  original_qnode = qml.QNode(circuit, original_dev)
+  ```
+
+  ```pycon
+  >>> weights = np.array([[0.4, 0.5, 0.6]])
+  >>> print(qml.draw(original_qnode, expansion_strategy="device")(weights))
+   0: ──RX(0.4)──╭C──────╭X──┤ ⟨Z⟩
+   1: ──RX(0.5)──╰X──╭C──│───┤
+   2: ──RX(0.6)──────╰X──╰C──┤
+  ```
+
+  Now, let's swap out the decomposition of the `CNOT` gate into `CZ`
+  and `Hadamard`, and furthermore the decomposition of `Hadamard` into
+  `RZ` and `RY` rather than the decomposition already available in PennyLane.
+  We define the two decompositions like so, and pass them to a device:
+
+  ```python
+  def custom_cnot(wires):
+      return [
+          qml.Hadamard(wires=wires[1]),
+          qml.CZ(wires=[wires[0], wires[1]]),
+          qml.Hadamard(wires=wires[1])
+      ]
+
+  def custom_hadamard(wires):
+      return [
+          qml.RZ(np.pi, wires=wires),
+          qml.RY(np.pi / 2, wires=wires)
+      ]
+
+  # Can pass the operation itself, or a string
+  custom_decomps = {qml.CNOT : custom_cnot, "Hadamard" : custom_hadamard}
+
+  decomp_dev = qml.device("default.qubit", wires=3, custom_decomps=custom_decomps)
+  decomp_qnode = qml.QNode(circuit, decomp_dev)
+  ```
+
+  Now when we draw or run a QNode on this device, the gates will be expanded
+  according to our specifications:
+
+  ```pycon
+  >>> print(qml.draw(decomp_qnode, expansion_strategy="device")(weights))
+   0: ──RX(0.4)──────────────────────╭C──RZ(3.14)──RY(1.57)──────────────────────────╭Z──RZ(3.14)──RY(1.57)──┤ ⟨Z⟩
+   1: ──RX(0.5)──RZ(3.14)──RY(1.57)──╰Z──RZ(3.14)──RY(1.57)──╭C──────────────────────│───────────────────────┤
+   2: ──RX(0.6)──RZ(3.14)──RY(1.57)──────────────────────────╰Z──RZ(3.14)──RY(1.57)──╰C──────────────────────┤
+  ```
+
+  A separate context manager, `set_decomposition`, has also been implemented to enable
+  application of custom decompositions on devices that have already been created.
+
+  ```pycon
+  >>> with qml.transforms.set_decomposition(custom_decomps, original_dev):
+  ...     print(qml.draw(original_qnode, expansion_strategy="device")(weights))
+   0: ──RX(0.4)──────────────────────╭C──RZ(3.14)──RY(1.57)──────────────────────────╭Z──RZ(3.14)──RY(1.57)──┤ ⟨Z⟩
+   1: ──RX(0.5)──RZ(3.14)──RY(1.57)──╰Z──RZ(3.14)──RY(1.57)──╭C──────────────────────│───────────────────────┤
+   2: ──RX(0.6)──RZ(3.14)──RY(1.57)──────────────────────────╰Z──RZ(3.14)──RY(1.57)──╰C──────────────────────┤
+  ```
+
 * PennyLane now supports drawing a QNode with matplotlib!
   [(#1803)](https://github.com/PennyLaneAI/pennylane/pull/1803)
 
diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index e789dff0d..e280f393a 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -169,6 +169,19 @@ def device(name, *args, **kwargs):
     the  `available plugins <https://pennylane.ai/plugins.html>`_ for more
     details.
 
+    Args:
+        name (str): the name of the device to load
+        wires (int): the number of wires (subsystems) to initialise
+            the device with
+
+    Keyword Args:
+        config (pennylane.Configuration): a PennyLane configuration object
+            that contains global and/or device specific configurations.
+        custom_decomps (Dict[Union(str, qml.Operator), Callable]): Custom
+            decompositions to be applied by the device at runtime.
+        decomp_depth (int): For when custom decompositions are specified,
+            the maximum expansion depth used by the expansion function.
+
     All devices must be loaded by specifying their **short-name** as listed above,
     followed by the **wires** (subsystems) you wish to initialize. The *wires*
     argument can be an integer, in which case the wires of the device are addressed
@@ -218,6 +231,49 @@ def device(name, *args, **kwargs):
     >>> circuit(0.8)  # back to default of 10 samples
     [ 1  1  1 -1 -1  1  1  1  1  1]
 
+    When constructing a device, we may optionally pass a dictionary of custom
+    decompositions to be applied to certain operations upon device execution.
+    This is useful for enabling support of gates on devices where they would normally
+    be unsupported.
+
+    For example, suppose we are running on an ion trap device which does not
+    natively implement the CNOT gate, but we would still like to write our
+    circuits in terms of CNOTs. On a ion trap device, CNOT can be implemented
+    using the ``IsingXX`` gate. We first define a decomposition function
+    (such functions have the signature ``decomposition(*params, wires)``):
+
+    .. code-block:: python
+
+        def ion_trap_cnot(wires):
+            return [
+                qml.RY(np.pi/2, wires=wires[0]),
+                qml.IsingXX(np.pi/2, wires=wires),
+                qml.RX(-np.pi/2, wires=wires[0]),
+                qml.RY(-np.pi/2, wires=wires[0]),
+                qml.RY(-np.pi/2, wires=wires[1])
+            ]
+
+    Next, we create a device, and a QNode for testing. When constructing the
+    QNode, we can set the expansion strategy to ``"device"`` to ensure the
+    decomposition is applied and will be viewable when we draw the circuit.
+
+    .. code-block:: python
+
+        # As the CNOT gate normally has no decomposition, we can use default.qubit
+        # here for expository purposes.
+        dev = qml.device(
+            'default.qubit', wires=2, custom_decomps={"CNOT" : ion_trap_cnot}
+        )
+
+        @qml.qnode(dev, expansion_strategy="device")
+        def run_cnot():
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliX(wires=1))
+
+    >>> print(qml.draw(run_cnot)())
+     0: ──RY(1.57)──╭IsingXX(1.57)──RX(-1.57)──RY(-1.57)──┤
+     1: ────────────╰IsingXX(1.57)──RY(-1.57)─────────────┤ ⟨X⟩
+
     Some devices may accept additional arguments. For instance,
     ``default.gaussian`` accepts the keyword argument ``hbar``, to set
     the convention used in the commutation relation :math:`[\x,\p]=i\hbar`
@@ -225,15 +281,6 @@ def device(name, *args, **kwargs):
 
     Please refer to the documentation for the individual devices to see any
     additional arguments that might be required or supported.
-
-    Args:
-        name (str): the name of the device to load
-        wires (int): the number of wires (subsystems) to initialise
-            the device with
-
-    Keyword Args:
-        config (pennylane.Configuration): a PennyLane configuration object
-            that contains global and/or device specific configurations.
     """
     if name not in plugin_devices:
         # Device does not exist in the loaded device list.
@@ -255,6 +302,11 @@ def device(name, *args, **kwargs):
             options.update(config[name.split(".")[0] + ".global"])
             options.update(config[name])
 
+        # Pop the custom decomposition keyword argument; we will use it here
+        # only and not pass it to the device.
+        custom_decomps = kwargs.pop("custom_decomps", None)
+        decomp_depth = kwargs.pop("decomp_depth", 10)
+
         kwargs.pop("config", None)
         options.update(kwargs)
 
@@ -269,8 +321,18 @@ def device(name, *args, **kwargs):
                 )
             )
 
-        # load device
-        return plugin_device_class(*args, **options)
+        # Construct the device
+        dev = plugin_device_class(*args, **options)
+
+        # 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)
+
+        return dev
 
     raise DeviceError("Device does not exist. Make sure the required plugin is installed.")
 
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index a5eaf12bd..e930a7e22 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -76,6 +76,7 @@ both transforms, and decompositions within the larger PennyLane codebase.
 
     ~transforms.zyz_decomposition
     ~transforms.two_qubit_decomposition
+    ~transforms.set_decomposition
 
 Transforms that act on tapes
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -103,6 +104,7 @@ to help build custom QNode, quantum function, and tape transforms:
     ~qfunc_transform
     ~transforms.make_tape
     ~transforms.create_expand_fn
+    ~transforms.create_decomp_expand_fn
     ~transforms.expand_invalid_trainable
     ~transforms.expand_multipar
     ~transforms.expand_nonunitary_gen
@@ -137,4 +139,6 @@ from .tape_expand import (
     expand_multipar,
     expand_nonunitary_gen,
     create_expand_fn,
+    create_decomp_expand_fn,
+    set_decomposition,
 )
diff --git a/pennylane/transforms/tape_expand.py b/pennylane/transforms/tape_expand.py
index b6666c0a3..f0c42de35 100644
--- a/pennylane/transforms/tape_expand.py
+++ b/pennylane/transforms/tape_expand.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 """This module contains tape expansion functions and stopping criteria to
 generate such functions from."""
+# pylint: disable=unused-argument
+import contextlib
 
 import pennylane as qml
 from pennylane.operation import (
@@ -25,6 +27,11 @@ from pennylane.operation import (
     not_tape,
 )
 
+# Needed for custom decomposition context manager
+from pennylane.transforms.qfunc_transforms import NonQueuingTape
+
+NonQueuingTape = type("NonQueuingTape", (NonQueuingTape, qml.tape.QuantumTape), {})
+
 
 def _update_trainable_params(tape):
     params = tape.get_parameters(trainable_only=False)
@@ -180,3 +187,165 @@ expand_invalid_trainable = create_expand_fn(
     stop_at=not_tape | is_measurement | (~is_trainable) | has_grad_method,
     docstring=_expand_invalid_trainable_doc,
 )
+
+
+@contextlib.contextmanager
+def _custom_decomp_context(custom_decomps):
+    """A context manager for applying custom decompositions of operations."""
+
+    # Creates an individual context
+    @contextlib.contextmanager
+    def _custom_decomposition(obj, fn):
+        # Covers the case where the user passes a string to indicate the Operator
+        if isinstance(obj, str):
+            obj = getattr(qml, obj)
+
+        original_decomp_method = obj.decompose
+
+        # This is the method that will override the operations .decompose method
+        def new_decomp_method(self):
+            with NonQueuingTape():
+                if self.num_params == 0:
+                    return fn(self.wires)
+                return fn(*self.parameters, self.wires)
+
+        try:
+            # Explicitly set the new .decompose method
+            obj.decompose = new_decomp_method
+            yield
+
+        finally:
+            obj.decompose = original_decomp_method
+
+    # Loop through the decomposition dictionary and create all the contexts
+    try:
+        with contextlib.ExitStack() as stack:
+            for obj, fn in custom_decomps.items():
+                # We enter a new context for each decomposition the user passes
+                stack.enter_context(_custom_decomposition(obj, fn))
+
+            stack = stack.pop_all()
+
+        yield
+
+    finally:
+        stack.close()
+
+
+def create_decomp_expand_fn(custom_decomps, dev, decomp_depth=10):
+    """Creates a custom expansion function 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 (qml.Device): A quantum device.
+        decomp_depth: The maximum depth of the expansion.
+
+    Returns:
+        Callable: A custom expansion function that a device can call to expand
+        its tapes within a context manager that applies custom decompositions.
+
+    **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}
+    >>> expand_fn = qml.transforms.create_decomp_expand_fn(custom_decomps, dev)
+    >>> dev.custom_expand(expand_fn)
+    """
+    custom_op_names = [op if isinstance(op, str) else op.__name__ for op in custom_decomps.keys()]
+
+    # Create a new expansion function; stop at things that do not have
+    # custom decompositions, or that satisfy the regular device stopping criteria
+    custom_fn = qml.transforms.create_expand_fn(
+        decomp_depth,
+        stop_at=qml.BooleanFn(lambda obj: obj.name not in custom_op_names),
+        device=dev,
+    )
+
+    # Finally, we set the device's custom_expand_fn to a new one that
+    # runs in a context where the decompositions have been replaced.
+    def custom_decomp_expand(self, circuit, max_expansion=decomp_depth):
+        with _custom_decomp_context(custom_decomps):
+            return custom_fn(circuit, max_expansion=max_expansion)
+
+    return custom_decomp_expand
+
+
+@contextlib.contextmanager
+def set_decomposition(custom_decomps, dev, decomp_depth=10):
+    """Context manager for setting custom decompositions.
+
+    Args:
+        custom_decomps (Dict[Union(str, qml.operation.Operation), Callable]): Custom
+            decompositions to be applied by the device at runtime.
+        dev (qml.Device): A quantum device.
+        decomp_depth: The maximum depth of the expansion.
+
+    **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])
+            ]
+
+    This context manager can be used to temporarily change a devices expansion
+    function to one that takes into account the custom decompositions.
+
+    .. code-block:: python
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, expansion_strategy="device")
+        def circuit():
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(wires=0))
+
+    >>> print(qml.draw(circuit)())
+     0: ──╭C──┤ ⟨Z⟩
+     1: ──╰X──┤
+
+    Now let's set up a context where the custom decomposition will be applied:
+
+    >>> with qml.transforms.set_decomposition({qml.CNOT : custom_cnot}, dev):
+    ...     print(qml.draw(circuit)())
+     0: ─────╭C─────┤ ⟨Z⟩
+     1: ──H──╰Z──H──┤
+    """
+    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
+    )
+
+    # 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


Test patch:
diff --git a/tests/transforms/test_tape_expand.py b/tests/transforms/test_tape_expand.py
index 514d59a9d..7f5e1d61d 100644
--- a/tests/transforms/test_tape_expand.py
+++ b/tests/transforms/test_tape_expand.py
@@ -17,6 +17,7 @@ Unit tests for tape expansion stopping criteria and expansion functions.
 import pytest
 import numpy as np
 import pennylane as qml
+from pennylane.wires import Wires
 
 
 class TestCreateExpandFn:
@@ -328,3 +329,403 @@ class TestExpandInvalidTrainable:
 
         assert new_tape is tape
         spy.assert_not_called()
+
+
+# Custom decomposition functions for testing.
+def custom_cnot(wires):
+    return [
+        qml.Hadamard(wires=wires[1]),
+        qml.CZ(wires=[wires[0], wires[1]]),
+        qml.Hadamard(wires=wires[1]),
+    ]
+
+
+def custom_hadamard(wires):
+    return [qml.RZ(np.pi, wires=wires), qml.RY(np.pi / 2, wires=wires)]
+
+
+# Incorrect, for testing purposes only
+def custom_rx(params, wires):
+    return [qml.RY(params, wires=wires), qml.Hadamard(wires=wires)]
+
+
+# To test the gradient; use circuit identity RY(theta) = X RY(-theta) X
+def custom_rot(phi, theta, omega, wires):
+    return [
+        qml.RZ(phi, wires=wires),
+        qml.PauliX(wires=wires),
+        qml.RY(-theta, wires=wires),
+        qml.PauliX(wires=wires),
+        qml.RZ(omega, wires=wires),
+    ]
+
+
+# Decompose a template into another template
+def custom_basic_entangler_layers(weights, wires):
+    return [
+        qml.AngleEmbedding(weights, wires=wires),
+        qml.broadcast(qml.CNOT, pattern="ring", wires=wires),
+    ]
+
+
+class TestCreateCustomDecompExpandFn:
+    """Tests for the gradient expand function"""
+
+    def test_no_custom_decomp(self):
+        """Test that sending an empty dictionary results in no decompositions."""
+
+        def circuit():
+            qml.Hadamard(wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0))
+
+        original_dev = qml.device("default.qubit", wires=3)
+        decomp_dev = qml.device("default.qubit", wires=3, custom_decomps={})
+
+        original_qnode = qml.QNode(circuit, original_dev, expansion_strategy="device")
+        decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
+
+        original_res = original_qnode()
+        decomp_res = decomp_qnode()
+
+        assert np.isclose(original_res, decomp_res)
+        assert [
+            orig_op.name == decomp_op.name
+            for orig_op, decomp_op in zip(
+                original_qnode.qtape.operations, decomp_qnode.qtape.operations
+            )
+        ]
+
+    def test_no_custom_decomp_template(self):
+        """Test that sending an empty dictionary results in no decomposition
+        when a template is involved, except the decomposition expected from the device."""
+
+        def circuit():
+            qml.BasicEntanglerLayers([[0.1, 0.2]], wires=[0, 1])
+            return qml.expval(qml.PauliZ(0))
+
+        original_dev = qml.device("default.qubit", wires=3)
+        decomp_dev = qml.device("default.qubit", wires=3, custom_decomps={})
+
+        original_qnode = qml.QNode(circuit, original_dev, expansion_strategy="device")
+        decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
+
+        original_res = original_qnode()
+        decomp_res = decomp_qnode()
+
+        assert np.isclose(original_res, decomp_res)
+        assert [
+            orig_op.name == decomp_op.name
+            for orig_op, decomp_op in zip(
+                original_qnode.qtape.operations, decomp_qnode.qtape.operations
+            )
+        ]
+
+    @pytest.mark.parametrize("device_name", ["default.qubit", "lightning.qubit"])
+    def test_one_custom_decomp(self, device_name):
+        """Test that specifying a single custom decomposition works as expected."""
+
+        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
+
+        assert len(decomp_ops) == 3
+
+        assert decomp_ops[0].name == "RZ"
+        assert np.isclose(decomp_ops[0].parameters[0], np.pi)
+
+        assert decomp_ops[1].name == "RY"
+        assert np.isclose(decomp_ops[1].parameters[0], np.pi / 2)
+
+        assert decomp_ops[2].name == "CNOT"
+
+    def test_no_decomp_with_depth_zero(self):
+        """Test that specifying a single custom decomposition works as expected."""
+
+        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", 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
+
+        assert len(decomp_ops) == 2
+        assert decomp_ops[0].name == "Hadamard"
+        assert decomp_ops[1].name == "CNOT"
+
+    def test_one_custom_decomp_gradient(self):
+        """Test that gradients are still correctly computed after a decomposition
+        that performs transpilation."""
+
+        def circuit(x):
+            qml.Hadamard(wires=0)
+            qml.Rot(x[0], x[1], x[2], wires=0)
+            qml.Hadamard(wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        original_dev = qml.device("default.qubit", wires=3)
+        decomp_dev = qml.device("default.qubit", 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")
+
+        x = qml.numpy.array([0.2, 0.3, 0.4], requires_grad=True)
+
+        original_res = original_qnode(x)
+        decomp_res = decomp_qnode(x)
+        assert np.allclose(original_res, decomp_res)
+
+        original_grad = qml.grad(original_qnode)(x)
+        decomp_grad = qml.grad(decomp_qnode)(x)
+        assert np.allclose(original_grad, decomp_grad)
+
+        expected_ops = ["Hadamard", "RZ", "PauliX", "RY", "PauliX", "RZ", "Hadamard"]
+        print(decomp_qnode.qtape.operations)
+        assert all(
+            [op.name == name for op, name in zip(decomp_qnode.qtape.operations, expected_ops)]
+        )
+
+    def test_nested_custom_decomp(self):
+        """Test that specifying two custom decompositions that have interdependence
+        works as expected."""
+
+        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", wires=2, custom_decomps=custom_decomps)
+        decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
+        _ = decomp_qnode()
+        decomp_ops = decomp_qnode.qtape.operations
+
+        assert len(decomp_ops) == 7
+
+        # Check the RZ gates are in the correct place
+        for idx in [0, 2, 5]:
+            assert decomp_ops[idx].name == "RZ"
+            assert np.isclose(decomp_ops[idx].parameters[0], np.pi)
+
+        assert decomp_ops[0].wires == Wires(0)
+        assert decomp_ops[2].wires == Wires(1)
+        assert decomp_ops[5].wires == Wires(1)
+
+        # Check RY are in the correct place
+        for idx in [1, 3, 6]:
+            assert decomp_ops[idx].name == "RY"
+            assert np.isclose(decomp_ops[idx].parameters[0], np.pi / 2)
+
+        assert decomp_ops[1].wires == Wires(0)
+        assert decomp_ops[3].wires == Wires(1)
+        assert decomp_ops[6].wires == Wires(1)
+
+        assert decomp_ops[4].name == "CZ"
+
+    def test_nested_custom_decomp_with_template(self):
+        """Test that specifying two custom decompositions that have interdependence
+        works as expected even when there is a template."""
+
+        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", wires=2, custom_decomps=custom_decomps)
+        decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
+        _ = decomp_qnode()
+        decomp_ops = decomp_qnode.qtape.operations
+
+        assert len(decomp_ops) == 7
+
+        assert decomp_ops[0].name == "RX"
+        assert decomp_ops[0].parameters[0] == 0.1
+        assert decomp_ops[0].wires == Wires(0)
+
+        assert decomp_ops[1].name == "RX"
+        assert decomp_ops[1].parameters[0] == 0.2
+        assert decomp_ops[1].wires == Wires(1)
+
+        assert decomp_ops[2].name == "RZ"
+        assert np.isclose(decomp_ops[2].parameters[0], np.pi)
+        assert decomp_ops[2].wires == Wires(1)
+
+        assert decomp_ops[3].name == "RY"
+        assert np.isclose(decomp_ops[3].parameters[0], np.pi / 2)
+        assert decomp_ops[3].wires == Wires(1)
+
+        assert decomp_ops[4].name == "CZ"
+        assert decomp_ops[4].wires == Wires([0, 1])
+
+        assert decomp_ops[5].name == "RZ"
+        assert np.isclose(decomp_ops[5].parameters[0], np.pi)
+        assert decomp_ops[5].wires == Wires(1)
+
+        assert decomp_ops[6].name == "RY"
+        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):
+        """Test that decomposing a template into another template and some
+        gates yields the correct results."""
+
+        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", wires=2, custom_decomps=custom_decomps)
+        decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
+        _ = decomp_qnode()
+        decomp_ops = decomp_qnode.qtape.operations
+
+        assert len(decomp_ops) == 5
+
+        assert decomp_ops[0].name == "RY"
+        assert decomp_ops[0].parameters[0] == 0.1
+        assert decomp_ops[0].wires == Wires(0)
+
+        assert decomp_ops[1].name == "Hadamard"
+        assert decomp_ops[1].wires == Wires(0)
+
+        assert decomp_ops[2].name == "RY"
+        assert np.isclose(decomp_ops[2].parameters[0], 0.2)
+        assert decomp_ops[2].wires == Wires(1)
+
+        assert decomp_ops[3].name == "Hadamard"
+        assert decomp_ops[3].wires == Wires(1)
+
+        assert decomp_ops[4].name == "CNOT"
+        assert decomp_ops[4].wires == Wires([0, 1])
+
+    def test_custom_decomp_different_depth(self):
+        """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 1, 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", wires=2, custom_decomps=custom_decomps, decomp_depth=1
+        )
+        decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
+        _ = decomp_qnode()
+        decomp_ops = decomp_qnode.qtape.operations
+
+        assert len(decomp_ops) == 3
+
+        assert decomp_ops[0].name == "RX"
+        assert np.isclose(decomp_ops[0].parameters[0], 0.1)
+        assert decomp_ops[0].wires == Wires(0)
+
+        assert decomp_ops[1].name == "RX"
+        assert np.isclose(decomp_ops[1].parameters[0], 0.2)
+        assert decomp_ops[1].wires == Wires(1)
+
+        assert decomp_ops[2].name == "CNOT"
+        assert decomp_ops[2].wires == Wires([0, 1])
+
+    def test_custom_decomp_with_adjoint(self):
+        """Test that applying an adjoint in the circuit results in the adjoint
+        undergoing the custom decomposition."""
+
+        def circuit():
+            # Adjoint is RX(-0.2), so expect RY(-0.2) H
+            qml.adjoint(qml.RX)(0.2, wires="a")
+            return qml.expval(qml.PauliZ("a"))
+
+        custom_decomps = {qml.RX: custom_rx}
+        decomp_dev = qml.device("default.qubit", wires="a", custom_decomps=custom_decomps)
+        decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
+        _ = decomp_qnode()
+        decomp_ops = decomp_qnode.qtape.operations
+
+        assert len(decomp_ops) == 2
+
+        assert decomp_ops[0].name == "RY"
+        assert decomp_ops[0].parameters[0] == -0.2
+        assert decomp_ops[0].wires == Wires("a")
+
+        assert decomp_ops[1].name == "Hadamard"
+        assert decomp_ops[1].wires == Wires("a")
+
+    def test_custom_decomp_with_control(self):
+        """Test that applying a controlled version of a gate results in the
+        controlled version of a decomposition."""
+
+        def circuit():
+            qml.ctrl(qml.Hadamard, control=0)(wires=1)
+            return qml.expval(qml.PauliZ(0))
+
+        custom_decomps = {qml.Hadamard: custom_hadamard}
+        decomp_dev = qml.device("default.qubit", wires=2, custom_decomps=custom_decomps)
+        decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
+        _ = decomp_qnode()
+        decomp_ops = decomp_qnode.qtape.operations
+
+        assert len(decomp_ops) == 2
+
+        assert decomp_ops[0].name == "CRZ"
+        assert np.isclose(decomp_ops[0].parameters[0], np.pi)
+        assert decomp_ops[0].wires == Wires([0, 1])
+
+        assert decomp_ops[1].name == "CRY"
+        assert np.isclose(decomp_ops[1].parameters[0], np.pi / 2)
+        assert decomp_ops[1].wires == Wires([0, 1])
+
+    def test_custom_decomp_in_separate_context(self):
+        """Test that the set_decomposition context manager works."""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, expansion_strategy="device")
+        def circuit():
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(wires=0))
+
+        # Initial test
+        _ = circuit()
+
+        assert len(circuit.qtape.operations) == 1
+        assert circuit.qtape.operations[0].name == "CNOT"
+        assert dev.custom_expand_fn is None
+
+        # Test within the context manager
+        with qml.transforms.set_decomposition({qml.CNOT: custom_cnot}, dev):
+            _ = circuit()
+            ops_in_context = circuit.qtape.operations
+
+            assert dev.custom_expand_fn is not None
+
+        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()
+
+        assert len(circuit.qtape.operations) == 1
+        assert circuit.qtape.operations[0].name == "CNOT"
+        assert dev.custom_expand_fn is None