Merge branch 'master' into master

doc/code/qml_tape.rst

@@ -67,4 +67,5 @@ and a reduction in unintended side effects, ``QuantumScript`` is strictly used i
 .. automodapi:: pennylane.tape
     :no-main-docstr:
     :include-all-objects:
+    :skip: QuantumTapeBatch
     :inheritance-diagram:

doc/releases/changelog-dev.md

@@ -59,6 +59,51 @@
 * Molecules and Hamiltonians can now be constructed for all the elements present in the periodic table.
   [(#5821)](https://github.com/PennyLaneAI/pennylane/pull/5821)
 
+* `qml.for_loop` and `qml.while_loop` now fallback to standard Python control
+  flow if `@qjit` is not present, allowing the same code to work with and without
+  `@qjit` without any rewrites.
+  [(#6014)](https://github.com/PennyLaneAI/pennylane/pull/6014)
+
+  ```python
+  dev = qml.device("lightning.qubit", wires=3)
+
+  @qml.qnode(dev)
+  def circuit(x, n):
+
+      @qml.for_loop(0, n, 1)
+      def init_state(i):
+          qml.Hadamard(wires=i)
+
+      init_state()
+
+      @qml.for_loop(0, n, 1)
+      def apply_operations(i, x):
+          qml.RX(x, wires=i)
+
+          @qml.for_loop(i + 1, n, 1)
+          def inner(j):
+              qml.CRY(x**2, [i, j])
+
+          inner()
+          return jnp.sin(x)
+
+      apply_operations(x)
+      return qml.probs()
+  ```
+
+  ```pycon
+  >>> print(qml.draw(circuit)(0.5, 3))
+  0: ──H──RX(0.50)─╭●────────╭●──────────────────────────────────────┤  Probs
+  1: ──H───────────╰RY(0.25)─│──────────RX(0.48)─╭●──────────────────┤  Probs
+  2: ──H─────────────────────╰RY(0.25)───────────╰RY(0.23)──RX(0.46)─┤  Probs
+  >>> circuit(0.5, 3)
+  array([0.125     , 0.125     , 0.09949758, 0.15050242, 0.07594666,
+       0.11917543, 0.08942104, 0.21545687])
+  >>> qml.qjit(circuit)(0.5, 3)
+  Array([0.125     , 0.125     , 0.09949758, 0.15050242, 0.07594666,
+       0.11917543, 0.08942104, 0.21545687], dtype=float64)
+  ```
+
 * The `qubit_observable` function is modified to return an ascending wire order for molecular 
   Hamiltonians.
   [(#5950)](https://github.com/PennyLaneAI/pennylane/pull/5950)
@@ -175,6 +220,7 @@ Lillian M. A. Frederiksen,
 Pietropaolo Frisoni,
 Emiliano Godinez,
 Renke Huang,
+Josh Izaac,
 Soran Jahangiri,
 Christina Lee,
 Austin Huang,

pennylane/_qubit_device.py

@@ -26,7 +26,7 @@
 import logging
 import warnings
 from collections import defaultdict
-from typing import List, Union
+from typing import Union
 
 import numpy as np
 
@@ -1777,7 +1777,7 @@ def _adjoint_jacobian_processing(jac):
         # must be 2-dimensional
         return tuple(tuple(np.array(j_) for j_ in j) for j in jac)
 
-    def _get_diagonalizing_gates(self, circuit: QuantumTape) -> List[Operation]:
+    def _get_diagonalizing_gates(self, circuit: QuantumTape) -> list[Operation]:
         """Returns the gates that diagonalize the measured wires such that they
         are in the eigenbasis of the circuit observables.
 

pennylane/capture/primitives.py

@@ -15,8 +15,9 @@
 This submodule defines the abstract classes and primitives for capture.
 """
 
+from collections.abc import Callable
 from functools import lru_cache
-from typing import Callable, Optional, Tuple, Type
+from typing import Optional, Type
 
 import pennylane as qml
 
@@ -107,7 +108,7 @@ def __init__(
             self._n_wires = n_wires
             self.has_eigvals: bool = has_eigvals
 
-        def abstract_eval(self, num_device_wires: int, shots: int) -> Tuple[Tuple, type]:
+        def abstract_eval(self, num_device_wires: int, shots: int) -> tuple[tuple, type]:
             """Calculate the shape and dtype for an evaluation with specified number of device
             wires and shots.
 

pennylane/capture/switches.py

@@ -15,7 +15,7 @@
 Contains the switches to (de)activate the capturing mechanism, and a
 status reporting function on whether it is enabled or not.
 """
-from typing import Callable
+from collections.abc import Callable
 
 has_jax = True
 try:
@@ -24,7 +24,7 @@
     has_jax = False
 
 
-def _make_switches() -> [Callable[[], None], Callable[[], None], Callable[[], bool]]:
+def _make_switches() -> tuple[Callable[[], None], Callable[[], None], Callable[[], bool]]:
     r"""Create three functions, corresponding to an activation switch, a deactivation switch
     and a status query, in that order.
 

pennylane/compiler/qjit_api.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """QJIT compatible quantum and compilation operations API"""
+from collections.abc import Callable
 
 from .compiler import (
     AvailableCompilers,
@@ -377,20 +378,54 @@ def loop_rx(x):
         ops_loader = compilers[active_jit]["ops"].load()
         return ops_loader.while_loop(cond_fn)
 
-    raise CompileError("There is no active compiler package.")  # pragma: no cover
+    # if there is no active compiler, simply interpret the while loop
+    # via the Python interpretor.
+    def _decorator(body_fn: Callable) -> Callable:
+        """Transform that will call the input ``body_fn`` until the closure variable ``cond_fn`` is met.
 
+        Args:
+            body_fn (Callable):
 
-def for_loop(lower_bound, upper_bound, step):
-    """A :func:`~.qjit` compatible for-loop for PennyLane programs.
+        Closure Variables:
+            cond_fn (Callable):
 
-    .. note::
+        Returns:
+            Callable: a callable with the same signature as ``body_fn`` and ``cond_fn``.
+        """
+        return WhileLoopCallable(cond_fn, body_fn)
 
-        This function only supports the Catalyst compiler. See
-        :func:`catalyst.for_loop` for more details.
+    return _decorator
 
-        Please see the Catalyst :doc:`quickstart guide <catalyst:dev/quick_start>`,
-        as well as the :doc:`sharp bits and debugging tips <catalyst:dev/sharp_bits>`
-        page for an overview of the differences between Catalyst and PennyLane.
+
+class WhileLoopCallable:  # pylint:disable=too-few-public-methods
+    """Base class to represent a while loop. This class
+    when called with an initial state will execute the while
+    loop via the Python interpreter.
+
+    Args:
+        cond_fn (Callable): the condition function in the while loop
+        body_fn (Callable): the function that is executed within the while loop
+    """
+
+    def __init__(self, cond_fn, body_fn):
+        self.cond_fn = cond_fn
+        self.body_fn = body_fn
+
+    def __call__(self, *init_state):
+        args = init_state
+        fn_res = args if len(args) > 1 else args[0] if len(args) == 1 else None
+
+        while self.cond_fn(*args):
+            fn_res = self.body_fn(*args)
+            args = fn_res if len(args) > 1 else (fn_res,) if len(args) == 1 else ()
+
+        return fn_res
+
+
+def for_loop(lower_bound, upper_bound, step):
+    """A :func:`~.qjit` compatible for-loop for PennyLane programs. When
+    used without :func:`~.qjit`, this function will fall back to a standard
+    Python for loop.
 
     This decorator provides a functional version of the traditional
     for-loop, similar to `jax.cond.fori_loop <https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.fori_loop.html>`__.
@@ -430,19 +465,14 @@ def for_loop(lower_bound, upper_bound, step, loop_fn, *args):
         across iterations is handled automatically by the provided loop bounds, it must not be
         returned from the function.
 
-    Raises:
-        CompileError: if the compiler is not installed
-
     .. seealso:: :func:`~.while_loop`, :func:`~.qjit`
 
     **Example**
 
-
     .. code-block:: python
 
         dev = qml.device("lightning.qubit", wires=1)
 
-        @qml.qjit
         @qml.qnode(dev)
         def circuit(n: int, x: float):
 
@@ -457,15 +487,84 @@ def loop_rx(i, x):
             # apply the for loop
             final_x = loop_rx(x)
 
-            return qml.expval(qml.Z(0)), final_x
+            return qml.expval(qml.Z(0))
 
     >>> circuit(7, 1.6)
-    (array(0.97926626), array(0.55395718))
+    array(0.97926626)
+
+    ``for_loop`` is also :func:`~.qjit` compatible; when used with the
+    :func:`~.qjit` decorator, the for loop will not be unrolled, and instead
+    will be captured as-is during compilation and executed during runtime:
+
+    >>> qml.qjit(circuit)(7, 1.6)
+    Array(0.97926626, dtype=float64)
+
+    .. note::
+
+        Please see the Catalyst :doc:`quickstart guide <catalyst:dev/quick_start>`,
+        as well as the :doc:`sharp bits and debugging tips <catalyst:dev/sharp_bits>`
+        page for an overview of using quantum just-in-time compilation.
+
     """
 
     if active_jit := active_compiler():
         compilers = AvailableCompilers.names_entrypoints
         ops_loader = compilers[active_jit]["ops"].load()
         return ops_loader.for_loop(lower_bound, upper_bound, step)
 
-    raise CompileError("There is no active compiler package.")  # pragma: no cover
+    # if there is no active compiler, simply interpret the for loop
+    # via the Python interpretor.
+    def _decorator(body_fn):
+        """Transform that will call the input ``body_fn`` within a for loop defined by the closure variables lower_bound, upper_bound, and step.
+
+        Args:
+            body_fn (Callable): The function called within the for loop. Note that the loop body
+                function must always have the iteration index as its first
+                argument, which can be used arbitrarily inside the loop body. As the value of the index
+                across iterations is handled automatically by the provided loop bounds, it must not be
+                returned from the function.
+
+        Closure Variables:
+            lower_bound (int): starting value of the iteration index
+            upper_bound (int): (exclusive) upper bound of the iteration index
+            step (int): increment applied to the iteration index at the end of each iteration
+
+        Returns:
+            Callable: a callable with the same signature as ``body_fn``
+        """
+        return ForLoopCallable(lower_bound, upper_bound, step, body_fn)
+
+    return _decorator
+
+
+class ForLoopCallable:  # pylint:disable=too-few-public-methods
+    """Base class to represent a for loop. This class
+    when called with an initial state will execute the while
+    loop via the Python interpreter.
+
+    Args:
+        lower_bound (int): starting value of the iteration index
+        upper_bound (int): (exclusive) upper bound of the iteration index
+        step (int): increment applied to the iteration index at the end of each iteration
+        body_fn (Callable): The function called within the for loop. Note that the loop body
+            function must always have the iteration index as its first
+            argument, which can be used arbitrarily inside the loop body. As the value of the index
+            across iterations is handled automatically by the provided loop bounds, it must not be
+            returned from the function.
+    """
+
+    def __init__(self, lower_bound, upper_bound, step, body_fn):
+        self.lower_bound = lower_bound
+        self.upper_bound = upper_bound
+        self.step = step
+        self.body_fn = body_fn
+
+    def __call__(self, *init_state):
+        args = init_state
+        fn_res = args if len(args) > 1 else args[0] if len(args) == 1 else None
+
+        for i in range(self.lower_bound, self.upper_bound, self.step):
+            fn_res = self.body_fn(i, *args)
+            args = fn_res if len(args) > 1 else (fn_res,) if len(args) == 1 else ()
+
+        return fn_res

pennylane/data/attributes/dictionary.py

@@ -15,9 +15,8 @@
 of Dataset attributes."""
 
 
-import typing
-from collections.abc import Mapping
-from typing import Dict, Generic, Union
+from collections.abc import Iterator, Mapping, MutableMapping
+from typing import Generic, Union
 
 from pennylane.data.base.attribute import DatasetAttribute
 from pennylane.data.base.hdf5 import HDF5Any, HDF5Group
@@ -27,35 +26,35 @@
 
 class DatasetDict(  # pylint: disable=too-many-ancestors
     Generic[T],
-    DatasetAttribute[HDF5Group, typing.Mapping[str, T], typing.Mapping[str, T]],
-    typing.MutableMapping[str, T],
+    DatasetAttribute[HDF5Group, Mapping[str, T], Mapping[str, T]],
+    MutableMapping[str, T],
     MapperMixin,
 ):
     """Provides a dict-like collection for Dataset attribute types. Keys must
     be strings."""
 
     type_id = "dict"
 
-    def __post_init__(self, value: typing.Mapping[str, T]):
+    def __post_init__(self, value: Mapping[str, T]):
         super().__post_init__(value)
         self.update(value)
 
     @classmethod
-    def default_value(cls) -> Dict:
+    def default_value(cls) -> dict:
         return {}
 
-    def hdf5_to_value(self, bind: HDF5Group) -> typing.MutableMapping[str, T]:
+    def hdf5_to_value(self, bind: HDF5Group) -> MutableMapping[str, T]:
         return self
 
     def value_to_hdf5(self, bind_parent: HDF5Group, key: str, value: None) -> HDF5Group:
         grp = bind_parent.create_group(key)
 
         return grp
 
-    def copy_value(self) -> Dict[str, T]:
+    def copy_value(self) -> dict[str, T]:
         return {key: attr.copy_value() for key, attr in self._mapper.items()}
 
-    def copy(self) -> Dict[str, T]:
+    def copy(self) -> dict[str, T]:
         """Returns a copy of this mapping as a builtin ``dict``, with all
         elements copied."""
         return self.copy_value()
@@ -93,7 +92,7 @@ def __eq__(self, __value: object) -> bool:
 
         return all(__value[key] == self[key] for key in __value.keys())
 
-    def __iter__(self) -> typing.Iterator[str]:
+    def __iter__(self) -> Iterator[str]:
         return (key for key in self.bind.keys())
 
     def __str__(self) -> str: