diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 1882d0de25..80427ddb1b 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -948,6 +948,9 @@
Bug fixes 🐛
+* Decomposition rules now reconstruct nested operators from their input parameters and wires.
+ This fixes symbolic wire placeholders when compiling GQSP with a nonzero polynomial degree.
+
* The `--convert-to-value-semantics` pass now correctly handles aliasing qubit references that are
used by multiple regions and/or operand lists on the same operation.
[(#3404)](https://github.com/PennyLaneAI/catalyst/pull/3404)
diff --git a/frontend/catalyst/decomposition/decomposition_rules.py b/frontend/catalyst/decomposition/decomposition_rules.py
index 57d8d05fa2..8d0d441d17 100644
--- a/frontend/catalyst/decomposition/decomposition_rules.py
+++ b/frontend/catalyst/decomposition/decomposition_rules.py
@@ -20,12 +20,12 @@
import warnings
from collections import deque
from dataclasses import dataclass
+from enum import Enum, auto
import jax.numpy as jnp
import numpy as np
import pennylane as qp
from jax._src.lib.mlir import ir
-from jax.tree_util import tree_flatten, tree_unflatten
from pennylane.core.operator import Operator2, abstractify
from pennylane.decomposition.utils import to_name
from pennylane.wires import Wires
@@ -224,44 +224,104 @@ def ordered_kwarg_names(call_kwargs, dynamic_shape) -> list:
return params + wires
-def flatten_hybrid_args(extra_data) -> tuple:
- """Split hybrid arguments into numeric ones (passed as operands) and the rest (closed over).
+class HybridLeafSource(Enum):
+ """Where a reconstructed hybrid argument gets each leaf."""
- A hybrid argument whose pytree is entirely numeric arrays (e.g. a ``CDFHamiltonian``) is
- flattened so its leaves can be passed into the rule call as operands, keeping the concrete
- values out of the rule body (where they would otherwise bake in as constants).
+ ORDINARY_PARAMETER = auto()
+ OPERATOR_PARAMETER = auto()
+ WIRE = auto()
+ STATIC = auto()
- Args:
- extra_data (dict): the operator's hybrid arguments, keyed by name
- Returns:
- list[tuple]: ``(name, treedef, num_leaves)`` per operand-passed hybrid argument
- list: the flat leaf dummies (shape/dtype of the concrete leaves), to pass as operands
- dict: the hybrid arguments to close over rather than pass as operands, keyed by name
- """
- specs, leaf_dummies, closed_over = [], [], {}
- for name, value in extra_data.items():
- leaves, treedef = tree_flatten(value)
- # TODO: this "all leaves are numeric arrays" test is a second source of truth for which
- # hybrid leaves become rule-function inputs, and it must agree with the operator lowering's
- # own partition (`_process_params` in from_plxpr/qref_operator2_primitives.py), which uses a
- # per-leaf `forward_mask` to split leaves into `param_map` vs `forward_args` (excluding
- # qubits).
- #
- # Note right now the two agree for all-numeric args (numeric hamiltonians) and for args
- # with a non-array leaf (an operator carrying Wires), but a mixed hybrid arg would diverge
- # causing the lowered op's params to outnumber the rule's inputs.
- # Drive this off `_process_params`/`forward_mask` instead, once the operator's mask is
- # passed through to rule compilation.
- if leaves and all(
- isinstance(leaf, (np.ndarray, np.generic, jnp.ndarray, qp.typing.AbstractArray))
- for leaf in leaves
- ):
- specs.append((name, treedef, len(leaves)))
- leaf_dummies.extend(jnp.zeros(leaf.shape, dtype=leaf.dtype) for leaf in leaves)
+@dataclass
+class HybridArgumentSpec:
+ """Describe how to rebuild a hybrid argument from rule inputs."""
+
+ name: str
+ argument_tree: object
+ original_leaves: list
+ leaf_sources: list[HybridLeafSource]
+
+ def rebuild(self, ordinary_parameters, operator_parameters, wires):
+ """Replace abstract leaves with the rule's parameter and wire inputs."""
+ reconstructed_leaves = []
+ wire_offset = 0
+ for leaf, source in zip(self.original_leaves, self.leaf_sources, strict=True):
+ if source == HybridLeafSource.ORDINARY_PARAMETER:
+ reconstructed_leaves.append(next(ordinary_parameters))
+ elif source == HybridLeafSource.OPERATOR_PARAMETER:
+ reconstructed_leaves.append(next(operator_parameters))
+ elif source == HybridLeafSource.WIRE:
+ wire_slice = wires[wire_offset : wire_offset + len(leaf)]
+ reconstructed_leaves.append(Wires(wire_slice))
+ wire_offset += len(leaf)
+ else:
+ reconstructed_leaves.append(leaf)
+
+ # Constructing the carrier object must not apply another gate.
+ with qp.capture.pause():
+ return qp.pytrees.unflatten(reconstructed_leaves, self.argument_tree)
+
+
+def _is_wire_leaf(value):
+ return isinstance(value, (Wires, qp.typing.AbstractWires))
+
+
+def flatten_hybrid_args(extra_data) -> tuple:
+ """Prepare hybrid arguments for reconstruction from rule inputs."""
+ argument_specs = []
+ ordinary_parameter_dummies = []
+ operator_parameter_dummies = []
+ static_arguments = {}
+
+ numeric_types = (np.ndarray, np.generic, jnp.ndarray, qp.typing.AbstractArray)
+
+ for argument_name, argument in extra_data.items():
+ # Keep the complete structure so we can rebuild the argument.
+ original_leaves, argument_tree = qp.pytrees.flatten(argument, is_leaf=_is_wire_leaf)
+
+ # Stop at operators to identify which parameters belong to them.
+ argument_parts, _ = qp.pytrees.flatten(
+ argument,
+ is_leaf=lambda value: isinstance(value, Operator2) or _is_wire_leaf(value),
+ )
+
+ leaf_sources = []
+ for argument_part in argument_parts:
+ is_operator = isinstance(argument_part, Operator2)
+ part_leaves, _ = qp.pytrees.flatten(argument_part, is_leaf=_is_wire_leaf)
+
+ for leaf in part_leaves:
+ if isinstance(leaf, numeric_types):
+ dummy_parameter = jnp.zeros(leaf.shape, dtype=leaf.dtype)
+ if is_operator:
+ source = HybridLeafSource.OPERATOR_PARAMETER
+ operator_parameter_dummies.append(dummy_parameter)
+ else:
+ source = HybridLeafSource.ORDINARY_PARAMETER
+ ordinary_parameter_dummies.append(dummy_parameter)
+ elif _is_wire_leaf(leaf):
+ source = HybridLeafSource.WIRE
+ else:
+ source = HybridLeafSource.STATIC
+
+ leaf_sources.append(source)
+
+ if any(source != HybridLeafSource.STATIC for source in leaf_sources):
+ argument_specs.append(
+ HybridArgumentSpec(
+ name=argument_name,
+ argument_tree=argument_tree,
+ original_leaves=original_leaves,
+ leaf_sources=leaf_sources,
+ )
+ )
else:
- closed_over[name] = value
- return specs, leaf_dummies, closed_over
+ static_arguments[argument_name] = argument
+
+ # Match gate lowering: ordinary parameters precede operator parameters.
+ dummy_parameters = ordinary_parameter_dummies + operator_parameter_dummies
+ return argument_specs, dummy_parameters, static_arguments
def rule_call_operands(
@@ -321,7 +381,7 @@ def unpack_rule_operands(
):
"""Recover ``(params, kwargs, control wires)`` from :func:`rule_call_operands`' flattening.
- Reconstructs each hybrid argument from its traced leaves (via its ``treedef``) so the rule body
+ Reconstructs each hybrid argument from its traced leaves and wires so the rule body
receives the same object the user passed, but built from the traced operands.
Args:
@@ -330,8 +390,8 @@ def unpack_rule_operands(
kwarg_names (list[str]): parameter names first, followed by wire names
wire_lens (dict[str, int]): wire names in operand order and their lengths
has_ctrl_wires (bool): whether a trailing control-wire operand is present
- hybrid_specs (list[tuple]): ``(name, treedef, num_leaves)`` per hybrid argument, sitting
- between the parameter kwargs and the grouped wire operand
+ hybrid_specs (list[HybridArgumentSpec]): reconstruction plans identifying each
+ argument's ordinary parameters, operator parameters, wires, and static leaves
Returns:
tuple: the rule's positional arguments
@@ -343,9 +403,16 @@ def unpack_rule_operands(
n_kwarg_params = len(param_names)
named = dict(zip(param_names, operands[n_params : n_params + n_kwarg_params], strict=True))
idx = n_params + n_kwarg_params
- for name, treedef, num_leaves in hybrid_specs:
- named[name] = tree_unflatten(treedef, operands[idx : idx + num_leaves])
- idx += num_leaves
+ num_ordinary_parameters = sum(
+ spec.leaf_sources.count(HybridLeafSource.ORDINARY_PARAMETER) for spec in hybrid_specs
+ )
+ num_operator_parameters = sum(
+ spec.leaf_sources.count(HybridLeafSource.OPERATOR_PARAMETER) for spec in hybrid_specs
+ )
+ ordinary_parameters = iter(operands[idx : idx + num_ordinary_parameters])
+ idx += num_ordinary_parameters
+ operator_parameters = iter(operands[idx : idx + num_operator_parameters])
+ idx += num_operator_parameters
grouped_wires = operands[idx]
idx += 1
@@ -354,6 +421,11 @@ def unpack_rule_operands(
named[name] = grouped_wires[offset : offset + length]
offset += length
+ for spec in hybrid_specs:
+ named[spec.name] = spec.rebuild(
+ ordinary_parameters, operator_parameters, wires=named.get(spec.name, ())
+ )
+
ctrl_wires = operands[idx] if has_ctrl_wires else None
ctrl_values = operands[idx + 1] if has_ctrl_wires and len(operands) > idx + 1 else None
if has_ctrl_wires:
diff --git a/frontend/test/lit/test_operator2/test_decomposition_rule_entry_point.py b/frontend/test/lit/test_operator2/test_decomposition_rule_entry_point.py
index 0d74d85acd..31d1d3df0b 100644
--- a/frontend/test/lit/test_operator2/test_decomposition_rule_entry_point.py
+++ b/frontend/test/lit/test_operator2/test_decomposition_rule_entry_point.py
@@ -461,21 +461,21 @@ def rule(cwires):
qp.add_decomps(HybridWires, rule)
result = compile_decomposition_rules_wrapper(
"HybridWires",
- "HybridWires{}{}{}[3742]",
+ "HybridWires{}{cwires:3}{}[3742]",
{},
+ {"cwires": 3},
{},
- {},
- extra_data={"cwires": [1, [2, qp.wires.Wires(3)]]},
+ extra_data={"cwires": [Wire[1], [Wire[1], Wire[1]]]},
)
print(result)
-# CHECK: func.func private @rule{{(_[0-9]+)?}}
+# CHECK: func.func private @"rule_HybridWires{}{cwires:3}{}[3742]"
# CHECK-SAME: resources = {operations = {"NoParams{}{reg:1}{}" = 3 : i64}}
-# CHECK-SAME: target_gate = "HybridWires{}{}{}[3742]"
-# CHECK: idx_attr = 1
-# CHECK: idx_attr = 2
-# CHECK: idx_attr = 3
+# CHECK-SAME: target_gate = "HybridWires{}{cwires:3}{}[3742]"
+# CHECK: stablehlo.slice %arg1 [0:1]
+# CHECK: stablehlo.slice %arg1 [1:2]
+# CHECK: stablehlo.slice %arg1 [2:3]
test_from_hybrid_wires()
@@ -562,23 +562,23 @@ def rule(angle, op, cwires, n_iters):
qp.add_decomps(HybridOpArg, rule)
result = compile_decomposition_rules_wrapper(
"HybridOpArg",
- "HybridOpArg{angle:[tensor]}{cwires:1}{}[5678]",
+ "HybridOpArg{angle:[tensor]}{cwires:1,op:3}{}[5678]",
{"angle": ["f64"]},
- {"cwires": 1},
+ {"cwires": 1, "op": 3},
{},
extra_data={
- "op": StaticDataMultiReg("hello", reg=[1], reg2=[2, 3], theta=0.2),
+ "op": StaticDataMultiReg("hello", reg=Wire[1], reg2=Wire[2], theta=Float),
"n_iters": 100,
},
)
print(result)
-# CHECK: func.func private @rule{{(_[0-9]+)?}}
+# CHECK: func.func private @"rule_HybridOpArg{angle:[tensor]}{cwires:1,op:3}{}[5678]"
# CHECK-SAME: resources = {operations = {
# CHECK-SAME: "NoParams{}{reg:1}{}" = 1 : i64,
# CHECK-SAME: "StaticDataMultiReg{theta:[tensor]}{reg:1,reg2:2}{}[[[uid:[0-9]+]]]" = 1 : i64
-# CHECK-SAME: target_gate = "HybridOpArg{angle:[tensor]}{cwires:1}{}[5678]"
+# CHECK-SAME: target_gate = "HybridOpArg{angle:[tensor]}{cwires:1,op:3}{}[5678]"
# CHECK: "qref.operator"
# CHECK-SAME: UID = [[uid]] : i64, op_name = "StaticDataMultiReg"
test_from_hybrid_op()
@@ -653,15 +653,15 @@ def rule(angle, op, cwires, n_iters):
qp.add_decomps(HybridOpArg, rule)
result = compile_decomposition_rules_wrapper(
"HybridOpArg",
- "HybridOpArg{angle:[tensor]}{cwires:1}{}[7654]",
+ "HybridOpArg{angle:[tensor]}{cwires:1,op:4}{}[7654]",
{"angle": ["f64"]},
- {"cwires": 1},
+ {"cwires": 1, "op": 4},
{},
extra_data={
"op": HybridOpArg(
- angle=0.2,
- op=StaticDataMultiReg("hello", reg=[1], reg2=[2, 3], theta=0.2),
- cwires=2,
+ angle=Float,
+ op=StaticDataMultiReg("hello", reg=Wire[1], reg2=Wire[2], theta=Float),
+ cwires=Wire[1],
n_iters=200,
),
"n_iters": 100,
@@ -670,12 +670,12 @@ def rule(angle, op, cwires, n_iters):
print(result)
-# CHECK: func.func private @rule{{(_[0-9]+)?}}
+# CHECK: func.func private @"rule_HybridOpArg{angle:[tensor]}{cwires:1,op:4}{}[7654]"
# CHECK-SAME: resources = {operations = {
# CHECK-SAME: "HybridOpArg{angle:[tensor]}{cwires:1,op:3}{}[[[uid_outer:[0-9]+]]]" = 1 : i64,
# CHECK-SAME: "NoParams{}{reg:1}{}" = 1 : i64,
# CHECK-SAME: "StaticDataMultiReg{theta:[tensor]}{reg:1,reg2:2}{}[[[uid_inner:[0-9]+]]]" = 1 : i64
-# CHECK-SAME: target_gate = "HybridOpArg{angle:[tensor]}{cwires:1}{}[7654]"
+# CHECK-SAME: target_gate = "HybridOpArg{angle:[tensor]}{cwires:1,op:4}{}[7654]"
# CHECK: "qref.operator"
# CHECK-SAME: UID = [[uid_outer]] : i64, op_name = "HybridOpArg"
# CHECK-SAME: qubit_map = {cwires = array, op = array}
@@ -772,28 +772,28 @@ def rule(reg1, reg2, angles1, angles2, pytree1, pytree2, op1, op2, hwires1, hwir
qp.add_decomps(MultipleFullArgs, rule)
result = compile_decomposition_rules_wrapper(
"MultipleFullArgs",
- "MultipleFullArgs{angles1:[tensor],angles2:[tensor<2xf64>]}{reg1:1,reg2:2}{}[4444]",
+ "MultipleFullArgs{angles1:[tensor],angles2:[tensor<2xf64>]}{hwires1:2,hwires2:1,op1:1,op2:1,reg1:1,reg2:2}{}[4444]",
{"angles1": ["f64"], "angles2": ["f64", "f64"]},
- {"reg1": 1, "reg2": 2},
+ {"reg1": 1, "reg2": 2, "op1": 1, "op2": 1, "hwires1": 2, "hwires2": 1},
{},
extra_data={
"pytree1": [1],
"pytree2": [2],
- "op1": SingleParam(x=1.1, reg=[1]),
- "op2": SingleParam(x=2, reg=[2]),
- "hwires1": [qp.wires.Wires(3), qp.wires.Wires(2)],
- "hwires2": [qp.wires.Wires(1)],
+ "op1": SingleParam(x=Float, reg=Wire[1]),
+ "op2": SingleParam(x=Int, reg=Wire[1]),
+ "hwires1": [Wire[1], Wire[1]],
+ "hwires2": [Wire[1]],
},
)
print(result)
-# CHECK: func.func private @rule{{(_[0-9]+)?}}
+# CHECK: func.func private @"rule_MultipleFullArgs{angles1:[tensor],angles2:[tensor<2xf64>]}{hwires1:2,hwires2:1,op1:1,op2:1,reg1:1,reg2:2}{}[4444]"
# CHECK-SAME: resources = {operations = {
# CHECK-SAME: "NoParams{}{reg:1}{}" = 1 : i64
# CHECK-SAME: "SingleParam{x:[tensor]}{reg:1}{}" = 1 : i64
# CHECK-SAME: "SingleParam{x:[tensor]}{reg:1}{}" = 1 : i64
-# CHECK-SAME: target_gate = "MultipleFullArgs{angles1:[tensor],angles2:[tensor<2xf64>]}{reg1:1,reg2:2}{}[4444]"
+# CHECK-SAME: target_gate = "MultipleFullArgs{angles1:[tensor],angles2:[tensor<2xf64>]}{hwires1:2,hwires2:1,op1:1,op2:1,reg1:1,reg2:2}{}[4444]"
# CHECK: "qref.operator"
# CHECK-SAME: op_name = "SingleParam"
# CHECK: "qref.operator"
diff --git a/frontend/test/pytest/decomposition/test_decomposition.py b/frontend/test/pytest/decomposition/test_decomposition.py
index 3ea75071b1..e70badaf97 100644
--- a/frontend/test/pytest/decomposition/test_decomposition.py
+++ b/frontend/test/pytest/decomposition/test_decomposition.py
@@ -59,6 +59,7 @@
compile_decomposition_rules_wrapper,
compile_reachable_decomposition_rules_wrapper,
compile_registered_symbolic_rules,
+ flatten_hybrid_args,
get_rule_strings_from_module,
materialize_reachable_rule_strings,
ordered_kwarg_names,
@@ -154,6 +155,53 @@ def test_unpack_rule_operands_restores_wire_kwargs(self):
np.testing.assert_array_equal(ctrl_wires, [5, 6])
np.testing.assert_array_equal(ctrl_values, [True, False])
+ def test_hybrid_operands_preserve_parameters_and_wires(self):
+ """Nested parameters follow ordinary hybrid parameters; wires retain their mapping."""
+ extra = {
+ "unitary": abstractify(qp.RX(0.1, 0) @ qp.RY(0.2, 1)),
+ "data": Float[2],
+ }
+ specs, dummies, closed = flatten_hybrid_args(extra)
+ assert not closed
+ assert [dummy.shape for dummy in dummies] == [(2,), (), ()]
+ _, kwargs, _, _ = unpack_rule_operands(
+ (jnp.array([4.0, 5.0]), jnp.array(0.7), jnp.array(-0.2), jnp.array([3, 1])),
+ n_params=0,
+ kwarg_names=["unitary"],
+ wire_lens={"unitary": 2},
+ has_ctrl_wires=False,
+ hybrid_specs=specs,
+ )
+ np.testing.assert_array_equal(kwargs["data"], [4.0, 5.0])
+ np.testing.assert_allclose(
+ qp.matrix(kwargs["unitary"], wire_order=[3, 1]),
+ qp.matrix(qp.RX(0.7, 3) @ qp.RY(-0.2, 1), wire_order=[3, 1]),
+ )
+
+ def test_hybrid_operands_preserve_static_leaves(self):
+ """Rebuilding a mixed hybrid argument preserves its static data and runtime values."""
+ extra = {
+ "mixed": {"angle": Float, "label": "rotation", "wires": Wire[2]},
+ "static": ("unchanged", 7),
+ }
+ specs, dummies, closed = flatten_hybrid_args(extra)
+ assert closed == {"static": ("unchanged", 7)}
+ assert len(dummies) == 1
+ assert dummies[0].shape == ()
+
+ for angle, wires in ((0.7, [3, 1]), (-0.2, [1, 4])):
+ _, kwargs, _, _ = unpack_rule_operands(
+ (jnp.array(angle), jnp.array(wires)),
+ n_params=0,
+ kwarg_names=["mixed"],
+ wire_lens={"mixed": 2},
+ has_ctrl_wires=False,
+ hybrid_specs=specs,
+ )
+ assert kwargs["mixed"]["label"] == "rotation"
+ np.testing.assert_allclose(kwargs["mixed"]["angle"], angle)
+ assert list(kwargs["mixed"]["wires"]) == wires
+
def test_probe_wires_dont_overlap(self):
"""Test that the helper for generating probe arguments doesnt create
overlapping wires.
@@ -860,6 +908,34 @@ def decomposed():
assert np.allclose(np.asarray(decomposed()), np.asarray(reference()))
+ @pytest.mark.parametrize("kind", ["x", "rx", "ising"])
+ @pytest.mark.parametrize("shift", [0, -1])
+ def test_gqsp_nested_unitary_operands(self, kind, shift):
+ """The controlled unitary in a GQSP loop uses the invocation's wires and parameters."""
+
+ def circuit(theta, angles):
+ qp.Hadamard(2)
+ if kind == "x":
+ unitary = qp.X(2)
+ elif kind == "rx":
+ unitary = qp.RX(theta, 2)
+ else:
+ unitary = qp.IsingXX(theta, wires=[2, 1])
+ qp.GQSP(unitary, angles, control=0, shift=shift)
+ return qp.probs(wires=[0, 1, 2])
+
+ reference = qp.QNode(circuit, qp.device("default.qubit", wires=3))
+ quantum_function = qp.QNode(circuit, qp.device("lightning.qubit", wires=3))
+ compiled = qjit(
+ qp.decompose(
+ gate_set={"RX", "RY", "RZ", "PauliX", "PauliZ", "Hadamard", "CNOT", "GlobalPhase"}
+ )(quantum_function),
+ capture=True,
+ )
+ angles = np.array([[0.1, -0.2, 0.4], [0.3, 0.5, -0.1], [-0.7, 0.2, 0.6]])
+ for theta in (0.17, -0.62):
+ np.testing.assert_allclose(compiled(theta, angles), reference(theta, angles), atol=1e-7)
+
class TestOnDemand:
"""Test the python wrapper functions used for on-demand,