Add autograph support for for loops

doc/releases/changelog-dev.md

@@ -90,9 +90,10 @@ added `binary_mapping()` function to map `BoseWord` and `BoseSentence` to qubit
 * Added submodule `devices.qubit_mixed.measure` as a necessary step for the new API, featuring a `measure` function for measuring qubits in mixed-state devices.
   [(#6637)](https://github.com/PennyLaneAI/pennylane/pull/6637)
 
-* Support is added for `if`/`else` statements and `while` loops in circuits executed with `qml.capture.enabled`, via `autograph`.
+* Support is added for `if`/`else` statements and `for` and `while` loops in circuits executed with `qml.capture.enabled`, via `autograph`
   [(#6406)](https://github.com/PennyLaneAI/pennylane/pull/6406)
   [(#6413)](https://github.com/PennyLaneAI/pennylane/pull/6413)
+  [(#6426)](https://github.com/PennyLaneAI/pennylane/pull/6426)
 
 * Added `christiansen_mapping()` function to map `BoseWord` and `BoseSentence` to qubit operators, using christiansen mapping.
   [(#6623)](https://github.com/PennyLaneAI/pennylane/pull/6623)

pennylane/capture/autograph/ag_primitives.py

@@ -18,24 +18,27 @@
 """
 import copy
 import functools
-from typing import Any, Callable, Tuple
+from typing import Any, Callable, Iterator, SupportsIndex, Tuple, Union
 
 from malt.core import config as ag_config
 from malt.impl import api as ag_api
 from malt.impl.api import converted_call as ag_converted_call
+from malt.operators import py_builtins as ag_py_builtins
 from malt.operators.variables import Undefined
 
 import pennylane as qml
 
 has_jax = True
 try:
     import jax
+    import jax.numpy as jnp
 except ImportError:  # pragma: no cover
     has_jax = False
 
 
 __all__ = [
     "if_stmt",
+    "for_stmt",
     "while_stmt",
     "converted_call",
 ]
@@ -58,8 +61,6 @@ def _assert_results(results, var_names):
 
 
 # pylint: disable=too-many-arguments, too-many-positional-arguments
-
-
 def if_stmt(
     pred: bool,
     true_fn: Callable[[], Any],
@@ -137,6 +138,157 @@ def _assert_iteration_inputs(inputs, symbol_names):
             ) from e
 
 
+def _assert_iteration_results(inputs, outputs, symbol_names):
+    """The results of a for loop should have the identical type as the inputs since they are
+    "passed" as inputs to the next iteration. A mismatch here may indicate that a loop-carried
+    variable was initialized with the wrong type.
+    """
+
+    for i, (inp, out) in enumerate(zip(inputs, outputs)):
+        inp_t, out_t = jax.api_util.shaped_abstractify(inp), jax.api_util.shaped_abstractify(out)
+        if inp_t.dtype != out_t.dtype or inp_t.shape != out_t.shape:
+            raise AutoGraphError(
+                f"The variable '{symbol_names[i]}' was initialized with the wrong type, or you may "
+                f"be trying to change its type from one iteration to the next. "
+                f"Expected: {out_t}, Got: {inp_t}"
+            )
+
+
+# pylint: disable=too-many-positional-arguments
+def _call_pennylane_for(
+    start,
+    stop,
+    step,
+    body_fn,
+    get_state,
+    set_state,
+    symbol_names,
+    enum_start=None,
+    array_iterable=None,
+):
+    """Dispatch to a PennyLane implementation of for loops."""
+
+    # Ensure iteration arguments are properly initialized. We cannot process uninitialized
+    # loop carried values as we need their type information for tracing.
+    init_iter_args = get_state()
+    _assert_iteration_inputs(init_iter_args, symbol_names)
+
+    @qml.for_loop(start, stop, step)
+    def functional_for(i, *iter_args):
+        # Assign tracers to the iteration variables identified by AutoGraph (iter_args in mlir).
+        set_state(iter_args)
+
+        # The iteration index/element (for <...> in) is already handled by the body function, e.g.:
+        #   def body_fn(itr):
+        #     i, x = itr
+        #     ...
+        if enum_start is None and array_iterable is None:
+            # for i in range(..)
+            body_fn(i)
+        elif enum_start is None:
+            # for x in array
+            body_fn(array_iterable[i])
+        else:
+            # for (i, x) in enumerate(array)
+            body_fn((i + enum_start, array_iterable[i]))
+
+        return get_state()
+
+    final_iter_args = functional_for(*init_iter_args)
+    _assert_iteration_results(init_iter_args, final_iter_args, symbol_names)
+    return final_iter_args
+
+
+# pylint: disable=too-many-statements
+def for_stmt(
+    iteration_target: Any,
+    _extra_test: Union[Callable[[], bool], None],
+    body_fn: Callable[[int], None],
+    get_state: Callable[[], Tuple],
+    set_state: Callable[[Tuple], None],
+    symbol_names: Tuple[str],
+    _opts: dict,
+):
+    """An implementation of the AutoGraph 'for .. in ..' statement. The interface is defined by
+    AutoGraph, here we merely provide an implementation of it in terms of PennyLane primitives."""
+
+    assert _extra_test is None
+
+    # The general approach is to convert as much code as possible into a graph-based form:
+    # - For loops over iterables will attempt a conversion of the iterable to array
+    # - For loops over a Python range will be converted to a native PennyLane for_loop. The now
+    #   dynamic iteration variable can cause issues in downstream user code that raise an error.
+    # - For loops over a Python enumeration use a combination of the above, providing a dynamic
+    #   iteration variable and conversion of the iterable to array.
+
+    # Any of these could fail depending on the compatibility of the user code. A failure could
+    # also occur because an exception is raised during the tracing of the loop body after conversion
+    # (for example because the user forgot to use a list instead of an array)
+    # The PennyLane autograph implementation does not currently fall back to a Python loop in this case,
+    # but this has been implemented in Catalyst and could be extended to this. It does, however, require an
+    # active qeueing context.
+
+    exception_raised = None
+    init_state = get_state()
+    assert len(init_state) == len(symbol_names)
+
+    if isinstance(iteration_target, PRange):
+        start, stop, step = iteration_target.get_raw_range()
+        enum_start = None
+        iteration_array = None
+    elif isinstance(iteration_target, PEnumerate):
+        start, stop, step = 0, len(iteration_target.iteration_target), 1
+        enum_start = iteration_target.start_idx
+        try:
+            iteration_array = jnp.asarray(iteration_target.iteration_target)
+        except Exception as e:  # pylint: disable=bare-except, broad-exception-caught, broad-except
+            exception_raised = e
+    else:
+        start, stop, step = 0, len(iteration_target), 1
+        enum_start = None
+        try:
+            iteration_array = jnp.asarray(iteration_target)
+        except Exception as e:  # pylint: disable=bare-except, broad-exception-caught, broad-except
+            exception_raised = e
+
+    if exception_raised:
+
+        raise AutoGraphError(
+            f"Could not convert the iteration target {iteration_target} to array while processing "
+            f"a for-loop with AutoGraph."
+        ) from exception_raised
+
+    try:
+        set_state(init_state)
+        results = _call_pennylane_for(
+            start,
+            stop,
+            step,
+            body_fn,
+            get_state,
+            set_state,
+            symbol_names,
+            enum_start,
+            iteration_array,
+        )
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        # pylint: disable=import-outside-toplevel
+        import textwrap
+
+        raise AutoGraphError(
+            f"Tracing of an AutoGraph converted for loop failed with an exception:\n"
+            f"  {type(e).__name__}:{textwrap.indent(str(e), '    ')}\n"
+            f"\n"
+            f"Make sure that loop variables are not used in tracing-incompatible ways, for instance "
+            f"by indexing a Python list with it (rather than a JAX array). Also ensure all variables "
+            f"are initialized before the loop begins, and that they don't change type across iterations.\n"
+            f"To understand different types of JAX tracing errors, please refer to the guide at: "
+            f"https://jax.readthedocs.io/en/latest/errors.html"
+        ) from e
+
+    set_state(results)
+
+
 def _call_pennylane_while(loop_test, loop_body, get_state, set_state, symbol_names):
     """Dispatch to a PennyLane implementation of while loops."""
 
@@ -226,6 +378,7 @@ def converted_call(fn, args, kwargs, caller_fn_scope=None, options=None):
     with Patcher(
         (ag_api, "_TRANSPILER", qml.capture.autograph.transformer.TRANSFORMER),
         (ag_config, "CONVERSION_RULES", module_allowlist),
+        (ag_py_builtins, "BUILTIN_FUNCTIONS_MAP", py_builtins_map),
     ):
         # HOTFIX: pass through calls of known PennyLane wrapper functions
         if fn in (
@@ -264,3 +417,89 @@ def qnode_call_wrapper():
             return new_qnode()
 
         return ag_converted_call(fn, args, kwargs, caller_fn_scope, options)
+
+
+class PRange:
+    """PennyLane range object. This class re-implements the built-in range class
+    (which can't be inherited from). The only change is saving and accessing the
+    inputs directly, to circumvent some JAX-unfriendly code in the Python range.
+    """
+
+    def __init__(self, start_stop, stop=None, step=None):
+        self._py_range = None
+        self._start = start_stop if stop is not None else 0
+        self._stop = stop if stop is not None else start_stop
+        self._step = step if step is not None else 1
+
+    def get_raw_range(self):
+        """Get the raw values defining this range: start, stop, step."""
+        return self._start, self._stop, self._step
+
+    @property
+    def py_range(self):
+        """Access the underlying Python range object. If it doesn't exist, create one."""
+        if self._py_range is None:
+            self._py_range = range(self._start, self._stop, self._step)
+        return self._py_range
+
+    # Interface of the Python range class.
+    # pylint: disable=missing-function-docstring
+
+    @property
+    def start(self) -> int:  # pragma: no cover
+        return self.py_range.start
+
+    @property
+    def stop(self) -> int:  # pragma: no cover
+        return self.py_range.stop
+
+    @property
+    def step(self) -> int:  # pragma: no cover
+        return self.py_range.step
+
+    def count(self, __value: int) -> int:  # pragma: no cover
+        return self.py_range.count(__value)
+
+    def index(self, __value: int) -> int:  # pragma: no cover
+        return self.py_range.index(__value)
+
+    def __len__(self) -> int:  # pragma: no cover
+        return self.py_range.__len__()
+
+    def __eq__(self, __value: object) -> bool:  # pragma: no cover
+        return self.py_range.__eq__(__value)
+
+    def __hash__(self) -> int:  # pragma: no cover
+        return self.py_range.__hash__()
+
+    def __contains__(self, __key: object) -> bool:  # pragma: no cover
+        return self.py_range.__contains__(__key)
+
+    def __iter__(self) -> Iterator[int]:  # pragma: no cover
+        return self.py_range.__iter__()
+
+    def __getitem__(
+        self, __key: Union[SupportsIndex, slice]
+    ) -> Union[int, range]:  # pragma: no cover
+        return self.py_range.__getitem__(__key)
+
+    def __reversed__(self) -> Iterator[int]:  # pragma: no cover
+        return self.py_range.__reversed__()
+
+
+# pylint: disable=too-few-public-methods, super-init-not-called
+class PEnumerate(enumerate):
+    """PennyLane enumeration object. Inherits from Python ``enumerate``, but adds storing the
+    input iteration_target and start_idx, which are used by the for-loop conversion.
+    """
+
+    def __init__(self, iterable, start=0):
+        self.iteration_target = iterable
+        self.start_idx = start
+
+
+py_builtins_map = {
+    **ag_py_builtins.BUILTIN_FUNCTIONS_MAP,
+    "range": PRange,
+    "enumerate": PEnumerate,
+}

tests/capture/autograph/test_autograph.py

@@ -108,6 +108,7 @@ def test_get_extra_locals(self):
 
         assert ag_fn_dict["if_stmt"].__module__ == "pennylane.capture.autograph.ag_primitives"
         assert ag_fn_dict["while_stmt"].__module__ == "pennylane.capture.autograph.ag_primitives"
+        assert ag_fn_dict["for_stmt"].__module__ == "pennylane.capture.autograph.ag_primitives"
         assert (
             ag_fn_dict["converted_call"].__module__ == "pennylane.capture.autograph.ag_primitives"
         )