Merge branch 'main' into ops-qp

Author: andrijapau

doc/releases/changelog-dev.md

@@ -175,7 +175,21 @@ The following classes have been ported over:
 
 <h3>Improvements 🛠</h3>
 
-* Replaced the O(n²) incremental ``@=`` operator chaining in ``qp.pauli.string_to_pauli_word`` and ``qp.pauli.binary_to_pauli`` with a single ``qp.prod(*tuple_of_ops)`` call, collecting operators via generator expressions. These operators are now much faster for large Pauli strings.
+* With program capture and `for_loop` and `while_loop`, const closure variables with dynamic shapes
+  can now be combined with explicit inputs with dynamic shapes when they have matching shapes.
+  [(#9275)](https://github.com/PennyLaneAI/pennylane/pull/9275)
+
+* Added another decomposition to `MultiControlledX` with two control wires and at least one zeroed
+  work wire that has been passed explicitly. It decomposes into a pair of `TemporaryAND` and a
+  `CNOT`.
+  [(#9291)](https://github.com/PennyLaneAI/pennylane/pull/9291)
+
+* Operations using ``FermiWord`` are now much faster due to various performance improvements to the class
+  [(#9283)](https://github.com/PennyLaneAI/pennylane/pull/9283)
+
+* Replaced the O(n²) incremental ``@=`` operator chaining in ``qp.pauli.string_to_pauli_word`` and
+  ``qp.pauli.binary_to_pauli`` with a single ``qp.prod(*tuple_of_ops)`` call, collecting operators via
+  generator expressions. These operators are now much faster for large Pauli strings.
   [(#9271)](https://github.com/PennyLaneAI/pennylane/pull/9271)
 
 * Operations using ``PauliSentence`` are now much faster due to additional memoization in ``PauliWord.__hash__``
@@ -789,6 +803,9 @@ The following classes have been ported over:
 
 <h3>Internal changes ⚙️</h3>
 
+* During program, `qml.for_loop` with negative step sizes is now handled immediately during capture time.
+  [(#9299)](https://github.com/PennyLaneAI/pennylane/pull/9299)
+
 * With program capture, arrays dynamic shapes with `qml.for_loop` and `qml.while_loop` can now be combined
   after the loop.
   [(#9245)](https://github.com/PennyLaneAI/pennylane/pull/9245)
@@ -975,6 +992,10 @@ The following classes have been ported over:
 
 <h3>Bug fixes 🐛</h3>
 
+* Global phases are now supported in `from_qasm3` so that QASM including the `gphase` instruction 
+  can be interpreted.
+  [(#9247)](https://github.com/PennyLaneAI/pennylane/pull/9247)
+
 * Fixes an issue with Catalyst and `qml.for_loop` and `qml.while_loop`, where it was defaulting
   to `allow_array_resizing=True` instead of `allow_array_resizing=False`.
   [(#9251)](https://github.com/PennyLaneAI/pennylane/pull/9251)
@@ -1094,6 +1115,10 @@ The following classes have been ported over:
   produced for a specific wire configuration.
   [(#9270)](https://github.com/PennyLaneAI/pennylane/pull/9270)
 
+* Fixes a bug where the `DecompositionGraph` underestimates the minimum number of work wires required to solve for a particular operator
+  when it has decomposition rules with a lower work wire budget but is unrecheable from the provided gate set.
+  [(#9298)](https://github.com/PennyLaneAI/pennylane/pull/9298)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):

pennylane/_version.py

@@ -16,4 +16,4 @@
 Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.45.0-dev77"
+__version__ = "0.45.0-dev80"

pennylane/control_flow/_loop_abstract_axes.py

@@ -29,6 +29,44 @@
 AbstractShapeLocation = namedtuple("AbstractShapeLocation", ("arg_idx", "shape_idx"))
 
 
+def promote_consts_to_inputs(f):
+    """This function extracts any closure variables with dynamic shapes from f.__closure__
+    and promotes them to being normal arguments. This produces a new function that
+    takes the original args and the new consts as explicit inputs. It also returns
+    the extracted consts.
+    """
+    indices = []
+    consts = []
+
+    if getattr(f, "__closure__", None) is not None:
+        for ind, cell in enumerate(f.__closure__):
+            val = cell.cell_contents
+            if hasattr(val, "shape") and not all(isinstance(s, int) for s in val.shape):
+                indices.append(ind)
+                consts.append(val)
+
+    def new_f(args, new_consts):
+        """A version of f where the consts with dynamic shapes have been promoted to inputs."""
+
+        # even deepcopy does not actually copy the closure for a function
+        # we don't have a way to produce a new function with independent closure
+        # so therefore we just need to make sure to clean up after ourselves after
+        # in-place modifying the closure contents.
+
+        try:
+            for ind, c in zip(indices, new_consts, strict=True):
+                f.__closure__[ind].cell_contents = c
+
+            f_results = f(*args)
+        finally:
+            for ind, c in zip(indices, consts, strict=True):
+                f.__closure__[ind].cell_contents = c
+
+        return f_results, new_consts
+
+    return new_f, consts
+
+
 def add_abstract_shapes(f, shape_locations: list[list[AbstractShapeLocation]]):  # pragma: no cover
     """Add the abstract shapes at the specified locations to the output of f.
 

pennylane/control_flow/for_loop.py

@@ -18,7 +18,7 @@
 import warnings
 from typing import Literal
 
-from pennylane import capture
+from pennylane import capture, math
 from pennylane.capture import FlatFn, enabled
 from pennylane.capture.dynamic_shapes import register_custom_staging_rule
 from pennylane.compiler.compiler import AvailableCompilers, active_compiler
@@ -29,13 +29,33 @@
     get_dummy_arg,
     handle_jaxpr_error,
     loop_determine_abstracted_axes,
+    promote_consts_to_inputs,
     validate_no_resizing_returns,
 )
 
 logger = logging.getLogger(__name__)
 logger.addHandler(logging.NullHandler())
 
 
+def _reverse_iterator(f, start, step):
+    """Produces a new f with positive steps."""
+
+    def new_f(*args):
+        new_i = start + step * args[0]
+        inputs = args[1:]
+        return f(new_i, *inputs)
+
+    return new_f
+
+
+def _is_reverse_iteration(start, stop, step):
+    # without the int() call, when we have a jnp array with a single int
+    # in it (jnp.array(1)), performing a comparison will produce a tracer
+    if not math.is_abstract(step):
+        return int(step) < 0
+    return not math.is_abstract(start) and not math.is_abstract(stop) and int(stop) < int(start)
+
+
 def for_loop(
     start, stop=None, step=1, *, allow_array_resizing: Literal["auto", True, False] = "auto"
 ):
@@ -397,8 +417,10 @@ def _get_jaxpr(self, init_state, allow_array_resizing):
 
         import jax  # pylint: disable=import-outside-toplevel
 
+        f_consts_extracted, dynamic_consts = promote_consts_to_inputs(self.body_fn)
+
         # need in_tree to include index so flat_fn will repack args correctly
-        flat_args, in_tree = jax.tree_util.tree_flatten((0, *init_state))
+        flat_args, in_tree = jax.tree_util.tree_flatten(((0, *init_state), dynamic_consts))
 
         # slice out the index so shape_locations indexes from non-index args/ results
         flat_args = flat_args[1:]
@@ -407,7 +429,7 @@ def _get_jaxpr(self, init_state, allow_array_resizing):
             tuple(flat_args), allow_array_resizing=tmp_array_resizing
         )
 
-        flat_fn = FlatFn(self.body_fn, in_tree=in_tree)
+        flat_fn = FlatFn(f_consts_extracted, in_tree=in_tree)
 
         if abstracted_axes:  # pragma: no cover
             new_body_fn = add_abstract_shapes(flat_fn, shape_locations)
@@ -417,6 +439,10 @@ def _get_jaxpr(self, init_state, allow_array_resizing):
             new_body_fn = flat_fn
             dummy_init_state = flat_args
 
+        if _is_reverse_iteration(self.start, self.stop, self.step):
+            # MLIR does not support reverse iteration of for loops
+            new_body_fn = _reverse_iterator(new_body_fn, self.start, self.step)
+
         try:
             jaxpr_body_fn = jax.make_jaxpr(new_body_fn, abstracted_axes=abstracted_axes)(
                 0, *dummy_init_state
@@ -469,10 +495,17 @@ def _call_capture_enabled(self, *init_state):
         abstract_shapes_slice = slice(consts_slice.stop, consts_slice.stop + len(abstract_shapes))
         args_slice = slice(abstract_shapes_slice.stop, None)
 
+        if _is_reverse_iteration(self.start, self.stop, self.step):
+            # mlir does not support reverse iteration of for loops
+            num_iterations = math.ceil((self.stop - self.start) / self.step).astype(int)
+            start, stop, step = 0, num_iterations, 1
+        else:
+            start, stop, step = self.start, self.stop, self.step
+
         results = for_loop_prim.bind(
-            self.start,
-            self.stop,
-            self.step,
+            start,
+            stop,
+            step,
             *jaxpr_body_fn.consts,
             *abstract_shapes,
             *flat_args,
@@ -483,7 +516,8 @@ def _call_capture_enabled(self, *init_state):
         )
 
         results = results[-out_tree.num_leaves :]
-        return jax.tree_util.tree_unflatten(out_tree, results)
+        # [0] to slice out the consts extracted by promote_consts_to_inputs
+        return jax.tree_util.tree_unflatten(out_tree, results)[0]
 
     def __call__(self, *init_state):
 

pennylane/control_flow/while_loop.py

@@ -27,6 +27,7 @@
     get_dummy_arg,
     handle_jaxpr_error,
     loop_determine_abstracted_axes,
+    promote_consts_to_inputs,
     validate_no_resizing_returns,
 )
 
@@ -43,6 +44,14 @@ def _new_cond_fn(*args, **kwargs):
     return _new_cond_fn
 
 
+def _body_consts_extracted_cond(cond_fn):
+    # pylint: disable=unused-argument
+    def new_cond_fn(args, body_consts):
+        return cond_fn(*args)
+
+    return new_cond_fn
+
+
 def while_loop(cond_fn, allow_array_resizing: Literal["auto", True, False] = "auto"):
     """A :func:`~.qjit` compatible while-loop for PennyLane programs. When
     used without :func:`~.qjit` or program capture, this function will fall back to a standard
@@ -317,14 +326,16 @@ def _call_capture_disabled(self, *init_state):
     def _get_jaxprs(self, init_state, allow_array_resizing):
         import jax  # pylint: disable=import-outside-toplevel
 
-        flat_args, in_tree = jax.tree_util.tree_flatten(init_state)
+        body_consts_extracted, dynamic_consts = promote_consts_to_inputs(self.body_fn)
+
+        flat_args, in_tree = jax.tree_util.tree_flatten((init_state, dynamic_consts))
         tmp_array_resizing = False if allow_array_resizing == "auto" else allow_array_resizing
         abstracted_axes, abstract_shapes, shape_locations = loop_determine_abstracted_axes(
             tuple(flat_args), allow_array_resizing=tmp_array_resizing
         )
 
-        flat_body_fn = FlatFn(self.body_fn, in_tree=in_tree)
-        flat_cond_fn = FlatFn(self.cond_fn, in_tree=in_tree)
+        flat_body_fn = FlatFn(body_consts_extracted, in_tree=in_tree)
+        flat_cond_fn = FlatFn(_body_consts_extracted_cond(self.cond_fn), in_tree=in_tree)
         bool_cond_fn = _to_bool_cond_fn(flat_cond_fn)
 
         if abstracted_axes:  # pragma: no cover
@@ -379,7 +390,8 @@ def _call_capture_enabled(self, *init_state):
         )
 
         results = results[-out_tree.num_leaves :]
-        return jax.tree_util.tree_unflatten(out_tree, results)
+        # [0] to slice out the consts extracted by promote_consts_to_inputs
+        return jax.tree_util.tree_unflatten(out_tree, results)[0]
 
     def __call__(self, *init_state):
 

pennylane/decomposition/decomposition_graph.py

@@ -97,6 +97,9 @@ class _OperatorNode:
     min_work_wires: int = 0
     """The minimum number of additional work wires required to decompose this operator."""
 
+    reachable: bool = False
+    """Whether the operator node can be reached from the gate set."""
+
     def __hash__(self) -> int:
         # If the decomposition of an operator does not depend on the availability of work wires
         # at all, we don't need to have multiple nodes representing the same operator with
@@ -130,6 +133,7 @@ class _DecompositionNode:
     num_work_wire_not_available: int
     work_wire_dependent: bool = False
     min_work_wires: int = 0
+    reachable: bool = True
 
     def __post_init__(self):
         self.min_work_wires = self.min_work_wires or self.work_wire_spec.total
@@ -309,6 +313,9 @@ def _add_op_node(self, op: CompressedResourceOp, num_used_work_wires: int) -> in
         if op_node in self._all_op_indices:
             return self._all_op_indices[op_node]
 
+        if op in self._gate_set_weights:
+            op_node = replace(op_node, reachable=True)
+
         op_node_idx = self._graph.add_node(op_node)
         self._all_op_indices[op_node] = op_node_idx
         self._op_to_op_nodes[op].add(op_node)
@@ -333,23 +340,32 @@ def _add_op_node(self, op: CompressedResourceOp, num_used_work_wires: int) -> in
             self._graph.add_edge(self._start, op_node_idx, math.inf)
             return op_node_idx
 
+        op_reachable = False
         work_wire_dependent = known_work_wire_dependent
         min_work_wires = -1  # use -1 to represent undetermined work wire requirement
         for decomposition in rules:
             d_node = self._add_decomp(decomposition, op_node, op_node_idx, num_used_work_wires)
+            if not d_node or not d_node.reachable:
+                continue
+            # If a decomposition is reachable, the operator is also reachable
+            op_reachable = True
             # If any of the operator's decompositions depend on work wires, this operator
             # should also depend on work wires.
-            if d_node and d_node.work_wire_dependent:
+            if d_node.work_wire_dependent:
                 work_wire_dependent = True
-            if d_node and (min_work_wires == -1 or d_node.min_work_wires < min_work_wires):
+            if min_work_wires == -1 or d_node.min_work_wires < min_work_wires:
                 min_work_wires = d_node.min_work_wires
 
+        if op_reachable:
+            op_node = replace(op_node, reachable=True)
+            self._replace_node(op_node_idx, op_node)
+
         # If we found that this operator depends on work wires, but it's currently recorded
         # as independent of work wires, we must replace every record of this operator node
         # with a new node with `work_wire_dependent` set to `True`.
         if not known_work_wire_dependent and work_wire_dependent:
-            new_op_node = replace(op_node, work_wire_dependent=True, min_work_wires=min_work_wires)
-            self._replace_node(op_node_idx, new_op_node)
+            op_node = replace(op_node, work_wire_dependent=True, min_work_wires=min_work_wires)
+            self._replace_node(op_node_idx, op_node)
             # Also record that this operator type depends on work wires, so in the future
             # when we encounter other instances of the same operator type, we correctly
             # identify it as work-wire dependent.
@@ -400,6 +416,8 @@ def _add_decomp(
             # decomposition is also dependent on work wires, even it itself does not use
             # any work wires.
             op_node = self._graph[op_node_idx]
+            if not op_node.reachable:
+                d_node.reachable = False
             if op_node.work_wire_dependent:
                 d_node.work_wire_dependent = True
             max_op_min_work_wires = max(op_node.min_work_wires, max_op_min_work_wires)