Repository: PennyLaneAI/pennylane
Pull request #7389 (2025-05-06 21:41:28)

Support for weighted gates in graph decomposition algorithm

------------------------------------------------------------------------------------------------------------

**Context:** In short, a decomposition system based on a graph traversal using Dijkstra's algorithm is applied to find the most optimal gate decompositions for non-native gates. The nodes in the graph represent either decompositions, operations in the target gateset, or higher order operations. The edges show either membership of a native op in a decomposition, or equivalence of a decomposition to a higher order operation. The cost of an edge between a native operation and a decomposition is the number of additional gates found in the decomposition. The cost of an edge between a decomposition and an equivalent higher order op is zero. For more context, see [[sc-84664](https://app.shortcut.com/xanaduai/story/84664)].

**Description of the Change:** Work in progress on implementing support for relatively weighted gates within a target gateset. The weights are provided by the user, and influence what gate decompositions are chosen by the graph decomposition algorithm. The implementation is such that the weights of the gates are represented by weights on the edges leading to and from native gates.

**Benefits:** Allows a user to more closely model the behaviour of their backends, and capture the relative cost of executing certain gates such as non-Clifford gates which are native, but more expensive.

**Possible Drawbacks:** The representation of weights in this way is possibly not intuitive.

**Related GitHub Issues:**


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7389?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.65%. Comparing base [(`353fb36`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/353fb36a73bb6e9c0df10dec3a5afb20e0d85424?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`976e43f`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/976e43f440de098a3c308546ca85ce610e05a8ff?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 24 commits behind head on symbolic-rules-01.

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


```diff
@@                Coverage Diff                 @@
##           symbolic-rules-01    #7389   +/-   ##
==================================================
  Coverage              99.65%   99.65%           
==================================================
  Files                    529      529           
  Lines                  50603    50604    +1     
==================================================
+ Hits                   50427    50428    +1     
  Misses                   176      176           
```

</details>

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

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 9bde96dfb..bdd7f2939 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -73,6 +73,62 @@
 
 <h4>Resource-efficient Decompositions 🔎</h4>
 
+* The :func:`~.transforms.decompose` transform now supports weighting gates in the target `gate_set`, allowing for 
+  preferential treatment of certain gates in a target `gate_set` over others.
+  [(#7389)](https://github.com/PennyLaneAI/pennylane/pull/7389)
+
+  Gates specified in `gate_set` can be given a numerical weight associated with their effective cost to have in a circuit:
+  
+  * Gate weights that are greater than 1 indicate a *greater cost* (less preferred).
+  * Gate weights that are less than 1 indicate a *lower cost* (more preferred).
+
+  Consider the following toy example.
+
+  ```python
+  qml.decomposition.enable_graph()
+  
+  @partial(
+    qml.transforms.decompose, gate_set={qml.Toffoli: 1.23, qml.RX: 4.56, qml.CZ: 0.01, qml.H: 420, qml.CRZ: 100}
+  )
+  @qml.qnode(qml.device("default.qubit"))
+  def circuit():
+      qml.CRX(0.1, wires=[0, 1])
+      qml.Toffoli(wires=[0, 1, 2])
+      return qml.expval(qml.Z(0))
+  ```
+
+  ```pycon
+  >>> print(qml.draw(circuit)())
+
+  0: ───────────╭●────────────╭●─╭●─┤  <Z>
+  1: ──RX(0.05)─╰Z──RX(-0.05)─╰Z─├●─┤     
+  2: ────────────────────────────╰X─┤     
+  ```
+
+  ```python
+  qml.decomposition.enable_graph()
+
+  @partial(
+      qml.transforms.decompose, gate_set={qml.Toffoli: 1.23, qml.RX: 4.56, qml.CZ: 0.01, qml.H: 0.1, qml.CRZ: 0.1}
+  )
+  @qml.qnode(qml.device("default.qubit"))
+  def circuit():
+      qml.CRX(0.1, wires=[0, 1])
+      qml.Toffoli(wires=[0, 1, 2])
+      return qml.expval(qml.Z(0))
+  ```
+
+  ```pycon
+  >>> print(qml.draw(circuit)())
+
+  0: ────╭●───────────╭●─┤  <Z>
+  1: ──H─╰RZ(0.10)──H─├●─┤     
+  2: ─────────────────╰X─┤  
+  ```
+
+  Here, when the Hadamard and ``CRZ`` have relatively high weights, a decomposition involving them is considered *less* 
+  efficient. When they have relatively low weights, a decomposition involving them is considered *more* efficient.
+
 * The decomposition of `qml.PCPhase` is now significantly more efficient for more than 2 qubits.
   [(#7166)](https://github.com/PennyLaneAI/pennylane/pull/7166)
 
@@ -156,6 +212,7 @@
     :func:`~.transforms.decompose` transform, allowing custom decomposition rules to be defined and
     registered for symbolic operators.
     [(#7347)](https://github.com/PennyLaneAI/pennylane/pull/7347)
+
     [(#7352)](https://github.com/PennyLaneAI/pennylane/pull/7352)
 
     ```python
@@ -504,6 +561,7 @@ This release contains contributions from (in alphabetical order):
 Guillermo Alonso-Linaje,
 Astral Cai,
 Yushao Chen,
+Marcus Edwards,
 Lillian Frederiksen,
 Pietropaolo Frisoni,
 Simone Gasperini,
diff --git a/pennylane/decomposition/decomposition_graph.py b/pennylane/decomposition/decomposition_graph.py
index 99ee9448a..d3066534d 100644
--- a/pennylane/decomposition/decomposition_graph.py
+++ b/pennylane/decomposition/decomposition_graph.py
@@ -69,13 +69,17 @@ class DecompositionGraph:  # pylint: disable=too-many-instance-attributes
 
     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 connect decomposition rules to operators have a weight of 0 because an operator can
-    be replaced with its decomposition at no additional cost.
+    decompose. The edge weights represent the difference in the total weights of the gates between
+    the two states. Edges that connect 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 ``CNOT`` to the following
+    estimate of the operator. Edges that connect an operator node to a decomposition node have a weight
+    calculated by the difference of the sum of the gate counts multiplied by their respective gate
+    weights in the decomposition, minus the weight of the operator of the operator node.
+
+    For example, if the graph was initialized with ``{qml.CNOT: 10.0, qml.H: 1.0}`` as the gate set, the edge that connects a ``CNOT`` to the following
     decomposition rule:
 
     .. code-block:: python
@@ -88,16 +92,19 @@ class DecompositionGraph:  # pylint: disable=too-many-instance-attributes
             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 one ``RX`` gate, the weight of this edge
-    becomes 2 * 3 = 6. This way, the total distance from the basis gate set 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.
+    will have a weight of (10.0 + 2 * 1.0) - 10.0 = 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 one ``RX`` gate, the weight of this edge
+    becomes 2 * 3 = 6, if ``RZ`` and ``RX`` have weights of 1.0 (the default). This way, the total distance
+    from the basis gate set to a high-level gate is by default 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. By specifying weights in the target gate set, the total distance calculation involves
+    a sum of weighted gate counts, which can represent the relative cost of executing a particular element
+    of the target gate set on the target hardware i.e. a ``T`` gate.
 
     Args:
         operations (list[Operator or CompressedResourceOp]): The list of operations to decompose.
-        gate_set (set[str]): The names of the gates in the target gate set.
+        gate_set (set[str | type] | dict[type | str, float]): A set of gates in the target gate set or a dictionary mapping gates in the target gate set to their respective weights. All weights must be positive.
         fixed_decomps (dict): A dictionary mapping operator names to fixed decompositions.
         alt_decomps (dict): A dictionary mapping operator names to alternative decompositions.
 
@@ -129,12 +136,16 @@ class DecompositionGraph:  # pylint: disable=too-many-instance-attributes
     def __init__(
         self,
         operations: list[Operator | CompressedResourceOp],
-        gate_set: set[str],
+        gate_set: set[type | str] | dict[type | str, float],
         fixed_decomps: dict = None,
         alt_decomps: dict = None,
     ):
-        # The names of the gates in the target gate set.
-        self._gate_set: set[str] = {translate_op_alias(op) for op in gate_set}
+        if isinstance(gate_set, set):
+            # The names of the gates in the target gate set.
+            self._weights = {_to_name(gate): 1.0 for gate in gate_set}
+        else:
+            # the gate_set is a dict
+            self._weights = {_to_name(gate): weight for gate, weight in gate_set.items()}
 
         # Tracks the node indices of various operators.
         self._original_ops_indices: set[int] = set()
@@ -217,8 +228,8 @@ class DecompositionGraph:  # pylint: disable=too-many-instance-attributes
         op_node_idx = self._graph.add_node(op_node)
         self._all_op_indices[op_node] = op_node_idx
 
-        if op_node.name in self._gate_set:
-            self._graph.add_edge(self._start, op_node_idx, 1)
+        if op_node.name in self._weights:
+            self._graph.add_edge(self._start, op_node_idx, self._weights[op_node.name])
             return op_node_idx
 
         for decomposition in self._get_decompositions(op_node):
@@ -314,7 +325,12 @@ class DecompositionGraph:  # pylint: disable=too-many-instance-attributes
                 entire graph will be explored.
 
         """
-        self._visitor = _DecompositionSearchVisitor(self._graph, self._original_ops_indices, lazy)
+        self._visitor = _DecompositionSearchVisitor(
+            self._graph,
+            self._weights,
+            self._original_ops_indices,
+            lazy,
+        )
         rx.dijkstra_search(
             self._graph,
             source=[self._start],
@@ -325,7 +341,7 @@ class DecompositionGraph:  # pylint: disable=too-many-instance-attributes
             unsolved_ops = [self._graph[op_idx] for op_idx in self._visitor.unsolved_op_indices]
             op_names = set(op.name for op in unsolved_ops)
             raise DecompositionError(
-                f"Decomposition not found for {op_names} to the gate set {self._gate_set}"
+                f"Decomposition not found for {op_names} to the gate set {set(self._weights)}"
             )
 
     def is_solved_for(self, op):
@@ -426,7 +442,13 @@ class DecompositionGraph:  # pylint: disable=too-many-instance-attributes
 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):
+    def __init__(
+        self,
+        graph: rx.PyDiGraph,
+        gate_set: dict,
+        original_op_indices: set[int],
+        lazy: bool = True,
+    ):
         self._graph = graph
         self._lazy = lazy
         # maps node indices to the optimal resource estimates
@@ -435,13 +457,15 @@ class _DecompositionSearchVisitor(DijkstraVisitor):
         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
+        self._gate_weights = gate_set
 
     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
+        return self.distances[d_node_idx].weighted_cost - self.distances[op_node_idx].weighted_cost
 
     def discover_vertex(self, v, _):
         """Triggered when a vertex is about to be explored during the Dijkstra search."""
@@ -476,7 +500,9 @@ class _DecompositionSearchVisitor(DijkstraVisitor):
         src_idx, target_idx, _ = edge
         target_node = self._graph[target_idx]
         if self._graph[src_idx] is None and not isinstance(target_node, _DecompositionNode):
-            self.distances[target_idx] = Resources({target_node: 1})
+            self.distances[target_idx] = Resources(
+                {target_node: 1}, self._gate_weights[_to_name(target_node)]
+            )
         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 c1771562c..65c5e80ac 100644
--- a/pennylane/decomposition/resources.py
+++ b/pennylane/decomposition/resources.py
@@ -25,20 +25,24 @@ import pennylane as qml
 from pennylane.operation import Operator
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=False)
 class Resources:
     r"""Stores resource estimates.
 
     Args:
         gate_counts (dict): dictionary mapping operator types to their number of occurrences.
-
+        weighted_cost (float): the cumulative weight of the gates.
     """
 
     gate_counts: dict[CompressedResourceOp, int] = field(default_factory=dict)
+    weighted_cost: Optional[float] = field(default=None)
 
     def __post_init__(self):
         """Verify that all gate counts are non-zero."""
         assert all(v > 0 for v in self.gate_counts.values())
+        if self.weighted_cost is None:
+            self.weighted_cost = sum(count for gate, count in self.gate_counts.items())
+        assert self.weighted_cost >= 0.0
 
     @cached_property
     def num_gates(self) -> int:
@@ -48,15 +52,18 @@ class Resources:
     def __add__(self, other: Resources):
         return Resources(
             _combine_dict(self.gate_counts, other.gate_counts),
+            weighted_cost=self.weighted_cost + other.weighted_cost,
         )
 
     def __mul__(self, scalar: int):
-        return Resources(_scale_dict(self.gate_counts, scalar))
+        return Resources(
+            _scale_dict(self.gate_counts, scalar), weighted_cost=self.weighted_cost * scalar
+        )
 
     __rmul__ = __mul__
 
     def __repr__(self):
-        return f"<num_gates={self.num_gates}, gate_counts={self.gate_counts}>"
+        return f"<num_gates={self.num_gates}, gate_counts={self.gate_counts}, weighted_cost={self.weighted_cost}>"
 
 
 def _combine_dict(dict1: dict, dict2: dict):
diff --git a/pennylane/transforms/decompose.py b/pennylane/transforms/decompose.py
index 823e3fcf9..698d7e9a8 100644
--- a/pennylane/transforms/decompose.py
+++ b/pennylane/transforms/decompose.py
@@ -32,6 +32,65 @@ from pennylane.operation import Operator
 from pennylane.transforms.core import transform
 
 
+def _resolve_gate_set(gate_set: set[type | str] | dict[type | str, float] = None):
+    """
+    From the provided gate_set dict, set (or None) resolves a standardized form for the gate_set,
+    as well as a tuple target_gate_types and set target_gate_names, and a function that
+    returns whether the gate set contains a given operator.
+
+    Args:
+        gate_set (set[type | str] | dict[type | str, float]): The gate_set to be resolved.
+
+    Returns:
+        gate_set (set[type | str] | dict[type | str, float]): The standardized gate_set.
+        target_gate_types tuple[type]: The types of the gates in the gate_set.
+        target_gate_names set[str]: The names of the gates in the gate_set.
+        gate_set_contains (Callable): A function that checks for the membership of a gate in the gate_set.
+    """
+    target_gate_types = tuple()
+    target_gate_names = set()
+
+    if gate_set is None:
+        gate_set: set[Type[Operator] | str] = set(qml.ops.__all__)
+
+    if isinstance(gate_set, (str, type)):
+        gate_set = {gate_set}
+
+    if isinstance(gate_set, dict):
+        for v in gate_set.values():
+            if v < 0.0:
+                raise ValueError(
+                    "Negative gate weights provided to gate_set in decompose" "are not supported."
+                )
+
+        if not qml.decomposition.enabled_graph():
+            raise TypeError(
+                "Specifying the gate_set with a dictionary of operator types and their weights is only supported "
+                "with the new experimental graph-based decomposition system. Enable the new system "
+                "using qml.decomposition.enable_graph()"
+            )
+
+    if isinstance(gate_set, Iterable):
+        target_gate_types = tuple(gate for gate in gate_set if isinstance(gate, type))
+        target_gate_names = set(
+            translate_op_alias(gate) for gate in gate_set if isinstance(gate, str)
+        )
+
+        def gate_set_contains(op):
+            return (op.name in target_gate_names) or isinstance(op, target_gate_types)
+
+    else:  # if isinstance(gate_set, Callable):
+        gate_set_contains = gate_set
+
+        if qml.decomposition.enabled_graph():
+            raise TypeError(
+                "Specifying gate_set as a function is not supported with the new "
+                "graph-based decomposition system enabled."
+            )
+
+    return gate_set, target_gate_types, target_gate_names, gate_set_contains
+
+
 def null_postprocessing(results):
     """A postprocessing function returned by a transform that only converts the batch of results
     into a result for a single ``QuantumTape``.
@@ -96,39 +155,27 @@ def _get_plxpr_decompose():  # pylint: disable=missing-docstring, too-many-state
             # The name is different from the _env in the parent class (a dictionary) to avoid confusion.
             self._env_map = ChainMap()
 
-            if gate_set is None:
-                gate_set: set[Type[Operator] | str] = set(qml.ops.__all__)
-
-            if isinstance(gate_set, (str, type)):
-                gate_set: set[Type[Operator] | str] = {gate_set}
+            gate_set, target_gate_types, target_gate_names, _in_gate_set = _resolve_gate_set(
+                gate_set
+            )
 
             if isinstance(gate_set, Iterable):
 
-                target_gate_types = tuple(gate for gate in gate_set if isinstance(gate, type))
-                target_gate_names = set(
-                    translate_op_alias(gate) for gate in gate_set if isinstance(gate, str)
-                )
-
-                def _in_gate_set(op: Operator) -> bool:
-                    return (op.name in target_gate_names) or isinstance(op, target_gate_types)
-
                 self.gate_set_contains = _in_gate_set
 
                 if qml.decomposition.enabled_graph():
 
                     type_to_names = {op_type.__name__ for op_type in target_gate_types}
-                    self._target_gate_names = target_gate_names | type_to_names
+                    self._gate_set = (
+                        gate_set
+                        if isinstance(gate_set, dict)
+                        else target_gate_names | type_to_names
+                    )
 
             else:  # isinstance(gate_set, Callable)
 
                 self.gate_set_contains: Callable[[Operator], bool] = gate_set
 
-                if qml.decomposition.enabled_graph():
-                    raise TypeError(
-                        "Specifying gate_set as a function is not supported with the new "
-                        "graph-based decomposition system enabled."
-                    )
-
         def setup(self) -> None:
             """Setup the environment for the interpreter by pushing a new environment frame."""
 
@@ -271,7 +318,7 @@ def _get_plxpr_decompose():  # pylint: disable=missing-docstring, too-many-state
                 if operations:
                     self._decomp_graph = _construct_and_solve_decomp_graph(
                         operations,
-                        self._target_gate_names,
+                        self._gate_set,
                         self._fixed_decomps,
                         self._alt_decomps,
                     )
@@ -403,11 +450,12 @@ def decompose(
 
     Args:
         tape (QuantumScript or QNode or Callable): a quantum circuit.
-        gate_set (Iterable[str or type] or Callable, optional): The target gate set specified as
-            either (1) a sequence of operator types and/or names or (2) a function that returns
-            ``True`` if the operator belongs to the target gate set. Defaults to ``None``, in which
-            case the gate set is considered to be all available :doc:`quantum operators </introduction/operations>`.
-        max_expansion (int, optional): The maximum depth of the decomposition. Defaults to None.
+        gate_set (Iterable[str or type], Dict[type or str, float], or Callable, optional): The target gate
+            set specified as either (1) a sequence of operator types and/or names, (2) a function that returns
+            ``True`` if the operator belongs to the target gate set or (3) a dictionary of operator
+            types and/or names with weights. Defaults to ``None``, in which case the gate set is
+            considered to be all available :doc:`quantum operators </introduction/operations>`.
+        max_expansion (int, optional): The maximum depth of the decomposition. Defaults to ``None``.
             If ``None``, the circuit will be decomposed until the target gate set is reached.
         fixed_decomps (Dict[Type[Operator], DecompositionRule]): a dictionary mapping operator types
             to custom decomposition rules. A decomposition rule is a quantum function decorated with
@@ -580,6 +628,52 @@ def decompose(
             >>> new_tape.operations
             [RX(0.25, wires=[1]), CZ(wires=[0, 1]), RX(-0.25, wires=[1]), CZ(wires=[0, 1])]
 
+        **Weighted Gate Sets**
+
+        With the graph based decomposition enabled, gate weights can be provided in the ``gate_set`` parameter. For example:
+
+        .. code-block:: python
+
+            @partial(
+                qml.transforms.decompose,
+                gate_set={qml.Toffoli: 1.23, qml.RX: 4.56, qml.CZ: 0.01, qml.H: 420, qml.CRZ: 100}
+            )
+            @qml.qnode(qml.device("default.qubit"))
+            def circuit():
+                qml.CRX(0.1, wires=[0, 1])
+                qml.Toffoli(wires=[0, 1, 2])
+                return qml.expval(qml.Z(0))
+
+        .. code-block:: pycon
+
+            >>> print(qml.draw(circuit)())
+            0: ───────────╭●────────────╭●─╭●─┤  <Z>
+            1: ──RX(0.05)─╰Z──RX(-0.05)─╰Z─├●─┤
+            2: ────────────────────────────╰X─┤
+
+        .. code-block:: python
+
+            @partial(
+                qml.transforms.decompose,
+                gate_set={qml.Toffoli: 1.23, qml.RX: 4.56, qml.CZ: 0.01, qml.H: 0.1, qml.CRZ: 0.1}
+            )
+            @qml.qnode(qml.device("default.qubit"))
+            def circuit():
+                qml.CRX(0.1, wires=[0, 1])
+                qml.Toffoli(wires=[0, 1, 2])
+                return qml.expval(qml.Z(0))
+
+        .. code-block:: pycon
+
+            >>> print(qml.draw(circuit)())
+            0: ────╭●───────────╭●─┤  <Z>
+            1: ──H─╰RZ(0.10)──H─├●─┤
+            2: ─────────────────╰X─┤
+
+
+        Here, when the Hadamard and ``CRZ`` have relatively high weights, a decomposition involving them is considered
+        *less* efficient. When they have relatively low weights, a decomposition involving them is considered *more*
+        efficient.
 
         **Customizing Decompositions**
 
@@ -646,36 +740,7 @@ def decompose(
 
     _decomp_graph_kwargs_checks(fixed_decomps, alt_decomps)
 
-    if isinstance(gate_set, (str, type)):
-        gate_set = {gate_set}
-
-    if isinstance(gate_set, Iterable):
-        target_gate_types = tuple(gate for gate in gate_set if isinstance(gate, type))
-        target_gate_names = set(
-            translate_op_alias(gate) for gate in gate_set if isinstance(gate, str)
-        )
-
-        def gate_set_contains(op):
-            return (op.name in target_gate_names) or isinstance(op, target_gate_types)
-
-    # If the gate_set is None, we don't need to iterate over
-    # all the ops to construct `target_gate_types` or `target_gate_names`
-    elif gate_set is None:
-
-        target_gate_types = tuple()
-        target_gate_names = set(qml.ops.__all__)
-
-        def gate_set_contains(op):
-            return op.name in target_gate_names
-
-    else:
-        gate_set_contains = gate_set
-
-        if qml.decomposition.enabled_graph():
-            raise TypeError(
-                "Specifying gate_set as a function is not supported with the new "
-                "graph-based decomposition system enabled."
-            )
+    gate_set, target_gate_types, target_gate_names, gate_set_contains = _resolve_gate_set(gate_set)
 
     def stopping_condition(op):
 
@@ -710,7 +775,7 @@ def decompose(
 
         decomp_graph = _construct_and_solve_decomp_graph(
             tape.operations,
-            target_gate_names,
+            target_gate_names if not isinstance(gate_set, dict) else gate_set,
             fixed_decomps=fixed_decomps,
             alt_decomps=alt_decomps,
         )
@@ -783,13 +848,13 @@ def _decomp_graph_kwargs_checks(fixed_decomps, alt_decomps):
         )
 
 
-def _construct_and_solve_decomp_graph(operations, target_gate_names, fixed_decomps, alt_decomps):
+def _construct_and_solve_decomp_graph(operations, target_gates, fixed_decomps, alt_decomps):
     """Create and solve a DecompositionGraph instance to optimize the decomposition."""
 
     # Create the decomposition graph
     decomp_graph = DecompositionGraph(
         operations,
-        target_gate_names,
+        target_gates,
         fixed_decomps=fixed_decomps,
         alt_decomps=alt_decomps,
     )


Test patch:
diff --git a/tests/decomposition/conftest.py b/tests/decomposition/conftest.py
index 4f8fd924b..a4aced002 100644
--- a/tests/decomposition/conftest.py
+++ b/tests/decomposition/conftest.py
@@ -31,9 +31,16 @@ from pennylane.decomposition.symbolic_decomposition import (
 decompositions = defaultdict(list)
 
 
-def to_resources(gate_count: dict) -> Resources:
+def to_resources(gate_count: dict, weighted_cost: float = None) -> 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})
+    return Resources(
+        {_auto_wrap(op): count for op, count in gate_count.items() if count >= 0},
+        (
+            sum(count for gate, count in gate_count.items())
+            if weighted_cost is None
+            else weighted_cost
+        ),
+    )
 
 
 @qml.register_resources({qml.Hadamard: 2, qml.CNOT: 1})
@@ -90,7 +97,12 @@ def _crx_to_rx_cz(*_, **__):
     raise NotImplementedError
 
 
-decompositions["CRX"] = [_crx_to_rx_cz]
+@qml.register_resources({qml.RX: 2, qml.CNOT: 2, qml.RY: 4, qml.GlobalPhase: 4, qml.RZ: 4})
+def _crx_to_rx_ry_cnot_ry_cnot_ry_cnot_rz(*_, **__):
+    raise NotImplementedError
+
+
+decompositions["CRX"] = [_crx_to_rx_cz, _crx_to_rx_ry_cnot_ry_cnot_ry_cnot_rz]
 
 
 @qml.register_resources({qml.RZ: 3, qml.CNOT: 2, qml.GlobalPhase: 1})
diff --git a/tests/decomposition/test_controlled_decomposition.py b/tests/decomposition/test_controlled_decomposition.py
index 7259d8280..547dfc53c 100644
--- a/tests/decomposition/test_controlled_decomposition.py
+++ b/tests/decomposition/test_controlled_decomposition.py
@@ -112,7 +112,7 @@ class TestControlledDecompositionRules:
                         "num_work_wires": 0,
                     },
                 ): 1,
-            }
+            },
         )
 
 
@@ -180,7 +180,7 @@ class TestControlledX:
                         "num_work_wires": 1,
                     },
                 ): 1
-            }
+            },
         )
 
         op = qml.ops.Controlled(
@@ -202,7 +202,8 @@ class TestControlledX:
                         "num_work_wires": 1,
                     },
                 ): 1,
-            }
+            },
+            1.0,
         )
 
 
@@ -399,7 +400,7 @@ class TestControlledBaseDecomposition:
                 ): 1,
                 qml.resource_rep(qml.CZ): 1,
                 qml.resource_rep(qml.CCZ): 1,
-            }
+            },
         )
 
     def test_double_control_wire(self):
@@ -511,7 +512,7 @@ class TestControlledBaseDecomposition:
                 qml.decomposition.controlled_resource_rep(
                     qml.Z, {}, num_control_wires=3, num_work_wires=1
                 ): 1,
-            }
+            },
         )
 
     def test_multi_control_wires(self):
@@ -632,5 +633,5 @@ class TestControlledBaseDecomposition:
                 qml.decomposition.controlled_resource_rep(
                     qml.Z, {}, num_control_wires=4, num_work_wires=1
                 ): 1,
-            }
+            },
         )
diff --git a/tests/decomposition/test_decomposition_graph.py b/tests/decomposition/test_decomposition_graph.py
index 96e434567..666f75fe7 100644
--- a/tests/decomposition/test_decomposition_graph.py
+++ b/tests/decomposition/test_decomposition_graph.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 
 """Unit tests for the decomposition graph."""
-
 # pylint: disable=protected-access,no-name-in-module
 
 from unittest.mock import patch
@@ -41,6 +40,51 @@ from pennylane.decomposition.decomposition_graph import _to_name
 )
 class TestDecompositionGraph:
 
+    def test_weighted_graph_solve(self, _):
+        """Tests solving a simple graph for the optimal decompositions with weighted gates."""
+
+        op = qml.CRX(2.5, wires=[0, 1])
+
+        # the RZ CZ RX CZ decomp is chosen when the RZ and CNOT weights are large.
+        gate_weights = {
+            "RX": 1.0,
+            "RY": 3.0,
+            "RZ": 10.0,
+            "GlobalPhase": 1.0,
+            "CNOT": 20.0,
+            "CZ": 1.0,
+        }
+
+        graph = DecompositionGraph(
+            operations=[op],
+            gate_set=gate_weights,
+        )
+        graph.solve()
+
+        expected_resource = to_resources({qml.CZ: 2, qml.RX: 2})
+        assert graph.resource_estimate(op) == expected_resource
+
+        # the RZ CZ RX CZ decomp is avoided when the CZ weight is large.
+        gate_weights = {
+            "RX": 1.0,
+            "RY": 1.0,
+            "RZ": 1.0,
+            "GlobalPhase": 1.0,
+            "CNOT": 1.0,
+            "CZ": 100.0,
+        }
+
+        graph = DecompositionGraph(
+            operations=[op],
+            gate_set=gate_weights,
+        )
+        graph.solve()
+
+        expected_resource = to_resources(
+            {qml.RX: 2, qml.CNOT: 2, qml.RY: 4, qml.GlobalPhase: 4, qml.RZ: 4}
+        )
+        assert graph.resource_estimate(op) == expected_resource
+
     def test_get_decomp_rule(self, _):
         """Tests the internal method that gets the decomposition rules for an operator."""
 
@@ -86,7 +130,10 @@ class TestDecompositionGraph:
         """Tests constructing a graph from a single Hadamard."""
 
         op = qml.Hadamard(wires=[0])
-        graph = DecompositionGraph(operations=[op], gate_set={"RX", "RZ", "GlobalPhase"})
+
+        graph = DecompositionGraph(
+            operations=[op], gate_set={"RX": 1.0, "RZ": 1.0, "GlobalPhase": 1.0}
+        )
         # 5 ops and 3 decompositions (2 for Hadamard and 1 for RY) and 1 dummy starting node
         assert len(graph._graph.nodes()) == 9
         # 8 edges from ops to decompositions, 3 from decompositions to ops, and 3 from the
@@ -191,9 +238,12 @@ class TestDecompositionGraph:
         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
+        assert graph.resource_estimate(op) == to_resources(
+            {qml.RY: 1, qml.GlobalPhase: 1, qml.RZ: 1},
+        )
+        assert graph.decomposition(op).compute_resources() == to_resources(
+            {qml.RY: 1, qml.GlobalPhase: 1, qml.RZ: 1},
+        )
 
     def test_decomposition_not_found(self, _):
         """Tests that the correct error is raised if a decomposition isn't found."""
@@ -294,16 +344,16 @@ class TestDecompositionGraph:
 
         graph.solve()
         assert graph.resource_estimate(op) == to_resources(
-            {qml.CZ: 14, qml.RZ: 59, qml.RX: 28, qml.GlobalPhase: 28}
+            {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}
+            {qml.RZ: 2, qml.RX: 1, qml.GlobalPhase: 1},
         )
 
 
@@ -524,7 +574,7 @@ class TestSymbolicDecompositions:
             qml.adjoint(qml.H(wires=0)),
         ]
         assert graph.resource_estimate(op) == to_resources(
-            {qml.H: 1, qml.CNOT: 2, qml.RX: 1, qml.PhaseShift: 1}
+            {qml.H: 1, qml.CNOT: 2, qml.RX: 1, qml.PhaseShift: 1},
         )
 
     def test_nested_powers(self, _):
diff --git a/tests/decomposition/test_decomposition_rule.py b/tests/decomposition/test_decomposition_rule.py
index 323bc0a3d..e065138db 100644
--- a/tests/decomposition/test_decomposition_rule.py
+++ b/tests/decomposition/test_decomposition_rule.py
@@ -63,7 +63,7 @@ class TestDecompositionRule:
         ]
 
         assert multi_rz_decomposition.compute_resources(num_wires=3) == Resources(
-            gate_counts={CompressedResourceOp(qml.RZ): 1, CompressedResourceOp(qml.CNOT): 4}
+            gate_counts={CompressedResourceOp(qml.RZ): 1, CompressedResourceOp(qml.CNOT): 4},
         )
 
     def test_decomposition_decorator(self):
@@ -97,7 +97,7 @@ class TestDecompositionRule:
         ]
 
         assert multi_rz_decomposition.compute_resources(num_wires=3) == Resources(
-            gate_counts={CompressedResourceOp(qml.RZ): 1, CompressedResourceOp(qml.CNOT): 4}
+            gate_counts={CompressedResourceOp(qml.RZ): 1, CompressedResourceOp(qml.CNOT): 4},
         )
 
     def test_decomposition_condition(self):
diff --git a/tests/decomposition/test_resources.py b/tests/decomposition/test_resources.py
index 2b87302ce..30ac57a17 100644
--- a/tests/decomposition/test_resources.py
+++ b/tests/decomposition/test_resources.py
@@ -37,6 +37,7 @@ class TestResources:
         resources = Resources()
         assert resources.num_gates == 0
         assert resources.gate_counts == {}
+        assert resources.weighted_cost == 0.0
 
     def test_negative_gate_counts(self):
         """Tests that an error is raised if the gate count is negative."""
@@ -48,14 +49,26 @@ class TestResources:
                 }
             )
 
+    def test_negative_weighted_cost(self):
+        """Tests that an error is raised if the cost is negative."""
+        with pytest.raises(AssertionError):
+            Resources(
+                gate_counts={
+                    CompressedResourceOp(qml.RX, {}): 2,
+                },
+                weighted_cost=-2.0,
+            )
+
     def test_add_resources(self):
         """Tests adding two Resources objects."""
 
         resources1 = Resources(
-            gate_counts={CompressedResourceOp(qml.RX, {}): 2, CompressedResourceOp(qml.RZ, {}): 1}
+            gate_counts={CompressedResourceOp(qml.RX, {}): 2, CompressedResourceOp(qml.RZ, {}): 1},
+            weighted_cost=6.0,
         )
         resources2 = Resources(
-            gate_counts={CompressedResourceOp(qml.RX, {}): 1, CompressedResourceOp(qml.RY, {}): 1}
+            gate_counts={CompressedResourceOp(qml.RX, {}): 1, CompressedResourceOp(qml.RY, {}): 1},
+            weighted_cost=2.0,
         )
 
         resources = resources1 + resources2
@@ -65,12 +78,14 @@ class TestResources:
             CompressedResourceOp(qml.RZ, {}): 1,
             CompressedResourceOp(qml.RY, {}): 1,
         }
+        assert resources.weighted_cost == 8.0
 
     def test_mul_resource_with_scalar(self):
         """Tests multiplying a Resources object with a scalar."""
 
         resources = Resources(
-            gate_counts={CompressedResourceOp(qml.RX, {}): 2, CompressedResourceOp(qml.RZ, {}): 1}
+            gate_counts={CompressedResourceOp(qml.RX, {}): 2, CompressedResourceOp(qml.RZ, {}): 1},
+            weighted_cost=2.0,
         )
 
         resources = resources * 2
@@ -79,14 +94,15 @@ class TestResources:
             CompressedResourceOp(qml.RX, {}): 4,
             CompressedResourceOp(qml.RZ, {}): 2,
         }
+        assert resources.weighted_cost == 4
 
     def test_repr(self):
         """Tests the __repr__ of a Resources object."""
 
         resources = Resources(
-            {CompressedResourceOp(qml.RX, {}): 2, CompressedResourceOp(qml.RZ, {}): 1}
+            {CompressedResourceOp(qml.RX, {}): 2, CompressedResourceOp(qml.RZ, {}): 1}, 5.0
         )
-        assert repr(resources) == "<num_gates=3, gate_counts={RX: 2, RZ: 1}>"
+        assert repr(resources) == "<num_gates=3, gate_counts={RX: 2, RZ: 1}, weighted_cost=5.0>"
 
 
 class DummyOp(qml.operation.Operator):  # pylint: disable=too-few-public-methods
diff --git a/tests/transforms/test_decompose_transform.py b/tests/transforms/test_decompose_transform.py
index 59d726dcf..8bbc8ee74 100644
--- a/tests/transforms/test_decompose_transform.py
+++ b/tests/transforms/test_decompose_transform.py
@@ -83,17 +83,26 @@ class TestDecompose:
             [qml.CNOT([0, 1]), qml.CNOT([1, 0]), qml.CNOT([0, 1])],
             None,
         ),
-        (
-            [qml.Toffoli([0, 1, 2])],
-            {qml.Toffoli},
-            [qml.Toffoli([0, 1, 2])],
-            None,
-        ),
+        ([qml.Toffoli([0, 1, 2])], {qml.Toffoli}, [qml.Toffoli([0, 1, 2])], None),
         (
             [qml.measurements.MidMeasureMP(0)],
             {},
             [qml.measurements.MidMeasureMP(0)],
-            "MidMeasureMP",
+            {
+                "type": TypeError,
+                "msg": "Specifying the gate_set with a dictionary of operator types and their weights is only supported "
+                "with the new experimental graph-based decomposition system. Enable the new system "
+                "using qml.decomposition.enable_graph()",
+            },
+        ),
+        (
+            [qml.Toffoli([0, 1, 2]), qml.measurements.MidMeasureMP(0)],
+            {qml.Toffoli},
+            [qml.Toffoli([0, 1, 2]), qml.measurements.MidMeasureMP(0)],
+            {
+                "type": UserWarning,
+                "msg": "MidMeasureMP",
+            },
         ),
     ]
 
@@ -147,20 +156,24 @@ class TestDecompose:
         with pytest.raises(RecursionError, match=r"Reached recursion limit trying to decompose"):
             decompose(tape, lambda obj: obj.has_matrix)
 
-    @pytest.mark.parametrize("initial_ops, gate_set, expected_ops, warning_pattern", iterables_test)
-    def test_iterable_gate_set(self, initial_ops, gate_set, expected_ops, warning_pattern):
+    @pytest.mark.parametrize(
+        "initial_ops, gate_set, expected_ops, warning_or_error_pattern", iterables_test
+    )
+    def test_iterable_gate_set(self, initial_ops, gate_set, expected_ops, warning_or_error_pattern):
         """Tests that gate sets defined with iterables decompose correctly"""
         tape = qml.tape.QuantumScript(initial_ops)
 
-        if warning_pattern is not None:
-            with pytest.warns(UserWarning, match=warning_pattern):
+        if warning_or_error_pattern is not None:
+            with pytest.raises(
+                warning_or_error_pattern["type"], match=warning_or_error_pattern["msg"]
+            ):
                 (decomposed_tape,), _ = decompose(tape, gate_set=gate_set)
+                expected_tape = qml.tape.QuantumScript(expected_ops)
+                qml.assert_equal(decomposed_tape, expected_tape)
         else:
             (decomposed_tape,), _ = decompose(tape, gate_set=gate_set)
-
-        expected_tape = qml.tape.QuantumScript(expected_ops)
-
-        qml.assert_equal(decomposed_tape, expected_tape)
+            expected_tape = qml.tape.QuantumScript(expected_ops)
+            qml.assert_equal(decomposed_tape, expected_tape)
 
     @pytest.mark.parametrize("initial_ops, gate_set, expected_ops, warning_pattern", callables_test)
     def test_callable_gate_set(self, initial_ops, gate_set, expected_ops, warning_pattern):
diff --git a/tests/transforms/test_decompose_transform_graph.py b/tests/transforms/test_decompose_transform_graph.py
index 239fb1909..bce71dadd 100644
--- a/tests/transforms/test_decompose_transform_graph.py
+++ b/tests/transforms/test_decompose_transform_graph.py
@@ -21,6 +21,20 @@ import pytest
 import pennylane as qml
 
 
+@pytest.mark.unit
+def test_weighted_graph_handles_negative_weight():
+    """Tests a DecompositionGraph raises a ValueError when given negative weights."""
+
+    tape = qml.tape.QuantumScript([])
+
+    # edge case: negative gate weight
+    with pytest.raises(
+        ValueError,
+        match="Negative gate weights provided to gate_set in decompose" "are not supported.",
+    ):
+        qml.transforms.decompose(tape, gate_set={"CNOT": -10.0, "RZ": 1.0})
+
+
 @pytest.mark.unit
 def test_fixed_alt_decomps_not_available():
     """Test that a TypeError is raised when graph is disabled and