Integrate new autograph package into catalyst (#401)

To remove our dependency on tensorflow, we will be releasing a
standalone autograph package. This PR makes the necessary changes in
Catalyst to use the new package.

Changes include:
- Substitute tensorflow imports with standalone autograph imports.
- Inherit from the existing transformer class to reuse upstream code
more effectively. This means we now run all AST passes that standard
AutoGraph runs.
- Eliminate the ag_utils module.
- Make `diastatic-malt` a hard dependency of `pennylane-catalyst`.

[sc-46389]
[sc-57401]

.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: |

.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

@@ -53,6 +53,10 @@
 
 <h3>Improvements</h3>
 
+* 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)

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 (
@@ -82,7 +82,7 @@
     vjp,
     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

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,
+}