Repository: PennyLaneAI/pennylane
Pull request #2240 (2022-02-28 08:38:22)

Add operator transform functionality
**Context:**

Currently, it is complex and unwieldy to write new functions that apply to operators in PennyLane, for several reasons:

- We aim to support multiple UIs, including applying the function directly to an instantiated operator,
  ```python
  op = some_function(op)
  ```
  as well as a functional approach
  ```python
  op = some_function(qml.Operation)(params, wires)
  ```

- The function must know how to deal with both a single `Operator` object, as well as multiple operators in the form of a tape

- The function must be usable both as a decorator and inline

- The function could support optional parameters that determine its usage

- operator functions sometimes need to take into account wire ordering

All of this results in a big barrier to contributions, while also requiring that contributors know about the internal implementation details of a lot of PennyLane's core.

In addition, for functions of operators that have already been implemented (e.g., `get_unitary_matrix()`, `adjoint()`, `ctrl()`), the same code boilerplate that enables the above UI is often repeated, not sufficiently tested, and error prone.

**Description of the Change:**

- Abstracts away all of the boilerplate required to construct the above UI into a single, well tested and well documented decorator `@qml.op_transform`.

- Decorating a function of an operator with `qml.op_transform` enables it to be used as per the context section, with minimal overhead to the developer.

- Wire ordering is taken into account automatically, and depends on the object the transform is applied to. For tapes and qfuncs, the wires available on the tape define the wire ordering. For QNodes, the wires available on the device define the wire ordering.

- By default, it creates transforms that accept a single operator. However, you can also register how the transform acts on a tape of _multiple_ operators. Once this is defined, the transform can be used _anywhere_ in PennyLane --- at the operator level for operator arithmetic, or at the qfunc/qnode level.

- There is no restriction to what the transform returns; it can return another operator, _or_ an arbitrary object such as a matrix or a string. E.g., `draw()`, `trace()`, `matrix()` etc. can all be rewritten using this framework.

- Operator transforms that return operators are special -- they are _also_ qfunc transforms, since the output of the operator transform will always be quantum in nature. As a result, they can be used within PennyLane's compilation frameworks.

**Benefits:** Will standardize a UI pattern that has been long assumed, ensure it is well tested, and make it a lot easier for contributors to contribute functions that act on operators.

**Possible Drawbacks:**

- The name

- Possible confusion with the other transform decorators.

Once this PR is merged, I plan to revisit the transform decorators, and see if they can be unified under a single `@qml.transform` decorator that _smartly_ knows which decorator to use under the hood.

