Merge remote-tracking branch 'origin/main' into native-quantum-contro…

…l-toml

.dep-versions

@@ -3,4 +3,4 @@ jax=0.4.23
 mhlo=4611968a5f6818e6bdfb82217b9e836e0400bba9
 llvm=cd9a641613eddf25d4b25eaa96b2c393d401d42c
 enzyme=1beb98b51442d50652eaa3ffb9574f4720d611f1
-pennylane=95129a0d6365b48cb4acfa828ceb6a8532e47ef5
+pennylane=f638bc53a09724a83d9d58964bf37bfd438e6aa3

.github/workflows/build-wheel-linux-x86_64.yaml

@@ -375,7 +375,6 @@ jobs:
     - name: Install Python dependencies
       run: |
         python${{ matrix.python_version }} -m pip install pytest pytest-xdist
-        python${{ matrix.python_version }} -m pip install tensorflow-cpu  # for autograph tests
 
     - name: Install Catalyst
       run: |

.github/workflows/build-wheel-macos-arm64.yaml

@@ -383,9 +383,7 @@ jobs:
 
     - name: Install Python dependencies
       run: |
-        # tensorflow-cpu is not distributed for macOS ARM
         python${{ matrix.python_version }} -m pip install pytest pytest-xdist
-        python${{ matrix.python_version }} -m pip install tensorflow  # for autograph tests
 
     - name: Install Catalyst
       run: |

.github/workflows/build-wheel-macos-x86_64.yaml

@@ -354,7 +354,6 @@ jobs:
     - name: Install Python dependencies
       run: |
         python${{ matrix.python_version }} -m pip install pytest pytest-xdist
-        python${{ matrix.python_version }} -m pip install tensorflow-cpu  # for autograph tests
 
     - name: Install Catalyst
       run: |

.github/workflows/check-pl-compat.yaml

@@ -157,7 +157,7 @@ jobs:
     - name: Install PennyLane (release-candidate)
       if: ${{ inputs.pennylane == 'release-candidate' }}
       run: |
-        pip install git+https://github.com/PennyLaneAI/pennylane.git@v0.34.0-rc0
+        pip install git+https://github.com/PennyLaneAI/pennylane.git@v0.35.0-rc0
 
 
     - name: Add Frontend Dependencies to PATH

.github/workflows/constants.yaml

@@ -70,7 +70,7 @@ jobs:
       - name: Python versions
         id: python_versions
         run: |
-          echo 'python_versions=["3.9", "3.10", "3.11"]' >> $GITHUB_OUTPUT
+          echo 'python_versions=["3.9", "3.10", "3.11", "3.12"]' >> $GITHUB_OUTPUT
 
       - name: Primary Python version
         id: primary_python_version

.github/workflows/utils/setup_self_hosted_macos_env/action.yml

@@ -97,7 +97,7 @@ runs:
       id: site_packages
       shell: bash
       run: |
-        echo "site_packages_dir=$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')" >> $GITHUB_OUTPUT
+        echo "site_packages_dir=$(python -c 'import sysconfig; print(sysconfig.get_path("platlib"))')" >> $GITHUB_OUTPUT
 
     - name: Display Python Setup
       shell: bash

.pylintrc

@@ -10,7 +10,7 @@ extension-pkg-allow-list=catalyst.utils.wrapper
 # (useful for modules/projects where namespaces are manipulated during runtime
 # and thus existing member attributes cannot be deduced by static analysis. It
 # supports qualified module names, as well as Unix pattern matching.
-ignored-modules=pennylane.ops,jaxlib.mlir.ir,jaxlib.xla_extension,tensorflow
+ignored-modules=pennylane.ops,jaxlib.mlir.ir,jaxlib.xla_extension
 
 # List of classes names for which member attributes should not be checked
 # (useful for classes with attributes dynamically set). This supports can work

doc/changelog.md

@@ -1,9 +1,27 @@
-# Release 0.5.0-dev
+# Release 0.6.0-dev
 
 <h3>New features</h3>
 
