Repository: PennyLaneAI/pennylane
Pull request #6952 (2025-02-11 20:20:10)

[Decomposition] Graph-based Decomposition with Dijkstra's Algorithm
**Description of the Change:**
Implements the `DecompositionGraph` that solves for an optimal decomposition.

**Related GitHub Issues:**
[sc-84329]

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6952?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 `26.36364%` with `81 lines` in your changes missing coverage. Please review.
> Project coverage is 99.23%. Comparing base [(`5cd437c`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/5cd437c158c0f5e2e47323c56486296b8b136502?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`8f81554`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/8f8155410d3c30a12c221749c362d8ed9cf2349c?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/6952?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.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6952?src=pr&el=tree&filepath=pennylane%2Fdecomposition%2Fdecomposition.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RlY29tcG9zaXRpb24vZGVjb21wb3NpdGlvbi5weQ==) | 24.29% | [81 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6952?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                 @@
##           decomp-some-ops    #6952      +/-   ##
===================================================
- Coverage            99.41%   99.23%   -0.18%     
===================================================
  Files                  483      484       +1     
  Lines                45704    45813     +109     
===================================================
+ Hits                 45436    45464      +28     
- Misses                 268      349      +81     
```

</details>

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

astralcai: @maliasadi This is ready for review
astralcai: > Robust repeated use of DecompositionGraph.solve on the same instance of a graph. If that does not make sense/is not supported (which does not sound crazy to me, given that there is only the lazy kwarg to solve), an explicit error should be raised and tested for :)

The solution to the graph traversal problem is stored in the `_visitor` attribute of the graph. Every time you call `solve`, it resets the `_visitor` to a fresh new instance, so you would just be doing the same thing again. 
astralcai: > use of lazy=False in DecompositionGraph.solve, optimally in a way so that it makes a difference for an op that is being explored, or not explored, depending on lazy kwarg.

Test added.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index ad8c0019d..f14b2587f 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -177,6 +177,10 @@
 * Decomposition rules are implemented using the new infrastructure for most PennyLane operators excluding templates.
   [(#6951)](https://github.com/PennyLaneAI/pennylane/pull/6951)
 
+* A graph-based decomposition solver has been implemented that solves for the optimal decomposition
+  rule to use for an operator towards a target gate set.
+  [(#6952)](https://github.com/PennyLaneAI/pennylane/pull/6952)
+
 <h3>Improvements 🛠</h3>
   
 * The decompositions of `qml.SX`, `qml.X` and `qml.Y` use `qml.GlobalPhase` instead of `qml.PhaseShift`.
diff --git a/pennylane/decomposition/__init__.py b/pennylane/decomposition/__init__.py
index 1360cedad..538768ac7 100644
--- a/pennylane/decomposition/__init__.py
+++ b/pennylane/decomposition/__init__.py
@@ -71,8 +71,32 @@ a list of known decomposition rules for a given operator.
 Graph-based Decomposition Solver
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-.. TODO::
-    Add section here explaining the decomposition graph [sc-84329]
+.. autosummary::
+    :toctree: api
+
+    ~DecompositionGraph
+
+The decomposition graph is a directed graph of operators and decomposition rules. Dijkstra's
+algorithm is used to explore the graph and find the most efficient decomposition of a given
+operator towards a target gate set.
+
+.. code-block:: python
+
+    op = qml.CRX(0.5, wires=[0, 1])
+    graph = DecompositionGraph(
+        operations=[op],
+        target_gate_set={"RZ", "RX", "CNOT", "GlobalPhase"},
+    )
+    graph.solve()
+
+.. code-block:: pycon
+
+    >>> with qml.queuing.AnnotatedQueue() as q:
+    ...     graph.decomposition(op)(0.5, wires=[0, 1])
+    >>> q.queue
+    [H(1), CRZ(0.5, wires=Wires([0, 1])), H(1)]
+    >>> graph.resource_estimate(op)
+    <num_gates=14, gate_counts={RZ: 6, GlobalPhase: 4, RX: 2, CNOT: 2}>
 
 Integration with the Decompose Transform
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -82,6 +106,7 @@ Integration with the Decompose Transform
 
 """
 
+from .decomposition_graph import DecompositionGraph
 from .resources import (
     Resources,
     # TODO: add CompressedResourceOp once the conflict with labs is resolved.
diff --git a/pennylane/decomposition/decomposition_graph.py b/pennylane/decomposition/decomposition_graph.py
new file mode 100644
index 000000000..e75e243ab
--- /dev/null
+++ b/pennylane/decomposition/decomposition_graph.py
@@ -0,0 +1,351 @@
+# Copyright 2025 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.
+
+"""Implements the class DecompositionGraph
+
+This module implements a graph-based decomposition algorithm that constructs a graph of operators
+connected by decomposition rules, and then traverses it using Dijkstra's algorithm to find the best
+decomposition for every operator.
+
+The architecture of this module utilizes design patterns similar to those present in Qiskit's
+implementation of the basis translator, the Boost Graph library, and RustworkX.
+
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import rustworkx as rx
+from rustworkx.visit import DijkstraVisitor, PruneSearch, StopSearch
+
+from pennylane.operation import Operator
+
+from .decomposition_rule import DecompositionRule, list_decomps
+from .resources import CompressedResourceOp, Resources, resource_rep
+
+
+class DecompositionError(Exception):
+    """Base class for decomposition errors."""
+
+
+@dataclass(frozen=True)
+class _DecompositionNode:
+    """A node that represents a decomposition rule."""
+
+    rule: DecompositionRule
+    resource_decomp: Resources
+
+    def count(self, op: CompressedResourceOp):
+        """Find the number of occurrences of an operator in the decomposition."""
+        return self.resource_decomp.gate_counts.get(op, 0)
+
+
+class DecompositionGraph:  # pylint: disable=too-many-instance-attributes
+    """A graph that models a decomposition problem.
+
+    The decomposition graph contains two types of nodes: operator nodes and decomposition nodes.
+    Each decomposition node is a :class:`pennylane.decomposition.DecompositionRule`, and each
+    operator node is a :class:`~pennylane.decomposition.resources.CompressedResourceOp` which
+    contains an operator type and any additional parameters that affects the resource requirements
+    of the operator. Essentially, two instances of the same operator type is represented by the
+    same node in the graph if they're expected to have the same decompositions.
+
+    There are also two types of directed edges: edges that connect operators to the decomposition
+    rules that contain them, and edges that connect decomposition rules to the operators that they
+    decompose. The edge weights represent the difference in the gate count between the two states.
+    Edges that connects decomposition rules to operators have a weight of 0 because an operator can
+    be replaced with its decomposition at no additional cost.
+
+    On the other hand, edges that connect operators to the decomposition rule that contains them
+    will have a weight that is the total resource estimate of the decomposition minus the resource
+    estimate of the operator. For example the edge that connects a ``qml.CNOT`` to the following
+    decomposition rule:
+
+    .. code-block:: python
+
+        import pennylane as qml
+
+        @qml.register_resources({qml.H: 2, qml.CNOT: 1})
+        def my_cz(wires):
+            qml.H(wires=wires[1])
+            qml.CNOT(wires=wires)
+            qml.H(wires=wires[1])
+
+    will have a weight of 2, because the decomposition rule contains 2 additional ``H`` gates.
+    Note that this gate count is in terms of gates in the target gate set. If ``H`` isn't supported
+    and is in turn decomposed to two ``RZ`` gates and a ``RX`` gate, the weight of this edge
+    becomes 2 * 3 = 6. This way, the total distance from a basis gate to a high-level gate is
+    conveniently the total number of basis gates required to decompose this high-level gate, which
+    allows us to use Dijkstra's algorithm to find the most efficient decomposition.
+
+    Args:
+        operations (list[Operator]): The list of operations to find decompositions for.
+        target_gate_set (set[str]): The names of the gates in the target gate set.
+        fixed_decomps (dict): A dictionary mapping operator names to fixed decompositions.
+        alt_decomps (dict): A dictionary mapping operator names to alternative decompositions.
+
+    **Example**
+
+    .. code-block:: python
+
+        op = qml.CRX(0.5, wires=[0, 1])
+        graph = DecompositionGraph(
+            operations=[op],
+            target_gate_set={"RZ", "RX", "CNOT", "GlobalPhase"},
+        )
+        graph.solve()
+
+    >>> with qml.queuing.AnnotatedQueue() as q:
+    ...     graph.decomposition(op)(0.5, wires=[0, 1])
+    >>> q.queue
+    [H(1), CRZ(0.5, wires=Wires([0, 1])), H(1)]
+    >>> graph.resource_estimate(op)
+    num_gates=14, gate_counts={RZ: 6, GlobalPhase: 4, RX: 2, CNOT: 2}
+
+    """
+
+    def __init__(
+        self,
+        operations,
+        target_gate_set: set[str],
+        fixed_decomps: dict = None,
+        alt_decomps: dict = None,
+    ):
+        self._original_ops = operations
+        self._target_gate_set = target_gate_set
+        self._original_ops_indices: set[int] = set()
+        self._target_gate_indices: set[int] = set()
+        self._graph = rx.PyDiGraph()
+        self._op_node_indices: dict[CompressedResourceOp, int] = {}
+        self._fixed_decomps = fixed_decomps or {}
+        self._alt_decomps = alt_decomps or {}
+        self._construct_graph()
+        self._visitor = None
+
+    def _get_decompositions(self, op_type) -> list[DecompositionRule]:
+        """Helper function to get a list of decomposition rules."""
+        if op_type in self._fixed_decomps:
+            return [self._fixed_decomps[op_type]]
+        return self._alt_decomps.get(op_type, []) + list_decomps(op_type)
+
+    def _construct_graph(self):
+        """Constructs the decomposition graph."""
+        for op in self._original_ops:
+            op_node = resource_rep(type(op), **op.resource_params)
+            idx = self._recursively_add_op_node(op_node)
+            self._original_ops_indices.add(idx)
+
+    def _recursively_add_op_node(self, op_node: CompressedResourceOp) -> int:
+        """Recursively adds an operation node to the graph.
+
+        An operator node is uniquely defined by its operator type and resource parameters, which
+        are conveniently wrapped in a ``CompressedResourceOp``.
+
+        """
+
+        if op_node in self._op_node_indices:
+            return self._op_node_indices[op_node]
+
+        op_node_idx = self._graph.add_node(op_node)
+        self._op_node_indices[op_node] = op_node_idx
+        if op_node.op_type.__name__ in self._target_gate_set:
+            self._target_gate_indices.add(op_node_idx)
+            return op_node_idx
+
+        for decomposition in self._get_decompositions(op_node.op_type):
+            d_node_idx = self._recursively_add_decomposition_node(decomposition, op_node.params)
+            self._graph.add_edge(d_node_idx, op_node_idx, 0)
+
+        return op_node_idx
+
+    def _recursively_add_decomposition_node(self, rule: DecompositionRule, params: dict) -> int:
+        """Recursively adds a decomposition node to the graph.
+
+        A decomposition node is defined by a decomposition rule and a first-order resource estimate
+        of this decomposition as computed with resource params passed from the operator node.
+
+        """
+
+        decomp_resources = rule.compute_resources(**params)
+        d_node = _DecompositionNode(rule, decomp_resources)
+        d_node_idx = self._graph.add_node(d_node)
+        all_ops = [op for op, count in decomp_resources.gate_counts.items() if count > 0]
+        for op in all_ops:
+            op_node_idx = self._recursively_add_op_node(op)
+            self._graph.add_edge(op_node_idx, d_node_idx, (op_node_idx, d_node_idx))
+        return d_node_idx
+
+    def solve(self, lazy=True):
+        """Solves the graph using the Dijkstra search algorithm.
+
+        Args:
+            lazy (bool): If True, the Dijkstra search will stop once optimal decompositions are
+                found for all operations that the graph was initialized with. Otherwise, the
+                entire graph will be explored.
+
+        """
+        self._visitor = _DecompositionSearchVisitor(self._graph, self._original_ops_indices, lazy)
+        start = self._graph.add_node("dummy")
+        self._graph.add_edges_from(
+            [(start, op_node_idx, 1) for op_node_idx in self._target_gate_indices]
+        )
+        rx.dijkstra_search(
+            self._graph,
+            source=[start],
+            weight_fn=self._visitor.edge_weight,
+            visitor=self._visitor,
+        )
+        self._graph.remove_node(start)
+        if self._visitor.unsolved_op_indices:
+            unsolved_ops = [self._graph[op_idx] for op_idx in self._visitor.unsolved_op_indices]
+            op_names = set(op.op_type.__name__ for op in unsolved_ops)
+            raise DecompositionError(
+                f"Decomposition not found for {op_names} to the gate set {self._target_gate_set}"
+            )
+
+    def is_solved_for(self, op):
+        """Tests whether the decomposition graph is solved for a given operator."""
+        op_node = resource_rep(type(op), **op.resource_params)
+        return (
+            op_node in self._op_node_indices
+            and self._op_node_indices[op_node] in self._visitor.distances
+        )
+
+    def resource_estimate(self, op) -> Resources:
+        """Returns the resource estimate for a given operator.
+
+        Args:
+            op (Operator): The operator for which to return the resource estimates.
+
+        Returns:
+            Resources: The resource estimate.
+
+        **Example**
+
+        The resource estimate is a gate count in terms of the target gate set, not the immediate
+        set of gates that the operator decomposes to.
+
+        .. code-block:: python
+
+            op = qml.CRX(0.5, wires=[0, 1])
+            graph = DecompositionGraph(
+                operations=[op],
+                target_gate_set={"RZ", "RX", "CNOT", "GlobalPhase"},
+            )
+            graph.solve()
+
+        >>> with qml.queuing.AnnotatedQueue() as q:
+        ...     graph.decomposition(op)(0.5, wires=[0, 1])
+        >>> q.queue
+        [H(1), CRZ(0.5, wires=Wires([0, 1])), H(1)]
+        >>> graph.resource_estimate(op)
+        num_gates=14, gate_counts={RZ: 6, GlobalPhase: 4, RX: 2, CNOT: 2}
+
+        """
+        op_node = resource_rep(type(op), **op.resource_params)
+        op_node_idx = self._op_node_indices[op_node]
+        return self._visitor.distances[op_node_idx]
+
+    def decomposition(self, op: Operator) -> DecompositionRule:
+        """Returns the optimal decomposition rule for a given operator.
+
+        Args:
+            op (Operator): The operator for which to return the optimal decomposition.
+
+        Returns:
+            DecompositionRule: The optimal decomposition.
+
+        **Example**
+
+        The decomposition rule is a quantum function that takes ``(*op.parameters, wires=op.wires, **op.hyperparameters)``
+        as arguments.
+
+        .. code-block:: python
+
+            op = qml.CRX(0.5, wires=[0, 1])
+            graph = DecompositionGraph(
+                operations=[op],
+                target_gate_set={"RZ", "RX", "CNOT", "GlobalPhase"},
+            )
+            graph.solve()
+            rule = graph.decomposition(op)
+
+        >>> with qml.queuing.AnnotatedQueue() as q:
+        ...     rule(*op.parameters, wires=op.wires, **op.hyperparameters)
+        >>> q.queue
+        [H(1), CRZ(0.5, wires=Wires([0, 1])), H(1)]
+
+        """
+        op_node = resource_rep(type(op), **op.resource_params)
+        op_node_idx = self._op_node_indices[op_node]
+        d_node_idx = self._visitor.predecessors[op_node_idx]
+        return self._graph[d_node_idx].rule
+
+
+class _DecompositionSearchVisitor(DijkstraVisitor):
+    """The visitor used in the Dijkstra search for the optimal decomposition."""
+
+    def __init__(self, graph: rx.PyDiGraph, original_op_indices: set[int], lazy: bool = True):
+        self._graph = graph
+        self._lazy = lazy
+        # maps node indices to the optimal resource estimates
+        self.distances: dict[int, Resources] = {}
+        # maps operator nodes to the optimal decomposition nodes
+        self.predecessors: dict[int, int] = {}
+        self.unsolved_op_indices = original_op_indices.copy()
+        self._num_edges_examined: dict[int, int] = {}  # keys are decomposition node indices
+
+    def edge_weight(self, edge_obj):
+        """Calculates the weight of an edge."""
+        if not isinstance(edge_obj, tuple):
+            return float(edge_obj)
+        op_node_idx, d_node_idx = edge_obj
+        return self.distances[d_node_idx].num_gates - self.distances[op_node_idx].num_gates
+
+    def discover_vertex(self, v, _):
+        """Triggered when a vertex is about to be explored during the Dijkstra search."""
+        self.unsolved_op_indices.discard(v)
+        if not self.unsolved_op_indices and self._lazy:
+            raise StopSearch
+
+    def examine_edge(self, edge):
+        """Triggered when an edge is examined during the Dijkstra search."""
+        src_idx, target_idx, _ = edge
+        src_node = self._graph[src_idx]
+        target_node = self._graph[target_idx]
+        if not isinstance(target_node, _DecompositionNode):
+            return  # nothing is to be done for edges leading to an operator node
+        if target_idx not in self.distances:
+            self.distances[target_idx] = Resources()  # initialize with empty resource
+        self.distances[target_idx] += self.distances[src_idx] * target_node.count(src_node)
+        if target_idx not in self._num_edges_examined:
+            self._num_edges_examined[target_idx] = 0
+        self._num_edges_examined[target_idx] += 1
+        if self._num_edges_examined[target_idx] < len(target_node.resource_decomp.gate_counts):
+            # Typically in Dijkstra's search, a vertex is discovered from any of its incoming
+            # edges. However, for a decomposition node, it requires all incoming edges to be
+            # examined before it can be discovered (each incoming edge represents a different
+            # operator that this decomposition depends on).
+            raise PruneSearch
+
+    def edge_relaxed(self, edge):
+        """Triggered when an edge is relaxed during the Dijkstra search."""
+        src_idx, target_idx, _ = edge
+        target_node = self._graph[target_idx]
+        if self._graph[src_idx] == "dummy":
+            self.distances[target_idx] = Resources({target_node: 1})
+        elif isinstance(target_node, CompressedResourceOp):
+            self.predecessors[target_idx] = src_idx
+            self.distances[target_idx] = self.distances[src_idx]
diff --git a/pennylane/decomposition/resources.py b/pennylane/decomposition/resources.py
index 068ebfefc..e1a398498 100644
--- a/pennylane/decomposition/resources.py
+++ b/pennylane/decomposition/resources.py
@@ -143,6 +143,11 @@ def _validate_resource_rep(op_type, params):
     if not issubclass(op_type, qml.operation.Operator):
         raise TypeError(f"op_type must be a type of Operator, got {op_type}")
 
+    if not isinstance(op_type.resource_keys, set):
+        raise TypeError(
+            f"{op_type.__name__}.resource_keys must be a set, not a {type(op_type.resource_keys)}"
+        )
+
     unexpected_arguments = set(params.keys()) - op_type.resource_keys
     if unexpected_arguments:
         raise TypeError(
diff --git a/pennylane/ops/op_math/controlled_ops.py b/pennylane/ops/op_math/controlled_ops.py
index c56dc8987..8f00f9d59 100644
--- a/pennylane/ops/op_math/controlled_ops.py
+++ b/pennylane/ops/op_math/controlled_ops.py
@@ -631,11 +631,22 @@ def _cz_to_cps_resources():
 
 
 @register_resources(_cz_to_cps_resources)
-def _cz(wires: WiresLike, **__):
+def _cz_to_cps(wires: WiresLike, **__):
     qml.ControlledPhaseShift(np.pi, wires=wires)
 
 
-add_decomps(CZ, _cz)
+def _cz_to_cnot_resources():
+    return {qml.H: 2, qml.CNOT: 1}
+
+
+@register_resources(_cz_to_cnot_resources)
+def _cz_to_cnot(wires: WiresLike, **__):
+    qml.H(wires=wires[1])
+    qml.CNOT(wires=wires)
+    qml.H(wires=wires[1])
+
+
+add_decomps(CZ, _cz_to_cps, _cz_to_cnot)
 
 
 class CSWAP(ControlledOp):


Test patch:
diff --git a/tests/decomposition/conftest.py b/tests/decomposition/conftest.py
new file mode 100644
index 000000000..92f3daa7d
--- /dev/null
+++ b/tests/decomposition/conftest.py
@@ -0,0 +1,79 @@
+# Copyright 2025 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.
+
+"""A fixed set of decomposition rules for testing purposes."""
+
+# pylint: disable=too-few-public-methods,protected-access
+
+from collections import defaultdict
+
+import pennylane as qml
+from pennylane.decomposition import Resources
+from pennylane.decomposition.decomposition_rule import _auto_wrap
+
+decompositions = defaultdict(list)
+
+
+def to_resources(gate_count: dict) -> Resources:
+    """Wrap a dictionary of gate counts in a Resources object."""
+    return Resources({_auto_wrap(op): count for op, count in gate_count.items() if count > 0})
+
+
+@qml.register_resources({qml.Hadamard: 2, qml.CNOT: 1})
+def _cz_to_cnot(*_, **__):
+    raise NotImplementedError
+
+
+decompositions[qml.CZ] = [_cz_to_cnot]
+
+
+@qml.register_resources({qml.Hadamard: 2, qml.CZ: 1})
+def _cnot_to_cz(*_, **__):
+    raise NotImplementedError
+
+
+decompositions[qml.CNOT] = [_cnot_to_cz]
+
+
+def _multi_rz_decomposition_resources(num_wires):
+    return {qml.RZ: 1, qml.CNOT: 2 * (num_wires - 1)}
+
+
+@qml.register_resources(_multi_rz_decomposition_resources)
+def _multi_rz_decomposition(*_, **__):
+    raise NotImplementedError
+
+
+decompositions[qml.MultiRZ] = [_multi_rz_decomposition]
+
+
+@qml.register_resources({qml.RZ: 2, qml.RX: 1, qml.GlobalPhase: 1})
+def _hadamard_to_rz_rx(*_, **__):
+    raise NotImplementedError
+
+
+@qml.register_resources({qml.RZ: 1, qml.RY: 1, qml.GlobalPhase: 1})
+def _hadamard_to_rz_ry(*_, **__):
+    raise NotImplementedError
+
+
+decompositions[qml.Hadamard] = [_hadamard_to_rz_rx, _hadamard_to_rz_ry]
+
+
+@qml.register_resources({qml.RX: 1, qml.RZ: 2})
+def _ry_to_rx_rz(*_, **__):
+    raise NotImplementedError
+
+
+decompositions[qml.RY] = [_ry_to_rx_rz]
diff --git a/tests/decomposition/test_decomposition_graph.py b/tests/decomposition/test_decomposition_graph.py
new file mode 100644
index 000000000..465d3a3b2
--- /dev/null
+++ b/tests/decomposition/test_decomposition_graph.py
@@ -0,0 +1,216 @@
+# Copyright 2025 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.
+
+"""Unit tests for the decomposition graph."""
+
+# pylint: disable=protected-access,no-name-in-module
+
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+from conftest import decompositions, to_resources
+
+import pennylane as qml
+from pennylane.decomposition import DecompositionGraph
+from pennylane.decomposition.decomposition_graph import DecompositionError
+
+
+@patch(
+    "pennylane.decomposition.decomposition_graph.list_decomps",
+    side_effect=lambda x: decompositions[x],
+)
+class TestDecompositionGraph:
+
+    @pytest.mark.unit
+    def test_get_decomp_rule(self, _):
+        """Tests the internal method that gets the decomposition rules for an operator."""
+
+        @qml.register_resources({qml.PhaseShift: 2, qml.RX: 1})
+        def custom_hadamard(wires):
+            qml.PhaseShift(np.pi / 2, wires=wires)
+            qml.RX(np.pi / 2, wires=wires)
+            qml.PhaseShift(np.pi / 2, wires=wires)
+
+        @qml.register_resources({qml.PhaseShift: 1, qml.RY: 1})
+        def custom_hadamard_2(wires):
+            qml.PhaseShift(np.pi / 2, wires=wires)
+            qml.RY(np.pi / 2, wires=wires)
+
+        graph = DecompositionGraph(operations=[qml.Hadamard(0)], target_gate_set={"RX", "RY", "RZ"})
+        assert graph._get_decompositions(qml.Hadamard) == decompositions[qml.Hadamard]
+
+        graph = DecompositionGraph(
+            operations=[qml.Hadamard(0)],
+            target_gate_set={"RX", "RY", "RZ"},
+            fixed_decomps={qml.Hadamard: custom_hadamard},
+        )
+        assert graph._get_decompositions(qml.Hadamard) == [custom_hadamard]
+
+        alt_dec = [custom_hadamard, custom_hadamard_2]
+        graph = DecompositionGraph(
+            operations=[qml.Hadamard(0)],
+            target_gate_set={"RX", "RY", "RZ"},
+            alt_decomps={qml.Hadamard: alt_dec},
+        )
+        exp_dec = alt_dec + decompositions[qml.Hadamard]
+        assert graph._get_decompositions(qml.Hadamard) == exp_dec
+
+        graph = DecompositionGraph(
+            operations=[qml.Hadamard(0)],
+            target_gate_set={"RX", "RY", "RZ"},
+            alt_decomps={qml.Hadamard: alt_dec},
+            fixed_decomps={qml.Hadamard: custom_hadamard},
+        )
+        assert graph._get_decompositions(qml.Hadamard) == [custom_hadamard]
+
+    @pytest.mark.unit
+    def test_graph_construction(self, _):
+        """Tests constructing a graph from a single Hadamard."""
+
+        op = qml.Hadamard(wires=[0])
+        graph = DecompositionGraph(operations=[op], target_gate_set={"RX", "RZ", "GlobalPhase"})
+        # 5 ops and 3 decompositions (2 for Hadamard and 1 for RY)
+        assert len(graph._graph.nodes()) == 8
+        # 8 edges from ops to decompositions and 3 from decompositions to ops
+        assert len(graph._graph.edges()) == 11
+
+        # Check that graph construction stops at gates in the target gate set.
+        graph2 = DecompositionGraph(operations=[op], target_gate_set={"RY", "RZ", "GlobalPhase"})
+        # 5 ops and 2 decompositions (RY is in the target gate set now)
+        assert len(graph2._graph.nodes()) == 7
+        # 6 edges from ops to decompositions and 2 from decompositions to ops
+        assert len(graph2._graph.edges()) == 8
+
+    @pytest.mark.unit
+    def test_graph_solve(self, _):
+        """Tests solving a simple graph for the optimal decompositions."""
+
+        op = qml.Hadamard(wires=[0])
+        graph = DecompositionGraph(
+            operations=[op],
+            target_gate_set={"RX", "RY", "RZ", "GlobalPhase"},
+        )
+        graph.solve()
+
+        # verify that the better decomposition rule is chosen when both are valid.
+        expected_resource = to_resources({qml.RZ: 1, qml.RY: 1, qml.GlobalPhase: 1})
+        assert graph.resource_estimate(op) == expected_resource
+        assert graph.decomposition(op).compute_resources() == expected_resource
+
+    @pytest.mark.unit
+    def test_decomposition_not_found(self, _):
+        """Tests that the correct error is raised if a decomposition isn't found."""
+
+        op = qml.Hadamard(wires=[0])
+        graph = DecompositionGraph(operations=[op], target_gate_set={"RX", "RY", "GlobalPhase"})
+        with pytest.raises(DecompositionError, match="Decomposition not found for {'Hadamard'}"):
+            graph.solve()
+
+    @pytest.mark.unit
+    def test_lazy_solve(self, _):
+        """Tests the lazy keyword argument."""
+
+        class CustomOp(qml.operation.Operation):  # pylint: disable=too-few-public-methods
+            """A custom operation."""
+
+            resource_keys = set()
+
+            @property
+            def resource_params(self):
+                return {}
+
+        class AnotherOp(qml.operation.Operation):  # pylint: disable=too-few-public-methods
+            """Another custom operation."""
+
+            resource_keys = set()
+
+            @property
+            def resource_params(self):
+                return {}
+
+        @qml.register_resources({qml.RZ: 1, qml.CNOT: 1})
+        def _custom_decomp(*_, **__):
+            raise NotImplementedError
+
+        @qml.register_resources({AnotherOp: 1, qml.CNOT: 1})
+        def _custom_decomp_2(*_, **__):
+            raise NotImplementedError
+
+        @qml.register_resources({qml.RZ: 2, qml.CNOT: 2})
+        def _another_decomp(*_, **__):
+            raise NotImplementedError
+
+        graph = DecompositionGraph(
+            operations=[CustomOp(wires=[0])],
+            target_gate_set={"CNOT", "RZ"},
+            alt_decomps={
+                CustomOp: [_custom_decomp, _custom_decomp_2],
+                AnotherOp: [_another_decomp],
+            },
+        )
+        graph.solve(lazy=True)
+        assert not graph.is_solved_for(AnotherOp(wires=[0, 1]))
+
+        graph.solve(lazy=False)
+        assert graph.is_solved_for(AnotherOp(wires=[0, 1]))
+
+    @pytest.mark.unit
+    def test_decomposition_with_resource_params(self, _):
+        """Tests operators with non-empty resource params."""
+
+        class CustomOp(qml.operation.Operation):  # pylint: disable=too-few-public-methods
+            """A custom operation."""
+
+            resource_keys = {"num_wires"}
+
+            @property
+            def resource_params(self):
+                return {"num_wires": len(self.wires)}
+
+        def _custom_resource(num_wires):
+            return {
+                qml.resource_rep(qml.MultiRZ, num_wires=num_wires): 1,
+                qml.resource_rep(qml.MultiRZ, num_wires=num_wires - 1): 2,
+            }
+
+        @qml.register_resources(_custom_resource)
+        def _custom_decomp(*_, **__):
+            raise NotImplementedError
+
+        op = CustomOp(wires=[0, 1, 2, 3])
+        graph = DecompositionGraph(
+            operations=[op],
+            target_gate_set={"RX", "RZ", "CZ", "GlobalPhase"},
+            alt_decomps={CustomOp: [_custom_decomp]},
+        )
+        # 10 ops (CustomOp, MultiRZ(4), MultiRZ(3), CNOT, CZ, RX, RY, RZ, Hadamard, GlobalPhase)
+        # 7 decompositions (1 for CustomOp, 1 for each of the two MultiRZs, 1 for CNOT, 2 for Hadamard, and 1 for RY)
+        assert len(graph._graph.nodes()) == 17
+        # 16 edges from ops to decompositions and 7 from decompositions to ops
+        assert len(graph._graph.edges()) == 23
+
+        graph.solve()
+        assert graph.resource_estimate(op) == to_resources(
+            {qml.CZ: 14, qml.RZ: 59, qml.RX: 28, qml.GlobalPhase: 28}
+        )
+        assert graph.decomposition(op).compute_resources(**op.resource_params) == to_resources(
+            {
+                qml.resource_rep(qml.MultiRZ, num_wires=4): 1,
+                qml.resource_rep(qml.MultiRZ, num_wires=3): 2,
+            },
+        )
+        assert graph.decomposition(qml.Hadamard(wires=[0])).compute_resources() == to_resources(
+            {qml.RZ: 2, qml.RX: 1, qml.GlobalPhase: 1}
+        )
diff --git a/tests/decomposition/test_resources.py b/tests/decomposition/test_resources.py
index 07141e3bb..8f36a3f5d 100644
--- a/tests/decomposition/test_resources.py
+++ b/tests/decomposition/test_resources.py
@@ -188,6 +188,17 @@ class TestResourceRep:
         with pytest.raises(TypeError, match="op_type must be a type of Operator"):
             resource_rep(int)
 
+        class CustomOp(qml.operation.Operator):  # pylint: disable=too-few-public-methods
+
+            resource_keys = {}
+
+            @property
+            def resource_params(self) -> dict:
+                return {}
+
+        with pytest.raises(TypeError, match="CustomOp.resource_keys must be a set"):
+            resource_rep(CustomOp)
+
     def test_params_mismatch(self):
         """Tests that an error is raised when parameters are missing."""