**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/2240?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2240](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2240?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (cb1fe52) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/9b61f17649c4b7c45dc3d66576e0727433726430?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (9b61f17) will **decrease** coverage by `3.11%`.
> The diff coverage is `20.00%`.

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

```diff
@@            Coverage Diff             @@
##           master    #2240      +/-   ##
==========================================
- Coverage   99.26%   96.15%   -3.12%     
==========================================
  Files         231      232       +1     
  Lines       18349    18444      +95     
==========================================
- Hits        18215    17735     -480     
- Misses        134      709     +575     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2240?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://codecov.io/gh/PennyLaneAI/pennylane/pull/2240/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19faW5pdF9fLnB5) | `100.00% <ø> (ø)` | |
| [pennylane/transforms/op\_transform.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2240/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvb3BfdHJhbnNmb3JtLnB5) | `19.14% <19.14%> (ø)` | |
| [pennylane/transforms/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2240/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/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2240/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/batch/jax\_jit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2240/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4X2ppdC5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2240/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/2240/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/2240/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/transforms/classical\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2240/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY2xhc3NpY2FsX2phY29iaWFuLnB5) | `77.08% <0.00%> (-22.92%)` | :arrow_down: |
| [pennylane/math/utils.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2240/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvdXRpbHMucHk=) | `84.41% <0.00%> (-15.59%)` | :arrow_down: |
| ... and [17 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2240/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/2240?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/2240?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [9b61f17...cb1fe52](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2240?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/pennylane/__init__.py b/pennylane/__init__.py
index f9c1224cb..4deb27a80 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -66,6 +66,7 @@ from pennylane.transforms import (
     metric_tensor,
     specs,
     qfunc_transform,
+    op_transform,
     single_tape_transform,
     quantum_monte_carlo,
     apply_controlled_Q,
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index a2815afb3..35ac77a12 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -142,6 +142,7 @@ to help build custom QNode, quantum function, and tape transforms:
     ~single_tape_transform
     ~batch_transform
     ~qfunc_transform
+    ~op_transform
     ~transforms.make_tape
     ~transforms.map_batch_transform
     ~transforms.create_expand_fn
@@ -154,6 +155,7 @@ to help build custom QNode, quantum function, and tape transforms:
 # Import the decorators first to prevent circular imports when used in other transforms
 from .batch_transform import batch_transform, map_batch_transform
 from .qfunc_transforms import make_tape, single_tape_transform, qfunc_transform
+from .op_transforms import op_transform
 from .adjoint import adjoint
 from .batch_params import batch_params
 from .classical_jacobian import classical_jacobian
diff --git a/pennylane/transforms/op_transforms.py b/pennylane/transforms/op_transforms.py
new file mode 100644
index 000000000..7cf1a745d
--- /dev/null
+++ b/pennylane/transforms/op_transforms.py
@@ -0,0 +1,505 @@
+# Copyright 2018-2022 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This module contains the op_transform decorator.
+"""
+# pylint: disable=protected-access
+import functools
+import inspect
+import os
+import warnings
+
+import pennylane as qml
+
+
+class OperationTransformError(Exception):
+    """Raised when there is an error with the op_transform logic"""
+
+
+class op_transform:
+    r"""Convert a function that applies to operators into a functional transform.
+
+    This allows the operator function to be used across PennyLane
+    on both instantiated operators as well as quantum functions.
+
+    By default, this decorator creates functional transforms that
+    accept a single operator. However, you can also register how the
+    transform acts on multiple operators. Once this is defined,
+    the transform can be used anywhere in PennyLane --- at the operator
+    level for operator arithmetic, or at the qfunc/QNode level.
+
+    .. warning::
+
+        This is an experimental feature, and is subject to change.
+
+    Args:
+        fn (function): The function to register as the operator transform.
+            It can have an arbitrary number of arguments, but the first argument
+            **must** be the input operator.
+
+    **Example**
+
+    Consider an operator function that computes the trace of an operator:
+
+    .. code-block:: python
+
+        @qml.op_transform
+        def trace(op):
+            try:
+                return qml.math.real(qml.math.sum(op.get_eigvals()))
+            except qml.operation.EigvalsUndefinedError:
+                return qml.math.real(qml.math.trace(op.get_matrix()))
+
+    We can use this function as written:
+
+    >>> op = qml.RX(0.5, wires=0)
+    >>> trace(op)
+    1.9378248434212895
+
+    By using the ``op_transform`` decorator, we also enable it to be used
+    as a functional transform:
+
+    >>> trace(qml.RX)(0.5, wires=0)
+    1.9378248434212895
+
+    Note that if we apply our function to an operation that does not define its
+    matrix or eigenvalues representation, we get an error:
+
+    >>> weights = np.array([[[0.7, 0.6, 0.5], [0.1, 0.2, 0.3]]])
+    >>> trace(qml.StronglyEntanglingLayers(weights, wires=[0, 1]))
+    pennylane.operation.EigvalsUndefinedError
+    During handling of the above exception, another exception occurred:
+    pennylane.operation.MatrixUndefinedError
+
+    The most powerful reason for using ``op_transform`` is the ability to define
+    how the transform behaves if applied to a datastructure that supports multiple
+    operations, such as a qfunc, tape, or QNode.
+
+    We do this by defining a tape transform:
+
+    .. code-block:: python
+
+        @trace.tape_transform
+        def trace(tape):
+            tr = qml.math.trace(qml.matrix(tape))
+            return qml.math.real(tr)
+
+    We can now apply this transform directly to a qfunc:
+
+    >>> def circuit(x, y):
+    ...     qml.RX(x, wires=0)
+    ...     qml.Hadamard(wires=1)
+    ...     qml.CNOT(wires=[0, 1])
+    ...     qml.CRY(y, wires=[1, 0])
+    >>> trace(circuit)(0.1, 0.8)
+    1.4124461636742214
+
+    Our example above, applying our function to an operation that does not
+    define the matrix or eigenvalues, will now work, since PennyLane will
+    decompose the operation automatically into multiple operations:
+
+    >>> trace(qml.StronglyEntanglingLayers)(weights, wires=[0, 1])
+    0.4253851061350833
+
+    .. note::
+
+        If the operator transform takes additional (optional) transform parameters,
+        then the registered tape transform should take the same transform parameters.
+
+        E.g., consider a transform that takes the transform parameter ``lower``:
+
+        .. code-block:: python
+
+            @qml.op_transform
+            def name(op, lower=True):
+                return op.name().lower() if lower else op.name()
+
+            @name.tape_transform
+            def name(tape, lower=True):
+                return [name(op, lower=lower) for op in tape.operations]
+
+    If the transformation has purely quantum output, we can register the tape transformation
+    as a qfunc transformation in addition:
+
+    .. code-block:: python
+
+        @qml.op_transform
+        def simplify_rotation(op):
+            if op.name == "Rot":
+                params = op.parameters
+                wires = op.wires
+
+                if qml.math.allclose(params, 0):
+                    return
+
+                if qml.math.allclose(params[1:2], 0):
+                    return qml.RZ(params[0], wires)
+
+            return op
+
+        @simplify_rotation.tape_transform
+        @qml.qfunc_transform
+        def simplify_rotation(tape):
+            for op in tape.operations + tape.measurements:
+                if op.name == "Rot":
+                    simplify_rotation(op)
+                else:
+                    qml.apply(op)
+
+    We can now use this combined operator and quantum function transform in compilation pipelines:
+
+    .. code-block:: python
+
+        @qml.qnode(dev)
+        @qml.compile(pipeline=[simplify_rotation])
+        def circuit(weights):
+            ansatz(weights)
+            qml.CNOT(wires=[0, 1])
+            qml.Rot(0.0, 0.0, 0.0, wires=0)
+            return qml.expval(qml.PauliX(1))
+    """
+
+    def __new__(cls, *args, **kwargs):  # pylint: disable=unused-argument
+        if os.environ.get("SPHINX_BUILD") == "1":
+            # If called during a Sphinx documentation build,
+            # simply return the original function rather than
+            # instantiating the object. This allows the signature to
+            # be correctly displayed in the documentation.
+
+            warnings.warn(
+                "Operator transformations have been disabled, as a Sphinx "
+                "build has been detected via SPHINX_BUILD='1'. If this is not the "
+                "case, please set the environment variable SPHINX_BUILD='0'.",
+                UserWarning,
+            )
+
+            args[0].tape_transform = lambda x: x
+            return args[0]
+
+        return super().__new__(cls)
+
+    def __init__(self, fn):
+        if not callable(fn):
+            raise OperationTransformError(
+                f"The operator function to register, {fn}, "
+                "does not appear to be a valid Python function or callable."
+            )
+
+        self._fn = fn
+        self._sig = inspect.signature(fn).parameters
+        self._tape_fn = None
+        functools.update_wrapper(self, fn)
+
+    def __call__(self, *targs, **tkwargs):
+        obj = None
+
+        if targs:
+            # assume the first argument passed to the transform
+            # is the object we wish to transform
+            obj, *targs = targs
+
+        if isinstance(obj, (qml.operation.Operator, qml.tape.QuantumTape)) or callable(obj):
+            return self._create_wrapper(obj, *targs, **tkwargs)
+
+        # Input is not an operator nor a QNode nor a quantum tape nor a qfunc.
+        # Assume Python decorator syntax:
+        #
+        # op_func = op_transform(op_func)
+        # result = op_func(*transform_args)(obj)(*obj_args)
+        #
+        # or
+        #
+        # @op_func(*transform_args)
+        # @qml.qnode(dev)
+        # def circuit(...):
+        #     ...
+        # result = circuit(*qnode_args)
+
+        # Prepend the input to the transform args,
+        # and create a wrapper function.
+        if obj is not None:
+            targs = (obj,) + tuple(targs)
+
+        def wrapper(obj):
+            return self._create_wrapper(obj, *targs, **tkwargs)
+
+        return wrapper
+
+    def fn(self, obj, *args, **kwargs):
+        """Evaluate the underlying operator transform function.
+
+        If a corresponding tape transform for the operator has been registered
+        using the :attr:`.op_transform.tape_transform` decorator,
+        then if an exception is raised while calling the transform function,
+        this method will attempt to decompose the provided object for the tape
+        transform.
+
+        Args:
+            obj (.Operator, pennylane.QNode, .QuantumTape, or Callable): An operator, quantum node, tape,
+                or function that applies quantum operations.
+            *args: positional arguments to pass to the function
+            **kwargs: keyword arguments to pass to the function
+
+        Returns:
+            any: the result of evaluating the transform
+        """
+        try:
+            return self._fn(obj, *args, **kwargs)
+
+        except Exception as e1:  # pylint: disable=broad-except
+
+            try:
+                # attempt to decompose the operation and call
+                # the tape transform function if defined
+                return self.tape_fn(obj.expand(), *args, **kwargs)
+
+            except (AttributeError, OperationTransformError):
+                # if obj.expand() does not exist, or the tape transform
+                # function does not exist, simply raise the original exception
+                raise e1 from None
+
+    def tape_fn(self, obj, *args, **kwargs):
+        """The tape transform function.
+
+        This is the function that is called if a datastructure is passed
+        that contains multiple operations.
+
+        Args:
+            obj (pennylane.QNode, .QuantumTape, or Callable): A quantum node, tape,
+                or function that applies quantum operations.
+            *args: positional arguments to pass to the function
+            **kwargs: keyword arguments to pass to the function
+
+        Returns:
+            any: the result of evaluating the transform
+
+        Raises:
+            .OperationTransformError: if no tape transform function is defined
+
+        .. seealso:: :meth:`.op_transform.tape_transform`
+        """
+        if self._tape_fn is None:
+            raise OperationTransformError(
+                "This transform does not support tapes or QNodes with multiple operations."
+            )
+
+        return self._tape_fn(obj, *args, **kwargs)
+
+    @property
+    def is_qfunc_transform(self):
+        """bool: Returns ``True`` if the operator transform is also a qfunc transform.
+        That is, it maps one or more quantum operations to one or more quantum operations, allowing
+        the output of the transform to be used as a quantum function.
+
+        .. seealso:: :func:`~.qfunc_transform`
+        """
+        return isinstance(getattr(self._tape_fn, "tape_fn", None), qml.single_tape_transform)
+
+    def tape_transform(self, fn):
+        """Register a tape transformation to enable the operator transform
+        to apply to datastructures containing multiple operations, such as QNodes, qfuncs,
+        and tapes.
+
+        .. note::
+
+            The registered tape transform should have the same parameters as the
+            original operation transform function.
+
+        .. note::
+
+            If the transformation maps a tape to a tape (or equivalently, a qfunc to a qfunc)
+            then the transformation is simultaneously a :func:`~.qfunc_transform`, and
+            can be declared as such. This enables additional functionality, for example
+            the ability to use the transform in a compilation pipeline.
+
+        Args:
+            fn (callable): The function to register as the tape transform. This function
+                should accept a :class:`~.QuantumTape` as the first argument.
+
+        **Example**
+
+        .. code-block:: python
+
+            @qml.op_transform
+            def name(op, lower=False):
+                if lower:
+                    return op.name.lower()
+                return op.name
+
+            @name.tape_transform
+            def name(tape, lower=True):
+                return [name(op, lower=lower) for op in tape.operations]
+
+        We can now use this function on a qfunc, tape, or QNode:
+
+        >>> def circuit(x, y):
+        ...     qml.RX(x, wires=0)
+        ...     qml.Hadamard(wires=1)
+        ...     qml.CNOT(wires=[0, 1])
+        ...     qml.CRY(y, wires=[1, 0])
+        >>> name(circuit, lower=True)(0.1, 0.8)
+        ['rx', 'hadamard', 'cnot', 'cry']
+
+        If the transformation has purely quantum output, we can register the tape transformation
+        as a qfunc transformation in addition:
+
+        .. code-block:: python
+
+            @qml.op_transform
+            def simplify_rotation(op):
+                if op.name == "Rot":
+                    params = op.parameters
+                    wires = op.wires
+
+                    if qml.math.allclose(params, 0):
+                        return
+
+                    if qml.math.allclose(params[1:2], 0):
+                        return qml.RZ(params[0], wires)
+
+                return op
+
+            @simplify_rotation.tape_transform
+            @qml.qfunc_transform
+            def simplify_rotation(tape):
+                for op in tape.operations + tape.measurements:
+                    if op.name == "Rot":
+                        simplify_rotation(op)
+                    else:
+                        qml.apply(op)
+
+        We can now use this combined operator and quantum function transform in compilation pipelines:
+
+        .. code-block:: python
+
+            @qml.qnode(dev)
+            @qml.compile(pipeline=[simplify_rotation])
+            def circuit(weights):
+                ansatz(weights)
+                qml.CNOT(wires=[0, 1])
+                qml.Rot(0.0, 0.0, 0.0, wires=0)
+                return qml.expval(qml.PauliX(1))
+        """
+        self._tape_fn = fn
+        return self
+
+    def _create_wrapper(self, obj, *targs, wire_order=None, **tkwargs):
+        """Create a wrapper function that, when evaluated, transforms
+        ``obj`` according to transform arguments ``*targs`` and ``**tkwargs``
+        """
+
+        if isinstance(obj, qml.operation.Operator):
+            # Input is a single operation.
+            # op_transform(obj, *transform_args)
+            if wire_order is not None:
+                tkwargs["wire_order"] = wire_order
+
+            wrapper = self.fn(obj, *targs, **tkwargs)
+
+        elif isinstance(obj, qml.tape.QuantumTape):
+            # Input is a quantum tape. Get the quantum tape.
+            tape, verified_wire_order = self._make_tape(obj, wire_order)
+
+            if wire_order is not None:
+                tkwargs["wire_order"] = verified_wire_order
+
+            wrapper = self.tape_fn(tape, *targs, **tkwargs)
+
+        elif callable(obj):
+            # Input is a QNode, or qfunc (including single-operation qfuncs).
+            # Get the quantum tape.
+            def wrapper(*args, **kwargs):
+                nonlocal wire_order
+                tape, verified_wire_order = self._make_tape(obj, wire_order, *args, **kwargs)
+
+                if wire_order is not None or (
+                    "wire_order" in self._sig and isinstance(obj, qml.QNode)
+                ):
+                    # Use the verified wire order if:
+                    # - wire_order was passed to the transform
+                    # - The object is a QNode, and the function takes a wire_order argument
+                    tkwargs["wire_order"] = verified_wire_order
+
+                if isinstance(tape, qml.operation.Operator):
+                    return self.fn(tape, *targs, **tkwargs)
+
+                if self.is_qfunc_transform:
+                    # we must evaluate the qfunc transform at the original
+                    # function arguments
+                    return self.tape_fn(obj, *kwargs, **tkwargs)(*args, **kwargs)
+
+                return self.tape_fn(tape, *targs, **tkwargs)
+
+        else:
+            raise OperationTransformError(
+                "Input is not an Operator, tape, QNode, or quantum function"
+            )
+
+        return wrapper
+
+    @staticmethod
+    def _make_tape(obj, wire_order, *args, **kwargs):
+        """Given an input object, which may be:
+
+        - an object such as a tape or a operation, or
+        - a callable such as a QNode or a quantum function
+          (alongside the callable arguments ``args`` and ``kwargs``),
+
+        this function constructs and returns the tape/operation
+        represented by the object.
+
+        The ``wire_order`` argument determines whether a custom wire ordering
+        should be used. If not provided, the wire ordering defaults to the
+        objects wire ordering accessed via ``obj.wires``.
+
+        Returns:
+            tuple[.QuantumTape, Wires]: returns the tape and the verified wire order
+        """
+        if isinstance(obj, qml.QNode):
+            # user passed a QNode, get the tape
+            obj.construct(args, kwargs)
+            tape = obj.qtape
+            wires = obj.device.wires
+
+        elif isinstance(obj, qml.tape.QuantumTape):
+            # user passed a tape
+            tape = obj
+            wires = tape.wires
+
+        elif inspect.isclass(obj) and issubclass(obj, qml.operation.Operator):
+            with qml.tape.stop_recording():
+                tape = obj(*args, **kwargs)
+
+            wires = tape.wires
+
+        elif callable(obj):
+            # user passed something that is callable but not a tape or QNode.
+            tape = qml.transforms.make_tape(obj)(*args, **kwargs)
+            wires = tape.wires
+
+            # raise exception if it is not a quantum function
+            if len(tape.operations) == 0 and len(tape.measurements) == 0:
+                raise OperationTransformError("Quantum function contains no quantum operations")
+
+        # if no wire ordering is specified, take wire list from tape/device
+        wire_order = wires if wire_order is None else qml.wires.Wires(wire_order)
+
+        # check that all wire labels in the circuit are contained in wire_order
+        if not set(tape.wires).issubset(wire_order):
+            raise OperationTransformError(
+                f"Wires in circuit {tape.wires.tolist()} are inconsistent with "
+                f"those in wire_order {wire_order.tolist()}"
+            )
+
+        return tape, wire_order


Test patch:
diff --git a/tests/transforms/test_op_transform.py b/tests/transforms/test_op_transform.py
new file mode 100644
index 000000000..ade1fe4f2
--- /dev/null
+++ b/tests/transforms/test_op_transform.py
@@ -0,0 +1,554 @@
+# Copyright 2018-2022 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for the @op_transform framework"""
+import pytest
+
+import pennylane as qml
+from pennylane import numpy as np
+from pennylane.transforms.op_transforms import OperationTransformError
+
+
+class TestValidation:
+    """Test for validation and exceptions"""
+
+    def test_sphinx_build(self, monkeypatch):
+        """Test that op transforms are not created during Sphinx builds"""
+
+        @qml.op_transform
+        def my_transform(op):
+            return op.name
+
+        assert isinstance(my_transform, qml.op_transform)
+
+        monkeypatch.setenv("SPHINX_BUILD", "1")
+
+        with pytest.warns(UserWarning, match="Operator transformations have been disabled"):
+
+            @qml.op_transform
+            def my_transform(op):
+                return op.name
+
+        assert not isinstance(my_transform, qml.batch_transform)
+
+    def test_error_invalid_callable(self):
+        """Test that an error is raised if the transform
+        is applied to an invalid function"""
+
+        with pytest.raises(
+            OperationTransformError, match="does not appear to be a valid Python function"
+        ):
+            qml.op_transform(5)
+
+    def test_unknown_object(self):
+        """Test that an error is raised if the transform
+        is applied to an unknown object"""
+
+        @qml.op_transform
+        def my_transform(op):
+            return op.name
+
+        with pytest.raises(
+            OperationTransformError,
+            match="Input is not an Operator, tape, QNode, or quantum function",
+        ):
+            my_transform(5)(5)
+
+    def test_empty_qfunc(self):
+        """Test that an error is raised if the qfunc has no quantum operations
+        (e.g., it is not a qfunc)"""
+
+        @qml.op_transform
+        def my_transform(op):
+            return op.name
+
+        def qfunc(x):
+            return x**2
+
+        with pytest.raises(
+            OperationTransformError,
+            match="Quantum function contains no quantum operations",
+        ):
+            my_transform(qfunc)(0.5)
+
+
+class TestUI:
+    """Test the user interface of the op_transform, and ensure it applies
+    and works well for all combinations of inputs and styles"""
+
+    def test_instantiated_operator(self):
+        """Test that a transform can be applied to an instantiated operator"""
+
+        @qml.op_transform
+        def my_transform(op):
+            return op.name
+
+        op = qml.CRX(0.5, wires=[0, 2])
+        res = my_transform(op)
+        assert res == "CRX"
+
+    def test_single_operator_qfunc(self, mocker):
+        """Test that a transform can be applied to a quantum function
+        that contains a single operation"""
+        spy = mocker.spy(qml.op_transform, "_make_tape")
+
+        @qml.op_transform
+        def my_transform(op):
+            return op.name
+
+        res = my_transform(qml.CRX)(0.5, wires=[0, "a"])
+        assert res == "CRX"
+
+        # check default wire order
+        assert spy.spy_return[1].tolist() == [0, "a"]
+
+    def test_multiple_operator_error(self):
+        """Test that an exception is raised if the transform
+        is applied to a multi-op quantum function, without
+        the corresponding behaviour being registered"""
+
+        @qml.op_transform
+        def my_transform(op):
+            return op.name
+
+        @my_transform
+        def multi_op_qfunc(x):
+            qml.RX(x, wires=0)
+            qml.RY(0.65, wires=1)
+
+        with pytest.raises(
+            OperationTransformError, match="transform does not support .+ multiple operations"
+        ):
+            multi_op_qfunc(1.5)
+
+    def test_multiple_operator_tape(self, mocker):
+        """Test that a transform can be applied to a quantum function
+        with multiple operations as long as it is registered _how_
+        the transform applies to multiple operations."""
+        spy = mocker.spy(qml.op_transform, "_make_tape")
+
+        @qml.op_transform
+        def my_transform(op):
+            return op.name
+
+        @my_transform.tape_transform
+        def my_transform(tape):
+            return [op.name for op in tape.operations]
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(1.6, wires=0)
+            qml.RY(0.65, wires="a")
+
+        res = my_transform(tape)
+        assert res == ["RX", "RY"]
+
+        # check default wire order
+        assert spy.spy_return[1].tolist() == [0, "a"]
+
+    def test_multiple_operator_qfunc(self, mocker):
+        """Test that a transform can be applied to a quantum function
+        with multiple operations as long as it is registered _how_
+        the transform applies to multiple operations."""
+        spy = mocker.spy(qml.op_transform, "_make_tape")
+
+        @qml.op_transform
+        def my_transform(op):
+            return op.name
+
+        @my_transform.tape_transform
+        def my_transform(tape):
+            return [op.name for op in tape.operations]
+
+        @my_transform
+        def multi_op_qfunc(x):
+            if x > 1:
+                qml.RX(x, wires=0)
+                qml.RY(0.65, wires="a")
+            else:
+                qml.RZ(x, wires="b")
+
+        res = multi_op_qfunc(1.5)
+        assert res == ["RX", "RY"]
+        # check default wire order
+        assert spy.spy_return[1].tolist() == [0, "a"]
+
+        res = multi_op_qfunc(0.5)
+        assert res == ["RZ"]
+        # check default wire order
+        assert spy.spy_return[1].tolist() == ["b"]
+
+    def test_qnode(self, mocker):
+        """Test that a transform can be applied to a QNode
+        with multiple operations as long as it is registered _how_
+        the transform applies to multiple operations."""
+        dev = qml.device("default.qubit", wires=["a", 0, 3])
+        spy = mocker.spy(qml.op_transform, "_make_tape")
+
+        @qml.op_transform
+        def my_transform(op):
+            return op.name
+
+        @my_transform.tape_transform
+        def my_transform(tape):
+            return [op.name for op in tape.operations]
+
+        @my_transform
+        @qml.qnode(dev)
+        def multi_op_qfunc(x):
+            if x > 1:
+                qml.RX(x, wires=0)
+                qml.RY(0.65, wires="a")
+            else:
+                qml.RZ(x, wires=0)
+
+            return qml.probs(wires="a")
+
+        res = multi_op_qfunc(1.5)
+        assert res == ["RX", "RY"]
+        # check default wire order
+        assert spy.spy_return[1] == dev.wires
+
+        res = multi_op_qfunc(0.5)
+        assert res == ["RZ"]
+        # check default wire order
+        assert spy.spy_return[1] == dev.wires
+
+
+class TestTransformParameters:
+    def test_instantiated_operator(self):
+        """Test that a transform can be applied to an instantiated operator"""
+
+        @qml.op_transform
+        def my_transform(op, lower=False):
+            if lower:
+                return op.name.lower()
+            return op.name
+
+        op = qml.RX(0.5, wires=0)
+        res = my_transform(op, lower=True)
+        assert res == "rx"
+
+    def test_single_operator_qfunc(self):
+        """Test that a transform can be applied to a quantum function"""
+
+        @qml.op_transform
+        def my_transform(op, lower=False):
+            if lower:
+                return op.name.lower()
+            return op.name
+
+        res = my_transform(qml.RX, lower=True)(0.5, wires=0)
+        assert res == "rx"
+
+    def test_transform_parameters_qfunc_decorator(self):
+        """Test that transform parameters correctly work
+        when used as a decorator"""
+
+        @qml.op_transform
+        def my_transform(op, lower=False):
+            if lower:
+                return op.name.lower()
+            return op.name
+
+        @my_transform.tape_transform
+        def my_transform(tape, lower=False):
+            if lower:
+                return [op.name.lower() for op in tape.operations]
+            return [op.name for op in tape.operations]
+
+        @my_transform(True)
+        def multi_op_qfunc(x):
+            if x > 1:
+                qml.RX(x, wires=0)
+                qml.RY(0.65, wires=1)
+            else:
+                qml.RZ(x, wires=0)
+
+        res = multi_op_qfunc(1.5)
+        assert res == ["rx", "ry"]
+
+        res = multi_op_qfunc(0.5)
+        assert res == ["rz"]
+
+
+@qml.op_transform
+def simplify_rotation(op):
+    """Simplify Rot(x, 0, 0) to RZ(x) or Rot(0,0,0) to Identity"""
+    if op.name == "Rot":
+        params = op.parameters
+        wires = op.wires
+
+        if qml.math.allclose(params, 0):
+            return
+
+        if qml.math.allclose(params[1:2], 0):
+            return qml.RZ(params[0], wires)
+
+    return op
+
+
+@simplify_rotation.tape_transform
+@qml.qfunc_transform
+def simplify_rotation(tape):
+    """Define how simplify rotation works on a tape"""
+    for op in tape.operations + tape.measurements:
+        if op.name == "Rot":
+            simplify_rotation(op)
+        else:
+            qml.apply(op)
+
+
+class TestQFuncTransformIntegration:
+    """Test that @qfunc_transform seamlessly integrates
+    with an operator transform."""
+
+    def test_instantiated_operator(self):
+        """Test a qfunc and operator transform applied to
+        an op"""
+        dev = qml.device("default.qubit", wires=2)
+        assert simplify_rotation.is_qfunc_transform
+
+        weights = np.array([0.5, 0, 0])
+        op = qml.Rot(*weights, wires=0)
+        res = simplify_rotation(op)
+        assert res.name == "RZ"
+
+        @qml.qnode(dev)
+        def circuit():
+            simplify_rotation(op)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliX(1))
+
+        circuit()
+        ops = circuit.tape.operations
+        assert len(ops) == 2
+        assert ops[0].name == "RZ"
+        assert ops[1].name == "CNOT"
+
+    def test_qfunc_inside(self):
+        """Test a qfunc and operator transform
+        applied to a qfunc inside a qfunc"""
+        dev = qml.device("default.qubit", wires=2)
+
+        def ansatz(weights):
+            qml.Rot(*weights, wires=0)
+            qml.CRX(0.5, wires=[0, 1])
+
+        @qml.qnode(dev)
+        def circuit(weights):
+            simplify_rotation(ansatz)(weights)  # <--- qfunc is applied within circuit (inside)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliX(1))
+
+        weights = np.array([0.1, 0.0, 0.0])
+        circuit(weights)
+
+        ops = circuit.tape.operations
+        assert len(ops) == 3
+        assert ops[0].name == "RZ"
+        assert ops[1].name == "CRX"
+        assert ops[2].name == "CNOT"
+
+        weights = np.array([0.0, 0.0, 0.0])
+        circuit(weights)
+
+        ops = circuit.tape.operations
+        assert len(ops) == 2
+        assert ops[0].name == "CRX"
+        assert ops[1].name == "CNOT"
+
+    def test_qfunc_outside(self):
+        """Test a qfunc and operator transform
+        applied to qfunc"""
+        dev = qml.device("default.qubit", wires=2)
+
+        def ansatz(weights):
+            qml.Rot(*weights, wires=0)
+            qml.CRX(0.5, wires=[0, 1])
+
+        @qml.qnode(dev)
+        @simplify_rotation  # <--- qfunc is applied to circuit (outside)
+        def circuit(weights):
+            ansatz(weights)
+            qml.CNOT(wires=[0, 1])
+            qml.Rot(0.0, 0.0, 0.0, wires=0)
+            return qml.expval(qml.PauliX(1))
+
+        weights = np.array([0.1, 0.0, 0.0])
+        circuit(weights)
+
+        ops = circuit.tape.operations
+        assert len(ops) == 3
+        assert ops[0].name == "RZ"
+        assert ops[1].name == "CRX"
+        assert ops[2].name == "CNOT"
+
+    def test_compilation_pipeline(self):
+        """Test a qfunc and operator transform
+        applied to qfunc"""
+        dev = qml.device("default.qubit", wires=2)
+
+        def ansatz(weights):
+            qml.Rot(*weights, wires=0)
+            qml.CRX(0.5, wires=[0, 1])
+
+        @qml.qnode(dev)
+        @qml.compile(pipeline=[simplify_rotation])
+        def circuit(weights):
+            ansatz(weights)
+            qml.CNOT(wires=[0, 1])
+            qml.Rot(0.0, 0.0, 0.0, wires=0)
+            return qml.expval(qml.PauliX(1))
+
+        weights = np.array([0.1, 0.0, 0.0])
+        circuit(weights)
+
+        ops = circuit.tape.operations
+        assert len(ops) == 3
+        assert ops[0].name == "RZ"
+        assert ops[1].name == "CRX"
+        assert ops[2].name == "CNOT"
+
+    def test_qnode_error(self):
+        """Since a qfunc transform always returns a qfunc,
+        it cannot be applied to a QNode."""
+        dev = qml.device("default.qubit", wires=2)
+
+        def ansatz(weights):
+            qml.Rot(*weights, wires=0)
+            qml.CRX(0.5, wires=[0, 1])
+
+        @simplify_rotation
+        @qml.qnode(dev)
+        def circuit(weights):
+            ansatz(weights)
+            qml.CNOT(wires=[0, 1])
+            qml.Rot(0.0, 0.0, 0.0, wires=0)
+            return qml.expval(qml.PauliX(1))
+
+        weights = np.array([0.1, 0.0, 0.0])
+
+        with pytest.raises(TypeError):
+            circuit(weights)
+
+
+class TestExpansion:
+    """Test for operator and tape expansion"""
+
+    def test_auto_expansion(self, mocker):
+        """Test that an operator is automatically expanded as needed"""
+
+        @qml.op_transform
+        def matrix(op):
+            return op.get_matrix()
+
+        weights = np.ones([2, 3, 3])
+        op = qml.StronglyEntanglingLayers(weights, wires=[0, 2, "a"])
+
+        # strongly entangling layers does not define a matrix representation
+
+        with pytest.raises(qml.operation.MatrixUndefinedError):
+            op.get_matrix()
+
+        # attempting to call our operator transform will fail
+
+        with pytest.raises(qml.operation.MatrixUndefinedError):
+            matrix(op)
+
+        # if we define how the transform acts on a tape,
+        # then pennylane will automatically expand the object
+        # and apply the tape transform
+
+        @matrix.tape_transform
+        def matrix_tape(tape):
+            n_wires = len(tape.wires)
+            unitary_matrix = np.eye(2**n_wires)
+
+            for op in tape.operations:
+                mat = qml.operation.expand_matrix(matrix(op), op.wires, tape.wires)
+                unitary_matrix = mat @ unitary_matrix
+
+            return unitary_matrix
+
+        res = matrix(op)
+        assert isinstance(res, np.ndarray)
+        assert res.shape == (2**3, 2**3)
+
+
+matrix = qml.op_transform(lambda op, wire_order=None: op.get_matrix(wire_order=wire_order))
+
+
+@matrix.tape_transform
+def matrix_tape(tape, wire_order=None):
+    n_wires = len(wire_order)
+    unitary_matrix = np.eye(2**n_wires)
+
+    for op in tape.operations:
+        unitary_matrix = matrix(op, wire_order=wire_order) @ unitary_matrix
+
+    return unitary_matrix
+
+
+class TestWireOrder:
+    """Test for wire re-ordering"""
+
+    def test_instantiated_operator(self):
+        """Test that wire order can be passed to an instantiated operator"""
+        op = qml.PauliZ(wires=0)
+        res = matrix(op, wire_order=[1, 0])
+        expected = np.kron(np.eye(2), np.diag([1, -1]))
+        assert np.allclose(res, expected)
+
+    def test_single_operator_qfunc(self, mocker):
+        """Test that wire order can be passed to a quantum function"""
+        spy = mocker.spy(qml.op_transform, "_make_tape")
+        res = matrix(qml.PauliZ, wire_order=["a", 0])(0)
+        expected = np.kron(np.eye(2), np.diag([1, -1]))
+        assert np.allclose(res, expected)
+        assert spy.spy_return[1].tolist() == ["a", 0]
+
+    def test_tape(self, mocker):
+        """Test that wire order can be passed to a tape"""
+        spy = mocker.spy(qml.op_transform, "_make_tape")
+
+        with qml.tape.QuantumTape() as tape:
+            qml.PauliZ(wires=0)
+
+        res = matrix(tape, wire_order=["a", 0])
+        expected = np.kron(np.eye(2), np.diag([1, -1]))
+        assert np.allclose(res, expected)
+        assert spy.spy_return[1].tolist() == ["a", 0]
+
+    def test_inconsistent_wires_tape(self, mocker):
+        """Test that an exception is raised if the wire order and tape wires are inconsistent"""
+        with qml.tape.QuantumTape() as tape:
+            qml.PauliZ(wires=0)
+            qml.PauliY(wires="b")
+
+        with pytest.raises(
+            OperationTransformError,
+            match=r"Wires in circuit .+ inconsistent with those in wire\_order",
+        ):
+            matrix(tape, wire_order=["b", "a"])
+
+    def test_qfunc(self, mocker):
+        """Test that wire order can be passed to a qfunc"""
+        spy = mocker.spy(qml.op_transform, "_make_tape")
+
+        def qfunc():
+            qml.PauliZ(wires=0)
+
+        res = matrix(qfunc, wire_order=["a", 0])()
+        expected = np.kron(np.eye(2), np.diag([1, -1]))
+        assert np.allclose(res, expected)
+        assert spy.spy_return[1].tolist() == ["a", 0]