+<h3>Improvements</h3>
+
+<h3>Breaking changes</h3>
+
+<h3>Bug fixes</h3>
+
+<h3>Contributors</h3>
+
+This release contains contributions from (in alphabetical order):
+
+# Release 0.5.0
+
+<h3>New features</h3>
+
+* Catalyst now supports Python 3.12
+  [(#532)](https://github.com/PennyLaneAI/catalyst/pull/532)
+
 * Catalyst now supports just-in-time compilation of static arguments.
   [(#476)](https://github.com/PennyLaneAI/catalyst/pull/476)
+  [(#550)](https://github.com/PennyLaneAI/catalyst/pull/550)
 
   The ``@qjit`` decorator can now be used to compile functions with static arguments with
   the ``static_argnums`` keyword argument. ``static_argnums`` can be an integer or an iterable
@@ -43,14 +61,52 @@
 * Catalyst now supports executing tapes in CUDA-Quantum simulators.
   [(#477)](https://github.com/PennyLaneAI/catalyst/pull/477)
   [(#536)](https://github.com/PennyLaneAI/catalyst/pull/536)
+  [(#547)](https://github.com/PennyLaneAI/catalyst/pull/547)
 
   It has added the following devices:
   * softwareq.qpp
   * nvidia.statevec (with support for multi-gpu)
   * nvidia.tensornet (with support for matrix product state)
 
+* Catalyst now supports QJIT compatible `catalyst.vmap` of hybrid programs.
+  `catalyst.vmap` offers the vectorization mapping backed by `catalyst.for_loop`.
+  [(#497)](https://github.com/PennyLaneAI/catalyst/pull/497)
+
+  For example,
+
+  ```py
+  @qjit
+  def workflow(x, y):
+      @qml.qnode(qml.device('lightning.qubit, wires=1))
+      def circuit(x, y):
+          qml.RX(jnp.pi * x['a'][0], wires=0)
+          qml.RY(x['a'][1] ** 2, wires=0)
+          qml.RX(x['b'][1] * x['b'][2] + y, wires=0)
+          return qml.state(), qml.probs(0)
+
+      return vmap(circuit, in_axes=(1, None), axis_size=5)(x, y)
+  ```
+
 <h3>Improvements</h3>
 
+* Catalyst now supports the differentiation of sliced arrays.
+  [(#552)](https://github.com/PennyLaneAI/catalyst/pull/552)
+
+  ```py
+  def f(x):
+    return jax.numpy.sum(x[::2])
+
+  x = jax.numpy.array([0.1, 0.2, 0.3, 0.4])
+  ```
+  ```pycon
+  >>> catalyst.qjit(catalyst.grad(f))(x)
+  [1. 0. 1. 0.]
+  ```
+
+* Catalyst no longer relies on a TensorFlow installation for its AutoGraph functionality. Instead,
+  the standalone `diastatic-malt` package is used and automatically installed as a dependency.
+  [(#401)](https://github.com/PennyLaneAI/catalyst/pull/401)
+
 * Catalyst will now remember previously compiled functions when the PyTree metadata of arguments
   changes, in addition to already rememebering compiled functions when static arguments change.
   [(#522)](https://github.com/PennyLaneAI/catalyst/pull/531)
@@ -89,10 +145,11 @@
   compiling
   ```
 
-* Keep the structure of the function return when taking the derivatives, JVP and VJP (pytrees support).
+* Keep the structure of the function arguments and return when taking the derivatives, JVP and VJP (pytrees support).
   [(#500)](https://github.com/PennyLaneAI/catalyst/pull/500)
   [(#501)](https://github.com/PennyLaneAI/catalyst/pull/501)
   [(#508)](https://github.com/PennyLaneAI/catalyst/pull/508)
+  [(#549)](https://github.com/PennyLaneAI/catalyst/pull/549)
 
   ```py
   dev = qml.device("lightning.qubit", wires=1)
@@ -181,6 +238,9 @@
 
 <h3>Breaking changes</h3>
 
+* The JAX version used by Catalyst has been updated to `v0.4.23`.
+  [(#428)](https://github.com/PennyLaneAI/catalyst/pull/428)
+
 * The Catalyst Python frontend has been partially refactored. The impact on user-facing
   functionality is minimal, but the location of certain classes and methods used by the package
   may have changed.
@@ -303,6 +363,12 @@
 
 <h3>Bug fixes</h3>
 
+* `QCtrl` now implements `map_wires`.
+  [(#555)](https://github.com/PennyLaneAI/catalyst/pull/555)
+
+  Applying `map_wires` to `QCtrl` lead to an incorrect circuit.
+  The wires in the nested regions remained unchanged.
+
 * Catalyst will no longer print a warning that recompilation is triggered when a `@qjit` decorated
   function with no arguments is invoke without having been compiled first, for example via the use
   of `target="mlir"`.

doc/dev/autograph.rst

@@ -29,19 +29,12 @@ restrictions and constraints you may discover.
 Using AutoGraph
 ---------------
 
-AutoGraph currently requires TensorFlow as a dependency; in most cases it can
-be installed via
+The AutoGraph feature in Catalyst is supported by the ``diastatic-malt`` package, a standalone
+fork of the AutoGraph module in TensorFlow (
+`official documentation <https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/index.md>`_
+).
 
-.. code-block:: console
-
-    pip install tensorflow
-
-but please refer to the
-`TensorFlow documentation <https://www.tensorflow.org/install>`__
-for specific details on installing TensorFlow for your platform.
-
-Once TensorFlow is available, AutoGraph can be enabled by passing
-``autograph=True`` to the ``@qjit`` decorator:
+To enable AutoGraph in Catalyst, simply pass ``autograph=True`` to the ``@qjit`` decorator:
 
 .. code-block:: python
 
@@ -661,7 +654,7 @@ of multiple measurements. For example,
         qml.RY(0.5, wires=1)
 
         m1 = measure(0)
-        m2 = measure(1)  
+        m2 = measure(1)
 
         if m1 and not m2:
             qml.Hadamard(wires=1)
@@ -713,17 +706,17 @@ Array arguments
 
 Note that, like with NumPy and JAX, logical operators apply elementwise to array arguments:
 
->>> @qjit(autograph=True)  
-... def f(x, y):  
+>>> @qjit(autograph=True)
+... def f(x, y):
 ...     return x and y
 >>> f(jnp.array([0, 1]), jnp.array([1, 1]))
 array([False,  True])
 
 Care must therefore be taken when using logical operators within conditional branches;
 ``jnp.all`` and ``jnp.any`` can be used to generate a single boolean for conditionals:
 
->>> @qjit(autograph=True)  
-... def f(x, y):  
+>>> @qjit(autograph=True)
+... def f(x, y):
 ...     if jnp.all(x and y):
 ...         z = 1
 ...     else:

frontend/catalyst/__init__.py

@@ -66,7 +66,7 @@
 )
 
 from catalyst import debug
-from catalyst.ag_utils import AutoGraphError, autograph_source
+from catalyst.autograph import autograph_source
 from catalyst.compiler import CompileOptions
 from catalyst.jit import QJIT, qjit
 from catalyst.pennylane_extensions import (
@@ -80,9 +80,10 @@
     measure,
     mitigate_with_zne,
     vjp,
+    vmap,
     while_loop,
 )
-from catalyst.utils.exceptions import CompileError
+from catalyst.utils.exceptions import AutoGraphError, CompileError
 
 autograph_ignore_fallbacks = False
 """bool: Specify whether AutoGraph should avoid raising
@@ -177,6 +178,7 @@
     "vjp",
     "jvp",
     "adjoint",
+    "vmap",
     "mitigate_with_zne",
     "debug",
     "autograph_source",

frontend/catalyst/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.5.0-dev"
+__version__ = "0.6.0-dev"

frontend/catalyst/ag_primitives.py

@@ -21,41 +21,22 @@
 
 import jax
 import jax.numpy as jnp
-
-# Use tensorflow implementations for handling function scopes and calls,
-# as well as various utility objects.
 import pennylane as qml
-import tensorflow.python.autograph.impl.api as tf_autograph_api
+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
+from malt.pyct.origin_info import LineLocation
 from pennylane.queuing import AnnotatedQueue
-from tensorflow.python.autograph.core import config
-from tensorflow.python.autograph.core.converter import STANDARD_OPTIONS as STD
-from tensorflow.python.autograph.core.converter import ConversionOptions
-from tensorflow.python.autograph.core.function_wrappers import (
-    FunctionScope,
-    with_function_scope,
-)
-from tensorflow.python.autograph.impl.api import autograph_artifact
-from tensorflow.python.autograph.impl.api import converted_call as tf_converted_call
-from tensorflow.python.autograph.operators.variables import (
-    Undefined,
-    UndefinedReturnValue,
-)
-from tensorflow.python.autograph.pyct.origin_info import LineLocation
 
 import catalyst
-from catalyst.ag_utils import AutoGraphError
 from catalyst.jax_extras import DynamicJaxprTracer, ShapedArray
 from catalyst.tracing.contexts import EvaluationContext
+from catalyst.utils.exceptions import AutoGraphError
 from catalyst.utils.patching import Patcher
 
 __all__ = [
-    "STD",
-    "ConversionOptions",
-    "Undefined",
-    "UndefinedReturnValue",
-    "autograph_artifact",
-    "FunctionScope",
-    "with_function_scope",
     "if_stmt",
     "for_stmt",
     "while_stmt",
@@ -522,19 +503,22 @@ def get_source_code_info(tb_frame):
 # issues such as always tracing through code that should only be executed conditionally. We might
 # have to be even more restrictive in the future to prevent issues if necessary.
 module_allowlist = (
-    config.DoNotConvert("pennylane"),
-    config.DoNotConvert("catalyst"),
-    config.DoNotConvert("jax"),
-) + config.CONVERSION_RULES
+    ag_config.DoNotConvert("pennylane"),
+    ag_config.DoNotConvert("catalyst"),
+    ag_config.DoNotConvert("jax"),
+    *ag_config.CONVERSION_RULES,
+)
 
 
 def converted_call(fn, args, kwargs, caller_fn_scope=None, options=None):
     """We want AutoGraph to use our own instance of the AST transformer when recursively
     transforming functions, but otherwise duplicate the same behaviour."""
 
+    # TODO: eliminate the need for patching by improving the autograph interface
     with Patcher(
-        (tf_autograph_api, "_TRANSPILER", catalyst.autograph.TRANSFORMER),
-        (config, "CONVERSION_RULES", module_allowlist),
+        (ag_api, "_TRANSPILER", catalyst.autograph.TRANSFORMER),
+        (ag_config, "CONVERSION_RULES", module_allowlist),
+        (ag_py_builtins, "BUILTIN_FUNCTIONS_MAP", py_builtins_map),
     ):
         # HOTFIX: pass through calls of known Catalyst wrapper functions
         if fn in (
@@ -557,14 +541,7 @@ def passthrough_wrapper(*args, **kwargs):
                 **(kwargs if kwargs is not None else {}),
             )
 
-        # Dispatch range calls to a custom range class that enables constructs like
-        # `for .. in range(..)` to be converted natively to `for_loop` calls. This is beneficial
-        # since the Python range function does not allow tracers as arguments.
-        if fn is range:
-            return CRange(*args, **(kwargs if kwargs is not None else {}))
-        elif fn is enumerate:
-            return CEnumerate(*args, **(kwargs if kwargs is not None else {}))
-
+        # TODO: find a way to handle custom decorators more effectively with autograph
         # We need to unpack nested QNode and QJIT calls as autograph will have trouble handling
         # them. Ideally, we only want the wrapped function to be transformed by autograph, rather
         # than the QNode or QJIT call method.
@@ -580,12 +557,12 @@ def passthrough_wrapper(*args, **kwargs):
 
             @functools.wraps(fn.func)
             def qnode_call_wrapper():
-                return tf_converted_call(fn.func, args, kwargs, caller_fn_scope, options)
+                return ag_converted_call(fn.func, args, kwargs, caller_fn_scope, options)
 
             new_qnode = qml.QNode(qnode_call_wrapper, device=fn.device, diff_method=fn.diff_method)
             return new_qnode()
 
-        return tf_converted_call(fn, args, kwargs, caller_fn_scope, options)
+        return ag_converted_call(fn, args, kwargs, caller_fn_scope, options)
 
 
 class CRange:
@@ -674,3 +651,10 @@ class CEnumerate(enumerate):
     def __init__(self, iterable, start=0):
         self.iteration_target = iterable
         self.start_idx = start
+
+
+py_builtins_map = {
+    **ag_py_builtins.BUILTIN_FUNCTIONS_MAP,
+    "range": CRange,
+    "enumerate": CEnumerate,
+}