Repository: PennyLaneAI/pennylane
Pull request #7439 (2025-05-15 16:50:38)

[Decomposition] Register conditions to decomposition rules
**Context:**

Sometimes a decomposition rule is only applicable when certain conditions are met.

**Description of the Change:**

Added a `register_condition` decorator that allows users to bind a condition function to a decomposition rule.

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**

[sc-91089]

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7439?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: Patch coverage is `95.00000%` with `3 lines` in your changes missing coverage. Please review.
> Project coverage is 99.68%. Comparing base [(`1116746`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/11167460d8e857450b85797371719e3f95365413?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`316d4c5`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/316d4c5603f4cb48b4658e957fb1bcc2fb2cc835?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

| [Files with missing lines](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7439?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/decomposition/decomposition\_rule.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7439?src=pr&el=tree&filepath=pennylane%2Fdecomposition%2Fdecomposition_rule.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RlY29tcG9zaXRpb24vZGVjb21wb3NpdGlvbl9ydWxlLnB5) | 93.10% | [2 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7439?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |
| [...s/op\_math/decompositions/unitary\_decompositions.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7439?src=pr&el=tree&filepath=pennylane%2Fops%2Fop_math%2Fdecompositions%2Funitary_decompositions.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL2RlY29tcG9zaXRpb25zL3VuaXRhcnlfZGVjb21wb3NpdGlvbnMucHk=) | 95.45% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7439?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

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


```diff
@@            Coverage Diff             @@
##           master    #7439      +/-   ##
==========================================
- Coverage   99.69%   99.68%   -0.01%     
==========================================
  Files         529      529              
  Lines       50738    50755      +17     
==========================================
+ Hits        50583    50597      +14     
- Misses        155      158       +3     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7439?dropdown=coverage&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).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
astralcai: > * You define 2 decomposition rules for an operator, one has a condition and one doesn't.
> * You define only 1 possible decomposition rule and then the condition is not met -- do you get an error?
These two scenarios are fundamentally the same. If a decomposition rule has a condition, it is checked when we construct the graph. If the condition isn't met, this decomposition will just be ignored. When we solve the graph, we work with the decomposition rules that are available to us. If an operator does not have a decomposition rule (either because none was defined for it, or because each of them has a condition that isn't met), then we won't be able to decompose this operator. The decomposition graph would raise an error saying that no decomposition can be found for this operator.


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index dfe0d5bcb..e52cd608e 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -101,33 +101,33 @@
   0: ──RX(0.00)──RY(1.57)──RX(3.14)──GlobalPhase(-1.57)─┤  <Z>
   ```
 
-* Decomposition rules can be marked as not-applicable with :class:`~.decomposition.DecompositionNotApplicable`, allowing for flexibility when creating conditional decomposition 
-  rules based on parameters that affects the rule's resources.
-  [(#7211)](https://github.com/PennyLaneAI/pennylane/pull/7211)
+* A :func:`~.decomposition.register_condition` decorator is added that allows users to bind a condition to a
+  decomposition rule for when it is applicable. The condition should be a function that takes the
+  resource parameters of an operator as arguments and returns `True` or `False` based on whether
+  these parameters satisfy the condition for when this rule can be applied.
+  [(#7439)](https://github.com/PennyLaneAI/pennylane/pull/7439)
 
   ```python
   import pennylane as qml
-  from pennylane.decomposition import DecompositionNotApplicable
   from pennylane.math.decomposition import zyz_rotation_angles
   
-  def _zyz_resource(num_wires):
-      if num_wires != 1:
-          # This decomposition is only applicable when num_wires is 1
-          raise DecompositionNotApplicable
-      return {qml.RZ: 2, qml.RY: 1, qml.GlobalPhase: 1}
+  # The parameters must be consistent with ``qml.QubitUnitary.resource_keys``
+  def _zyz_condition(num_wires):
+    return num_wires == 1
 
-  @qml.register_resources(_zyz_resource)
+  @qml.register_condition(_zyz_condition)
+  @qml.register_resources({qml.RZ: 2, qml.RY: 1, qml.GlobalPhase: 1})
   def zyz_decomposition(U, wires, **__):
+      # Assumes that U is a 2x2 unitary matrix
       phi, theta, omega, phase = zyz_rotation_angles(U, return_global_phase=True)
       qml.RZ(phi, wires=wires[0])
       qml.RY(theta, wires=wires[0])
       qml.RZ(omega, wires=wires[0])
       qml.GlobalPhase(-phase)
   
-  qml.add_decomps(QubitUnitary, zyz_decomposition)
+  # This decomposition will be ignored for `QubitUnitary` on more than one wire.
+  qml.add_decomps(qml.QubitUnitary, zyz_decomposition)
   ```
-  
-  This decomposition will be ignored for `QubitUnitary` on more than one wire.
 
 * The :func:`~.transforms.decompose` transform now supports symbolic operators (e.g., `Adjoint` and `Controlled`) specified as strings in the `gate_set` argument
   when the new graph-based decomposition system is enabled.
@@ -318,6 +318,9 @@ Here's a list of deprecations made this release. For a more detailed breakdown o
 
 <h3>Internal changes ⚙️</h3>
 
+* Remove and reduce excessively expensive test cases in `tests/templates/test_subroutines/` that do not add value.
+  [(#7436)](https://github.com/PennyLaneAI/pennylane/pull/7436)
+
 * Stop using `pytest-timeout` in the PennyLane CI/CD pipeline.
   [(#7451)](https://github.com/PennyLaneAI/pennylane/pull/7451)
 
diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index 4c47482fe..28da4ed25 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -33,6 +33,7 @@ import pennylane.operation
 import pennylane.decomposition
 from pennylane.decomposition import (
     register_resources,
+    register_condition,
     add_decomps,
     list_decomps,
     resource_rep,
diff --git a/pennylane/decomposition/__init__.py b/pennylane/decomposition/__init__.py
index a0c4459d4..055e91660 100644
--- a/pennylane/decomposition/__init__.py
+++ b/pennylane/decomposition/__init__.py
@@ -58,6 +58,7 @@ Defining Decomposition Rules
     :toctree: api
 
     ~register_resources
+    ~register_condition
     ~resource_rep
     ~controlled_resource_rep
     ~adjoint_resource_rep
@@ -243,13 +244,11 @@ Utility Classes
     :toctree: api
 
     ~DecompositionError
-    ~DecompositionNotApplicable
 
 """
 
 from .utils import (
     DecompositionError,
-    DecompositionNotApplicable,
     enable_graph,
     disable_graph,
     enabled_graph,
@@ -265,6 +264,7 @@ from .resources import (
 )
 from .decomposition_rule import (
     register_resources,
+    register_condition,
     DecompositionRule,
     add_decomps,
     list_decomps,
diff --git a/pennylane/decomposition/decomposition_graph.py b/pennylane/decomposition/decomposition_graph.py
index cdbbba327..013da02d9 100644
--- a/pennylane/decomposition/decomposition_graph.py
+++ b/pennylane/decomposition/decomposition_graph.py
@@ -52,7 +52,7 @@ from .symbolic_decomposition import (
     same_type_adjoint_decomp,
     same_type_adjoint_ops,
 )
-from .utils import DecompositionError, DecompositionNotApplicable, translate_op_alias
+from .utils import DecompositionError, translate_op_alias
 
 
 class DecompositionGraph:  # pylint: disable=too-many-instance-attributes
@@ -201,12 +201,11 @@ class DecompositionGraph:  # pylint: disable=too-many-instance-attributes
         self, rule: DecompositionRule, op_node: CompressedResourceOp, op_node_idx: int
     ):
         """Adds a special decomposition rule to the graph."""
-        try:
-            decomp_resource = rule.compute_resources(**op_node.params)
-            d_node_idx = self._recursively_add_decomposition_node(rule, decomp_resource)
-            self._graph.add_edge(d_node_idx, op_node_idx, 0)
-        except DecompositionNotApplicable:
-            pass  # ignore decompositions that are not applicable to the given op params.
+        if not rule.is_applicable(**op_node.params):
+            return  # skip the decomposition rule if it is not applicable
+        decomp_resource = rule.compute_resources(**op_node.params)
+        d_node_idx = self._recursively_add_decomposition_node(rule, decomp_resource)
+        self._graph.add_edge(d_node_idx, op_node_idx, 0)
 
     def _add_adjoint_decomp_node(self, op_node: CompressedResourceOp, op_node_idx: int) -> int:
         """Adds an adjoint decomposition node."""
diff --git a/pennylane/decomposition/decomposition_rule.py b/pennylane/decomposition/decomposition_rule.py
index 9d1f79b24..c4511cf78 100644
--- a/pennylane/decomposition/decomposition_rule.py
+++ b/pennylane/decomposition/decomposition_rule.py
@@ -26,6 +26,74 @@ from pennylane.operation import Operator
 from .resources import CompressedResourceOp, Resources, resource_rep
 
 
+@overload
+def register_condition(condition: Callable) -> Callable[[Callable], DecompositionRule]: ...
+@overload
+def register_condition(condition: Callable, qfunc: Callable) -> DecompositionRule: ...
+def register_condition(
+    condition: Callable[..., bool], qfunc: Optional[Callable] = None
+) -> Callable[[Callable], DecompositionRule] | DecompositionRule:
+    """Binds a condition to a decomposition rule for when it is applicable.
+
+    .. note::
+
+        This function is only relevant when the new experimental graph-based decomposition system
+        (introduced in v0.41) is enabled via :func:`~pennylane.decomposition.enable_graph`. This new way of
+        performing decompositions is generally more resource-efficient and accommodates multiple alternative
+        decomposition rules for an operator. In this new system, custom decomposition rules are
+        defined as quantum functions, and it is currently required that every decomposition rule
+        declares its required resources using :func:`~.register_resources`.
+
+    Args:
+        condition (Callable): a function which takes the resource parameters of an operator as
+            arguments and returns ``True`` or ``False`` based on whether the decomposition rule
+            is applicable to an operator with the given resource parameters.
+        qfunc (Callable): the quantum function that implements the decomposition. If ``None``,
+            returns a decorator for acting on a function.
+
+    Returns:
+        DecompositionRule:
+            a data structure that represents a decomposition rule, which contains a PennyLane
+            quantum function representing the decomposition, and its resource function.
+
+    **Example**
+
+    This function can be used as a decorator to bind a condition function to a quantum function
+    that implements a decomposition rule.
+
+    .. code-block:: python
+
+        import pennylane as qml
+        from pennylane.math.decomposition import zyz_rotation_angles
+
+        # The parameters must be consistent with ``qml.QubitUnitary.resource_keys``
+        def _zyz_condition(num_wires):
+            return num_wires == 1
+
+        @qml.register_condition(_zyz_condition)
+        @qml.register_resources({qml.RZ: 2, qml.RY: 1, qml.GlobalPhase: 1})
+        def zyz_decomposition(U, wires, **__):
+            # Assumes that U is a 2x2 unitary matrix
+            phi, theta, omega, phase = zyz_rotation_angles(U, return_global_phase=True)
+            qml.RZ(phi, wires=wires[0])
+            qml.RY(theta, wires=wires[0])
+            qml.RZ(omega, wires=wires[0])
+            qml.GlobalPhase(-phase)
+
+        # This decomposition will be ignored for `QubitUnitary` on more than one wire.
+        qml.add_decomps(qml.QubitUnitary, zyz_decomposition)
+
+    """
+
+    def _decorator(_qfunc) -> DecompositionRule:
+        if isinstance(_qfunc, DecompositionRule):
+            _qfunc.set_condition(condition)
+            return _qfunc
+        return DecompositionRule(_qfunc, condition=condition)
+
+    return _decorator(qfunc) if qfunc else _decorator
+
+
 @overload
 def register_resources(resources: Callable | dict) -> Callable[[Callable], DecompositionRule]: ...
 @overload
@@ -178,7 +246,10 @@ def register_resources(
     """
 
     def _decorator(_qfunc) -> DecompositionRule:
-        return DecompositionRule(_qfunc, resources)
+        if isinstance(_qfunc, DecompositionRule):
+            _qfunc.set_resources(resources)
+            return _qfunc
+        return DecompositionRule(_qfunc, resources=resources)
 
     return _decorator(qfunc) if qfunc else _decorator
 
@@ -186,18 +257,32 @@ def register_resources(
 class DecompositionRule:  # pylint: disable=too-few-public-methods
     """Represents a decomposition rule for an operator."""
 
-    def __init__(self, func: Callable, resources: Callable | dict):
+    def __init__(
+        self,
+        func: Callable,
+        resources: Optional[Callable | dict] = None,
+        condition: Optional[Callable[..., bool]] = None,
+    ):
+
         self._impl = func
-        self._source = inspect.getsource(func)
+
+        try:
+            self._source = inspect.getsource(func)
+        except OSError:  # pragma: no cover
+            # OSError is raised if the source code cannot be retrieved
+            self._source = ""  # pragma: no cover
+
         if isinstance(resources, dict):
 
-            def resource_fn():
+            def resource_fn(*_, **__):
                 return resources
 
             self._compute_resources = resource_fn
         else:
             self._compute_resources = resources
 
+        self._condition = condition
+
     def __call__(self, *args, **kwargs):
         return self._impl(*args, **kwargs)
 
@@ -213,6 +298,28 @@ class DecompositionRule:  # pylint: disable=too-few-public-methods
         gate_counts = {_auto_wrap(op): count for op, count in gate_counts.items() if count > 0}
         return Resources(gate_counts)
 
+    def is_applicable(self, *args, **kwargs) -> bool:
+        """Checks whether this decomposition rule is applicable."""
+        if self._condition is None:
+            return True
+        return self._condition(*args, **kwargs)
+
+    def set_condition(self, condition: Callable[..., bool]) -> None:
+        """Sets the condition for this decomposition rule."""
+        self._condition = condition
+
+    def set_resources(self, resources: Callable | dict) -> None:
+        """Sets the resources for this decomposition rule."""
+
+        if isinstance(resources, dict):
+
+            def resource_fn(*_, **__):
+                return resources
+
+            self._compute_resources = resource_fn
+        else:
+            self._compute_resources = resources
+
 
 def _auto_wrap(op_type):
     """Conveniently wrap an operator type in a resource representation."""
diff --git a/pennylane/decomposition/utils.py b/pennylane/decomposition/utils.py
index 3203ee217..12c7f7ff4 100644
--- a/pennylane/decomposition/utils.py
+++ b/pennylane/decomposition/utils.py
@@ -24,10 +24,6 @@ class DecompositionError(Exception):
     """Base class for decomposition errors."""
 
 
-class DecompositionNotApplicable(Exception):
-    """Exception raised when a decomposition is not applicable to the given operator."""
-
-
 OP_NAME_ALIASES = {
     "X": "PauliX",
     "Y": "PauliY",
diff --git a/pennylane/ops/functions/assert_valid.py b/pennylane/ops/functions/assert_valid.py
index 2cdad355b..7a18e89ef 100644
--- a/pennylane/ops/functions/assert_valid.py
+++ b/pennylane/ops/functions/assert_valid.py
@@ -25,7 +25,7 @@ import numpy as np
 import scipy.sparse
 
 import pennylane as qml
-from pennylane.decomposition import DecompositionNotApplicable, DecompositionRule
+from pennylane.decomposition import DecompositionRule
 from pennylane.operation import EigvalsUndefinedError
 
 
@@ -116,12 +116,11 @@ def _check_decomposition_new(op, heuristic_resources=False):
 def _test_decomposition_rule(op, rule: DecompositionRule, heuristic_resources=False):
     """Tests that a decomposition rule is consistent with the operator."""
 
-    # Test that the resource function is correct
-    try:
-        resources = rule.compute_resources(**op.resource_params)
-    except DecompositionNotApplicable:
+    if not rule.is_applicable(**op.resource_params):
         return
 
+    # Test that the resource function is correct
+    resources = rule.compute_resources(**op.resource_params)
     gate_counts = resources.gate_counts
 
     with qml.queuing.AnnotatedQueue() as q:
diff --git a/pennylane/ops/op_math/decompositions/unitary_decompositions.py b/pennylane/ops/op_math/decompositions/unitary_decompositions.py
index ee23839b3..42ea3b81c 100644
--- a/pennylane/ops/op_math/decompositions/unitary_decompositions.py
+++ b/pennylane/ops/op_math/decompositions/unitary_decompositions.py
@@ -20,9 +20,8 @@ import numpy as np
 import scipy.sparse as sp
 
 from pennylane import capture, compiler, math, ops, queuing
-from pennylane.decomposition.decomposition_rule import DecompositionRule, register_resources
+from pennylane.decomposition.decomposition_rule import register_condition, register_resources
 from pennylane.decomposition.resources import resource_rep
-from pennylane.decomposition.utils import DecompositionNotApplicable
 from pennylane.math.decomposition import (
     xyx_rotation_angles,
     xzx_rotation_angles,
@@ -238,41 +237,20 @@ def two_qubit_decomposition(U, wires):
     return q.queue
 
 
-class OneQubitUnitaryDecomposition(DecompositionRule):  # pylint: disable=too-few-public-methods
-    """Wrapper around naive one-qubit decomposition rules that adds a global phase.
+def make_one_qubit_unitary_decomposition(su2_rule, su2_resource):
+    """Wrapper around a naive one-qubit decomposition rule that adds a global phase."""
 
-    Args:
-        su2_rule (callable): A function that implements the naive decomposition rule which
-            assumes that the unitary is SU(2)
-        su2_resource (callable): A function that returns the resources required by the naive
-            decomposition rule, without the GlobalPhase.
-
-    """
+    def _resource_fn(num_wires):  # pylint: disable=unused-argument
+        return su2_resource() | {ops.GlobalPhase: 1}
 
-    def __init__(self, su2_rule, su2_resource):
-        self._naive_rule = su2_rule
-        self._naive_resources = su2_resource
-        super().__init__(self._get_impl(), self._get_resource_fn())
-
-    def _get_impl(self):
-        """The implementation of the decomposition rule."""
-
-        def _impl(U, wires, **__):
-            U, global_phase = math.convert_to_su2(U, return_global_phase=True)
-            self._naive_rule(U, wires=wires)
-            ops.GlobalPhase(-global_phase)
+    @register_condition(lambda num_wires: num_wires == 1)
+    @register_resources(_resource_fn)
+    def _impl(U, wires, **__):
+        U, global_phase = math.convert_to_su2(U, return_global_phase=True)
+        su2_rule(U, wires=wires)
+        ops.cond(math.logical_not(math.allclose(global_phase, 0)), _global_phase)(global_phase)
 
-        return _impl
-
-    def _get_resource_fn(self):
-        """The resource function of the decomposition rule."""
-
-        def _resource_fn(num_wires):
-            if num_wires != 1:
-                raise DecompositionNotApplicable
-            return self._naive_resources() | {ops.GlobalPhase: 1}
-
-        return _resource_fn
+    return _impl
 
 
 def _su2_rot_resource():
@@ -336,11 +314,11 @@ def _su2_zxz_decomp(U, wires, **__):
     ops.RZ(omega, wires=wires[0])
 
 
-rot_decomp_rule = OneQubitUnitaryDecomposition(_su2_rot_decomp, _su2_rot_resource)
-zyz_decomp_rule = OneQubitUnitaryDecomposition(_su2_zyz_decomp, _su2_zyz_resource)
-xyx_decomp_rule = OneQubitUnitaryDecomposition(_su2_xyx_decomp, _su2_xyx_resource)
-xzx_decomp_rule = OneQubitUnitaryDecomposition(_su2_xzx_decomp, _su2_xzx_resource)
-zxz_decomp_rule = OneQubitUnitaryDecomposition(_su2_zxz_decomp, _su2_zxz_resource)
+rot_decomp_rule = make_one_qubit_unitary_decomposition(_su2_rot_decomp, _su2_rot_resource)
+zyz_decomp_rule = make_one_qubit_unitary_decomposition(_su2_zyz_decomp, _su2_zyz_resource)
+xyx_decomp_rule = make_one_qubit_unitary_decomposition(_su2_xyx_decomp, _su2_xyx_resource)
+xzx_decomp_rule = make_one_qubit_unitary_decomposition(_su2_xzx_decomp, _su2_xzx_resource)
+zxz_decomp_rule = make_one_qubit_unitary_decomposition(_su2_zxz_decomp, _su2_zxz_resource)
 
 ###################################################################################
 # Developer notes:
@@ -726,11 +704,8 @@ def _decompose_3_cnots(U, wires, initial_phase):
     return math.cast_like(-np.pi / 4, initial_phase)
 
 
-def _two_qubit_resource(num_wires):
+def _two_qubit_resource(**_):
     """A worst-case over-estimate for the resources of two-qubit unitary decomposition."""
-    if num_wires != 2:
-        raise DecompositionNotApplicable
-    # Assume the 3-CNOT case.
     return {
         resource_rep(ops.QubitUnitary, num_wires=1): 4,
         ops.CNOT: 3,
@@ -743,6 +718,7 @@ def _two_qubit_resource(num_wires):
     }
 
 
+@register_condition(lambda num_wires: num_wires == 2)
 @register_resources(_two_qubit_resource)
 def two_qubit_decomp_rule(U, wires, **__):
     """The decomposition rule for a two-qubit unitary."""
@@ -760,3 +736,7 @@ def two_qubit_decomp_rule(U, wires, **__):
     )(U, wires, initial_phase)
     total_phase = initial_phase + additional_phase
     ops.GlobalPhase(-total_phase)
+
+
+def _global_phase(phase):
+    ops.GlobalPhase(-phase)
diff --git a/pennylane/templates/subroutines/fable.py b/pennylane/templates/subroutines/fable.py
index 490b1d1c9..7846cecf8 100644
--- a/pennylane/templates/subroutines/fable.py
+++ b/pennylane/templates/subroutines/fable.py
@@ -137,7 +137,6 @@ class FABLE(Operation):
         Returns:
             list[.Operator]: list of gates for efficient circuit
         """
-        op_list = []
         alphas = qml.math.arccos(input_matrix).flatten()
         thetas = compute_theta(alphas)
 
@@ -150,8 +149,7 @@ class FABLE(Operation):
 
         wire_map = dict(enumerate(wires_j + wires_i))
 
-        for w in wires_i:
-            op_list.append(qml.Hadamard(w))
+        op_list = [qml.Hadamard(w) for w in wires_i]
 
         nots = {}
         for theta, control_index in zip(thetas, control_wires):
diff --git a/pennylane/templates/subroutines/kupccgsd.py b/pennylane/templates/subroutines/kupccgsd.py
index 0095ac0a9..63fb2a425 100644
--- a/pennylane/templates/subroutines/kupccgsd.py
+++ b/pennylane/templates/subroutines/kupccgsd.py
@@ -16,6 +16,7 @@ Contains the k-UpCCGSD template.
 """
 # pylint: disable-msg=too-many-branches,too-many-arguments,protected-access
 import copy
+from itertools import product
 
 import numpy as np
 
@@ -31,17 +32,15 @@ def generalized_singles(wires, delta_sz):
         \hat{T_1} = \sum_{pq} t_{p}^{q} \hat{c}^{\dagger}_{q} \hat{c}_{p}
 
     """
-    sz = np.array(
-        [0.5 if (i % 2 == 0) else -0.5 for i in range(len(wires))]
-    )  # alpha-beta electrons
+    n = len(wires)
+    sz = 0.5 * (-1) ** np.arange(n)  # alpha-beta electrons
     gen_singles_wires = []
-    for r in range(len(wires)):
-        for p in range(len(wires)):
-            if sz[p] - sz[r] == delta_sz and p != r:
-                if r < p:
-                    gen_singles_wires.append(wires[r : p + 1])
-                else:
-                    gen_singles_wires.append(wires[p : r + 1][::-1])
+    for r, p in product(range(n), repeat=2):
+        if sz[p] - sz[r] == delta_sz and p != r:
+            if r < p:
+                gen_singles_wires.append(wires[r : p + 1])
+            else:
+                gen_singles_wires.append(wires[p : r + 1][::-1])
     return gen_singles_wires
 
 
@@ -54,12 +53,9 @@ def generalized_pair_doubles(wires):
 
     """
     pair_gen_doubles_wires = [
-        [
-            wires[r : r + 2],
-            wires[p : p + 2],
-        ]  # wires for [wires[r], wires[r+1], wires[p], wires[p+1]] terms
-        for r in range(0, len(wires) - 1, 2)
-        for p in range(0, len(wires) - 1, 2)
+        # wires for [wires[r], wires[r+1], wires[p], wires[p+1]] terms
+        [wires[r : r + 2], wires[p : p + 2]]
+        for r, p in product(range(0, len(wires) - 1, 2), repeat=2)
         if p != r  # remove redundant terms
     ]
     return pair_gen_doubles_wires
@@ -125,14 +121,14 @@ class kUpCCGSD(Operation):
             # Build the electronic Hamiltonian
             symbols = ["H", "H"]
             coordinates = np.array([0.0, 0.0, -0.6614, 0.0, 0.0, 0.6614])
-            H, qubits = qml.qchem.molecular_hamiltonian(symbols, coordinates)
+            H, wires = qml.qchem.molecular_hamiltonian(symbols, coordinates)
 
             # Define the Hartree-Fock state
             electrons = 2
-            ref_state = qml.qchem.hf_state(electrons, qubits)
+            ref_state = qml.qchem.hf_state(electrons, wires)
 
             # Define the device
-            dev = qml.device('default.qubit', wires=qubits)
+            dev = qml.device('default.qubit', wires=wires)
 
             # Define the ansatz
             @qml.qnode(dev)
@@ -144,7 +140,7 @@ class kUpCCGSD(Operation):
             # Get the shape of the weights for this template
             layers = 1
             shape = qml.kUpCCGSD.shape(k=layers,
-                                n_wires=qubits, delta_sz=0)
+                                n_wires=wires, delta_sz=0)
 
             # Initialize the weight tensors
             np.random.seed(24)
@@ -230,12 +226,9 @@ class kUpCCGSD(Operation):
         d_wires = generalized_pair_doubles(list(wires))
 
         shape = qml.math.shape(weights)
-        if shape != (
-            k,
-            len(s_wires) + len(d_wires),
-        ):
+        if shape != (k, len(s_wires) + len(d_wires)):
             raise ValueError(
-                f"Weights tensor must be of shape {(k, len(s_wires) + len(d_wires),)}; got {shape}."
+                f"Weights tensor must be of shape {(k, len(s_wires) + len(d_wires))}; got {shape}."
             )
 
         init_state = qml.math.toarray(init_state)
@@ -319,7 +312,7 @@ class kUpCCGSD(Operation):
         r"""Returns the shape of the weight tensor required for this template.
         Args:
             k (int): Number of layers
-            n_wires (int): Number of qubits
+            n_wires (int): Number of wires
             delta_sz (int): Specifies the selection rules ``sz[p] - sz[r] = delta_sz``
             for the spin-projection ``sz`` of the orbitals involved in the single excitations.
             ``delta_sz`` can take the values :math:`0` and :math:`\pm 1`.
@@ -329,12 +322,12 @@ class kUpCCGSD(Operation):
 
         if n_wires < 4:
             raise ValueError(
-                f"This template requires the number of qubits to be greater than four; got 'n_wires' = {n_wires}"
+                f"This template requires the number of wires to be greater than four; got 'n_wires' = {n_wires}"
             )
 
         if n_wires % 2:
             raise ValueError(
-                f"This template requires an even number of qubits; got 'n_wires' = {n_wires}"
+                f"This template requires an even number of wires; got 'n_wires' = {n_wires}"
             )
 
         s_wires = generalized_singles(range(n_wires), delta_sz)


Test patch:
diff --git a/tests/decomposition/test_decomposition_graph.py b/tests/decomposition/test_decomposition_graph.py
index 67c777464..63b948832 100644
--- a/tests/decomposition/test_decomposition_graph.py
+++ b/tests/decomposition/test_decomposition_graph.py
@@ -26,7 +26,6 @@ import pennylane as qml
 from pennylane.decomposition import (
     DecompositionError,
     DecompositionGraph,
-    DecompositionNotApplicable,
     adjoint_resource_rep,
     controlled_resource_rep,
     pow_resource_rep,
@@ -99,7 +98,7 @@ class TestDecompositionGraph:
         assert len(graph2._graph.edges()) == 8
 
     def test_graph_construction_non_applicable_rules(self, _):
-        """Tests rules that raise DecompositionNotApplicable are skipped."""
+        """Tests rules which are not applicable are skipped."""
 
         class CustomOp(qml.operation.Operation):  # pylint: disable=too-few-public-methods
             """A custom op"""
@@ -110,20 +109,15 @@ class TestDecompositionGraph:
             def resource_params(self):
                 return {"num_wires": len(self.wires)}
 
-        def _some_resource(num_wires):
-            if num_wires > 1:
-                raise DecompositionNotApplicable
-            return {qml.RZ: 1, qml.CNOT: 1}
-
-        @qml.register_resources(_some_resource)
+        @qml.register_condition(lambda num_wires: num_wires == 1)
+        @qml.register_resources({qml.RZ: 1, qml.CNOT: 1})
         def some_rule(*_, **__):
             raise NotImplementedError
 
         def _some_other_resource(num_wires):
-            if num_wires < 2:
-                raise DecompositionNotApplicable
             return {qml.RZ: 1, qml.CNOT: num_wires - 1}
 
+        @qml.register_condition(lambda num_wires: num_wires >= 2)
         @qml.register_resources(_some_other_resource)
         def some_other_rule(*_, **__):
             raise NotImplementedError
diff --git a/tests/decomposition/test_decomposition_rule.py b/tests/decomposition/test_decomposition_rule.py
index 22b251fe0..6865313b3 100644
--- a/tests/decomposition/test_decomposition_rule.py
+++ b/tests/decomposition/test_decomposition_rule.py
@@ -100,6 +100,57 @@ class TestDecompositionRule:
             gate_counts={CompressedResourceOp(qml.RZ): 1, CompressedResourceOp(qml.CNOT): 4}
         )
 
+    def test_decomposition_condition(self):
+        """Tests that the register_condition works."""
+
+        @qml.register_resources({qml.H: 2, qml.Toffoli: 1})
+        @qml.register_condition(lambda num_wires: num_wires == 3)
+        def rule_1(wires, **__):
+            raise NotImplementedError
+
+        assert isinstance(rule_1, DecompositionRule)
+        assert rule_1.is_applicable(num_wires=3)
+        assert not rule_1.is_applicable(num_wires=2)
+        assert rule_1.compute_resources(num_wires=3) == Resources(
+            {
+                CompressedResourceOp(qml.H): 2,
+                CompressedResourceOp(qml.Toffoli): 1,
+            }
+        )
+
+        @qml.register_condition(lambda num_wires: num_wires == 3)
+        @qml.register_resources({qml.H: 2, qml.Toffoli: 1})
+        def rule_2(wires, **__):
+            raise NotImplementedError
+
+        assert isinstance(rule_2, DecompositionRule)
+        assert rule_2.is_applicable(num_wires=3)
+        assert not rule_2.is_applicable(num_wires=2)
+        assert rule_2.compute_resources(num_wires=3) == Resources(
+            {
+                CompressedResourceOp(qml.H): 2,
+                CompressedResourceOp(qml.Toffoli): 1,
+            }
+        )
+
+        def _resource_fn(**_):
+            return {qml.H: 2, qml.Toffoli: 1}
+
+        @qml.register_resources(_resource_fn)
+        @qml.register_condition(lambda num_wires: num_wires == 3)
+        def rule_3(wires, **__):
+            raise NotImplementedError
+
+        assert isinstance(rule_3, DecompositionRule)
+        assert rule_3.is_applicable(num_wires=3)
+        assert not rule_3.is_applicable(num_wires=2)
+        assert rule_3.compute_resources(num_wires=3) == Resources(
+            {
+                CompressedResourceOp(qml.H): 2,
+                CompressedResourceOp(qml.Toffoli): 1,
+            }
+        )
+
     def test_inspect_decomposition_rule(self):
         """Tests that the source code for a decomposition rule can be inspected."""
 
diff --git a/tests/templates/test_subroutines/test_fable.py b/tests/templates/test_subroutines/test_fable.py
index d2ba5f249..6dc85a6a8 100644
--- a/tests/templates/test_subroutines/test_fable.py
+++ b/tests/templates/test_subroutines/test_fable.py
@@ -44,10 +44,7 @@ class TestFable:
     @pytest.mark.parametrize(
         ("input", "wires"),
         [
-            (
-                np.random.random((2, 2)),
-                3,
-            ),
+            (np.random.random((2, 2)), 3),
             (
                 np.array(
                     [
@@ -59,17 +56,10 @@ class TestFable:
                 ),
                 5,
             ),
-            (
-                np.random.random((8, 8)),
-                7,
-            ),
-            (
-                np.random.random((16, 16)),
-                9,
-            ),
+            (np.random.random((8, 8)), 7),
         ],
     )
-    def test_fable_real_for_variety_of_input_matrices(self, input, wires):
+    def test_fable_real_input_matrices(self, input, wires):
         """Test that FABLE produces the right circuit given a square, real-valued matrix"""
         dev = qml.device("default.qubit")
 
@@ -78,10 +68,8 @@ class TestFable:
             qml.FABLE(input_matrix=input, wires=range(wires), tol=0)
             return qml.state()
 
-        expected = (
-            len(input)
-            * qml.matrix(circuit, wire_order=range(wires))().real[0 : len(input), 0 : len(input)]
-        )
+        dim = len(input)
+        expected = dim * qml.matrix(circuit, wire_order=range(wires))().real[:dim, :dim]
         assert np.allclose(input, expected)
 
     def test_fable_imaginary_error(self, input_matrix):
@@ -391,89 +379,36 @@ class TestFable:
     @pytest.mark.parametrize(
         ("input", "wires"),
         [
-            (
-                np.random.random((1, 2)),
-                3,
-            ),
-            (
-                np.random.random((1, 1)),
-                3,
-            ),
-            (
-                np.random.random((2, 1)),
-                3,
-            ),
-            (
-                np.random.random((3, 2)),
-                5,
-            ),
-            (
-                np.random.random((4, 2)),
-                5,
-            ),
-            (
-                np.random.random((2, 3)),
-                5,
-            ),
-            (
-                np.random.random((2, 4)),
-                5,
-            ),
-            (
-                np.random.random((3, 4)),
-                5,
-            ),
-            (
-                np.random.random((4, 3)),
-                5,
-            ),
-            (
-                np.random.random((3, 5)),
-                7,
-            ),
-            (
-                np.random.random((3, 6)),
-                7,
-            ),
-            (
-                np.random.random((3, 7)),
-                7,
-            ),
-            (
-                np.random.random((4, 5)),
-                7,
-            ),
-            (
-                np.random.random((5, 5)),
-                7,
-            ),
-            (
-                np.random.random((6, 5)),
-                7,
-            ),
-            (
-                np.random.random((2, 9)),
-                9,
-            ),
-            (
-                np.random.random((9, 3)),
-                9,
-            ),
+            (np.random.random((1, 2)), 3),
+            (np.random.random((1, 1)), 3),
+            (np.random.random((2, 1)), 3),
+            (np.random.random((3, 2)), 5),
+            (np.random.random((4, 2)), 5),
+            (np.random.random((2, 3)), 5),
+            (np.random.random((2, 4)), 5),
+            (np.random.random((3, 4)), 5),
+            (np.random.random((4, 3)), 5),
+            (np.random.random((3, 5)), 7),
+            (np.random.random((3, 6)), 7),
+            (np.random.random((3, 7)), 7),
+            (np.random.random((4, 5)), 7),
+            (np.random.random((5, 5)), 7),
+            (np.random.random((6, 5)), 7),
         ],
     )
     def test_variety_of_matrix_shapes(self, input, wires):
-        """Test that FABLE runs without error for variety of cases."""
+        """Test that FABLE runs without error for a variety of input shapes."""
         dev = qml.device("default.qubit")
-        s = int(qml.math.ceil(qml.math.log2(max(len(input), len(input[0])))))
-        dimension = s**2
+        s = int(qml.math.ceil(qml.math.log2(max(input.shape))))
+        s = max(s, 1)
+        dim = 2**s
 
         @qml.qnode(dev)
         def circuit():
             qml.FABLE(input_matrix=input, wires=range(wires), tol=0)
             return qml.state()
 
-        expected = (
-            dimension
-            * qml.matrix(circuit, wire_order=range(wires))().real[0:dimension, 0:dimension]
-        )
-        assert qml.math.shape(expected) == (dimension, dimension)
+        circuit_mat = dim * qml.matrix(circuit, wire_order=range(wires))().real[:dim, :dim]
+        # Test that the matrix was encoded up to a constant
+        submat = circuit_mat[: input.shape[0], : input.shape[1]]
+        assert np.allclose(input, submat)
diff --git a/tests/templates/test_subroutines/test_kupccgsd.py b/tests/templates/test_subroutines/test_kupccgsd.py
index 98599f443..a83e3975f 100644
--- a/tests/templates/test_subroutines/test_kupccgsd.py
+++ b/tests/templates/test_subroutines/test_kupccgsd.py
@@ -54,7 +54,8 @@ def test_standard_validity(k, delta_sz, init_state, wires):
 
     op = qml.kUpCCGSD(weights, wires=wires, k=k, delta_sz=delta_sz, init_state=init_state)
 
-    qml.ops.functions.assert_valid(op)
+    skip_diff = len(wires) > 4
+    qml.ops.functions.assert_valid(op, skip_differentiation=skip_diff)
 
 
 class TestDecomposition:
@@ -160,7 +161,7 @@ class TestDecomposition:
         assert np.allclose(state1, state2, atol=tol, rtol=0)
 
     @pytest.mark.parametrize(
-        ("num_qubits", "k", "exp_state"),
+        ("num_wires", "k", "exp_state"),
         [
             (
                 4,
@@ -243,17 +244,17 @@ class TestDecomposition:
             ),
         ],
     )
-    def test_k_layers_upccgsd(self, num_qubits, k, exp_state, tol):
+    def test_k_layers_upccgsd(self, num_wires, k, exp_state, tol):
         """Test that the k-UpCCGSD template with multiple layers works correctly asserting the prepared state."""
 
-        wires = range(num_qubits)
+        wires = range(num_wires)
 
-        shape = qml.kUpCCGSD.shape(k=k, n_wires=num_qubits, delta_sz=0)
+        shape = qml.kUpCCGSD.shape(k=k, n_wires=num_wires, delta_sz=0)
         weight = np.pi / 2 * qml.math.ones(shape)
 
         dev = qml.device("default.qubit", wires=wires)
 
-        init_state = qml.math.array([1 if x < num_qubits // 2 else 0 for x in wires])
+        init_state = qml.math.array([1 if x < num_wires // 2 else 0 for x in wires])
 
         @qml.qnode(dev)
         def circuit(weight):
@@ -465,18 +466,18 @@ class TestAttributes:
         shape = qml.kUpCCGSD.shape(k, n_wires, delta_sz)
         assert shape == expected_shape
 
-    def test_shape_exception_not_enough_qubits(self):
-        """Test that the shape function warns if there are not enough qubits."""
+    def test_shape_exception_not_enough_wires(self):
+        """Test that the shape function warns if there are not enough wires."""
 
         with pytest.raises(
-            ValueError, match="This template requires the number of qubits to be greater than four"
+            ValueError, match="This template requires the number of wires to be greater than four"
         ):
             qml.kUpCCGSD.shape(k=2, n_wires=1, delta_sz=0)
 
-    def test_shape_exception_not_even_qubits(self):
-        """Test that the shape function warns if the number of qubits are not even."""
+    def test_shape_exception_not_even_wires(self):
+        """Test that the shape function warns if the number of wires are not even."""
 
-        with pytest.raises(ValueError, match="This template requires an even number of qubits"):
+        with pytest.raises(ValueError, match="This template requires an even number of wires"):
             qml.kUpCCGSD.shape(k=2, n_wires=5, delta_sz=0)
 
 
diff --git a/tests/templates/test_subroutines/test_trotter.py b/tests/templates/test_subroutines/test_trotter.py
index 4787c974a..82143dfad 100644
--- a/tests/templates/test_subroutines/test_trotter.py
+++ b/tests/templates/test_subroutines/test_trotter.py
@@ -1726,8 +1726,6 @@ class TestTrotterizedQfuncIntegration:
                 qml.ControlledPhaseShift(time * arg2, wires=[wires[1], wires[0]]),
                 qml.CNOT([wires[0], wires[1]]),
                 qml.QFT(wires=wires[1:-1]),
-                qml.QFT(wires=wires[1:-1]),
-                qml.QFT(wires=wires[1:-1]),
             ]
 
             if reverse:
@@ -1742,8 +1740,6 @@ class TestTrotterizedQfuncIntegration:
                 qml.ControlledPhaseShift((time / 2) * arg2, wires=[wires[1], wires[0]]),
                 qml.CNOT([wires[0], wires[1]]),
                 qml.QFT(wires=wires[1:-1]),
-                qml.QFT(wires=wires[1:-1]),
-                qml.QFT(wires=wires[1:-1]),
             ]
 
             if reverse:
@@ -1760,8 +1756,6 @@ class TestTrotterizedQfuncIntegration:
                 qml.ControlledPhaseShift((p_4 * time / 2) * arg2, wires=[wires[1], wires[0]]),
                 qml.CNOT([wires[0], wires[1]]),
                 qml.QFT(wires=wires[1:-1]),
-                qml.QFT(wires=wires[1:-1]),
-                qml.QFT(wires=wires[1:-1]),
             ]
 
             expected_decomp1 = (
@@ -1778,8 +1772,6 @@ class TestTrotterizedQfuncIntegration:
                 qml.ControlledPhaseShift((p4_comp * time / 2) * arg2, wires=[wires[1], wires[0]]),
                 qml.CNOT([wires[0], wires[1]]),
                 qml.QFT(wires=wires[1:-1]),
-                qml.QFT(wires=wires[1:-1]),
-                qml.QFT(wires=wires[1:-1]),
             ]
 
             expected_decomp2 = (
@@ -1803,7 +1795,7 @@ class TestTrotterizedQfuncIntegration:
         arg1 = 2.34
         arg2 = -6.78
         args = (arg1, arg2)
-        kwargs = {"kwarg1": True, "kwarg2": 3}
+        kwargs = {"kwarg1": True, "kwarg2": 1}
 
         def my_qfunc(time, arg1, arg2, wires, kwarg1=False, kwarg2=None):
             """Arbitrarily complex qfunc"""
@@ -1867,7 +1859,7 @@ class TestTrotterizedQfuncIntegration:
         arg1 = jnp.array(2.34)
         arg2 = jnp.array(-6.78)
         args = (arg1, arg2)
-        kwargs = {"kwarg1": True, "kwarg2": 3}
+        kwargs = {"kwarg1": True, "kwarg2": 1}
 
         def my_qfunc(time, arg1, arg2, wires, kwarg1=False, kwarg2=None):
             """Arbitrarily complex qfunc"""
@@ -1935,7 +1927,7 @@ class TestTrotterizedQfuncIntegration:
         arg1 = jnp.array(2.34)
         arg2 = jnp.array(-6.78)
         args = (arg1, arg2)
-        kwargs = {"kwarg1": True, "kwarg2": 3}
+        kwargs = {"kwarg1": True, "kwarg2": 1}
 
         def my_qfunc(time, arg1, arg2, wires, kwarg1=False, kwarg2=None):
             """Arbitrarily complex qfunc"""
@@ -2000,7 +1992,7 @@ class TestTrotterizedQfuncIntegration:
         wires = ["aux1", "aux2", 0, 1, "target"]
         arg1 = qnp.array(2.34)
         arg2 = qnp.array(-6.78)
-        kwargs = {"kwarg1": True, "kwarg2": 3}
+        kwargs = {"kwarg1": True, "kwarg2": 1}
 
         def my_qfunc(time, arg1, arg2, wires, kwarg1=False, kwarg2=None):
             """Arbitrarily complex qfunc"""
@@ -2070,7 +2062,7 @@ class TestTrotterizedQfuncIntegration:
         arg1 = jnp.array(2.34)
         arg2 = jnp.array(-6.78)
         args = (arg1, arg2)
-        kwargs = {"kwarg1": True, "kwarg2": 3}
+        kwargs = {"kwarg1": True, "kwarg2": 1}
 
         def my_qfunc(time, arg1, arg2, wires, kwarg1=False, kwarg2=None):
             """Arbitrarily complex qfunc"""