Add support to AutoGraph for

doc/changelog.md

@@ -71,7 +71,7 @@
   [(#725)](https://github.com/PennyLaneAI/catalyst/pull/725)
 
   Although library code is not meant to be targeted by Autograph conversion,
-  it sometimes make sense to enable it for specific submodules that might 
+  it sometimes make sense to enable it for specific submodules that might
   benefit from such conversion:
 
   ```py
@@ -81,6 +81,30 @@
 
   ```
 
+* Support for usage of single index JAX array operator update
+  inside Autograph annotated functions.
+  [(#769)](https://github.com/PennyLaneAI/catalyst/pull/769)
+
+  Using operator assignment syntax in favor of at...operation expressions is now possible for the following operations:
+  * `x[i] += y` in favor of `x.at[i].add(y)`
+  * `x[i] -= y` in favor of `x.at[i].add(-y)`
+  * `x[i] *= y` in favor of `x.at[i].multiply(y)`
+  * `x[i] /= y` in favor of `x.at[i].divide(y)`
+  * `x[i] **= y` in favor of `x.at[i].power(y)`
+
+  ```py
+  @qjit(autograph=True)
+  def f(x):
+    first_dim = x.shape[0]
+    result = jnp.copy(x)
+
+    for i in range(first_dim):
+      result[i] *= 2  # This is now supported
+
+    return result
+
+  ```
+
 <h3>Improvements</h3>
 
 * Catalyst now has support for `qml.sample(m)` where `m` is the result of a mid-circuit
@@ -103,7 +127,7 @@
   [(#751)](https://github.com/PennyLaneAI/catalyst/pull/751)
 
 * Refactored `vmap`,`qjit`, `mitigate_with_zne` and gradient decorators in order to follow
-  a unified pattern that uses a callable class implementing the decorator's logic. This 
+  a unified pattern that uses a callable class implementing the decorator's logic. This
   prevents having to excessively define functions in a nested fashion.
   [(#758)](https://github.com/PennyLaneAI/catalyst/pull/758)
   [(#761)](https://github.com/PennyLaneAI/catalyst/pull/761)
@@ -127,7 +151,7 @@
 
 * Correctly linking openblas routines necessary for `jax.scipy.linalg.expm`.
   In this bug fix, four openblas routines were newly linked and are now discoverable by `stablehlo.custom_call@<blas_routine>`. They are `blas_dtrsm`, `blas_ztrsm`, `lapack_dgetrf`, `lapack_zgetrf`.
-  [(#752)](https://github.com/PennyLaneAI/catalyst/pull/752)    
+  [(#752)](https://github.com/PennyLaneAI/catalyst/pull/752)
 
 <h3>Internal changes</h3>
 

doc/dev/autograph.rst

@@ -921,8 +921,8 @@ Notice that ``autograph=True`` must be set in order to process the
 ``autograph_include`` list. Otherwise an error will be reported.
 
 
-In-place JAX array assignments
-------------------------------
+In-place JAX array updates
+--------------------------
 
 To update array values when using JAX, the `JAX syntax for array assignment
 <https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#array-updates-x-at-idx-set-y>`__
@@ -952,3 +952,33 @@ standard Python array assignment syntax:
 ...     return result
 
 Under the hood, Catalyst converts anything coming in the latter notation into the former one.
+
+Similarly, to update array values with an operation when using JAX, the JAX syntax for array
+update (which uses the array `at` and the `add`, `multiply`, etc. methods) must be used:
+
+>>> @qjit(autograph=True)
+... def f(x):
+...     first_dim = x.shape[0]
+...     result = jnp.copy(x)
+...
+...     for i in range(first_dim):
+...         result = result.at[i].multiply(2)
+...
+...     return result
+
+Again, if updating a single index of the array, Autograph supports conversion of
+standard Python array operator assignment syntax for the equivalent in-place expressions
+listed in the `JAX documentation for jax.numpy.ndarray.at
+<https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html#jax.numpy.ndarray.at>`__:
+
+>>> @qjit(autograph=True)
+... def f(x):
+...     first_dim = x.shape[0]
+...     result = jnp.copy(x)
+...
+...     for i in range(first_dim):
+...         result[i] *= 2
+...
+...     return result
+
+Under the hood, Catalyst converts anything coming in the latter notation into the former one.

frontend/catalyst/autograph/ag_primitives.py

@@ -47,6 +47,11 @@
     "or_",
     "not_",
     "set_item",
+    "update_item_with_add",
+    "update_item_with_sub",
+    "update_item_with_mult",
+    "update_item_with_div",
+    "update_item_with_pow",
 ]
 
 
@@ -600,6 +605,131 @@ def set_item(target, i, x):
     return target
 
 
+def update_item_with_add(target, i, x):
+    """An implementation of the 'update_item_with_add' function from operator_update. The interface
+    is defined in operator_update.SingleIndexArrayOperatorUpdateTransformer, here we provide an
+    implementation in terms of Catalyst primitives. The idea is to accept the simpler single index
+    operator assignment syntax for Jax arrays, to subsequently transform it under the hood into the
+    set of 'at' and 'add' calls that Autograph supports. E.g.:
+        target[i] += x -> target = target.at[i].add(x)
+
+    .. note::
+        For this feature to work, 'converter.Feature.LISTS' had to be added to the
+        TOP_LEVEL_OPTIONS and NESTED_LEVEL_OPTIONS conversion options of our own Catalyst
+        Autograph transformer. If you create a new transformer and want to support this feature,
+        make sure you enable such option there as well.
+    """
+
+    # Apply the 'at...add' transformation only to Jax arrays.
+    # Otherwise, fallback to Python's default syntax.
+    if isinstance(target, DynamicJaxprTracer):
+        target = target.at[i].add(x)
+    else:
+        target[i] += x
+
+    return target
+
+
+def update_item_with_sub(target, i, x):
+    """An implementation of the 'update_item_with_sub' function from operator_update. The interface
+    is defined in operator_update.SingleIndexArrayOperatorUpdateTransformer, here we provide an
+    implementation in terms of Catalyst primitives. The idea is to accept the simpler single index
+    operator assignment syntax for Jax arrays, to subsequently transform it under the hood into the
+    set of 'at' and 'add' calls that Autograph supports. E.g.:
+        target[i] -= x -> target = target.at[i].add(-x)
+
+    .. note::
+        For this feature to work, 'converter.Feature.LISTS' had to be added to the
+        TOP_LEVEL_OPTIONS and NESTED_LEVEL_OPTIONS conversion options of our own Catalyst
+        Autograph transformer. If you create a new transformer and want to support this feature,
+        make sure you enable such option there as well.
+    """
+
+    # Apply the 'at...add' transformation only to Jax arrays.
+    # Otherwise, fallback to Python's default syntax.
+    if isinstance(target, DynamicJaxprTracer):
+        target = target.at[i].add(-x)
+    else:
+        target[i] -= x
+
+    return target
+
+
+def update_item_with_mult(target, i, x):
+    """An implementation of the 'update_item_with_mult' function from operator_update. The interface
+    is defined in operator_update.SingleIndexArrayOperatorUpdateTransformer, here we provide an
+    implementation in terms of Catalyst primitives. The idea is to accept the simpler single index
+    operator assignment syntax for Jax arrays, to subsequently transform it under the hood into the
+    set of 'at' and 'multiply' calls that Autograph supports. E.g.:
+        target[i] *= x -> target = target.at[i].multiply(x)
+
+    .. note::
+        For this feature to work, 'converter.Feature.LISTS' had to be added to the
+        TOP_LEVEL_OPTIONS and NESTED_LEVEL_OPTIONS conversion options of our own Catalyst
+        Autograph transformer. If you create a new transformer and want to support this feature,
+        make sure you enable such option there as well.
+    """
+
+    # Apply the 'at...multiply' transformation only to Jax arrays.
+    # Otherwise, fallback to Python's default syntax.
+    if isinstance(target, DynamicJaxprTracer):
+        target = target.at[i].multiply(x)
+    else:
+        target[i] *= x
+
+    return target
+
+
+def update_item_with_div(target, i, x):
+    """An implementation of the 'update_item_with_div' function from operator_update. The interface
+    is defined in operator_update.SingleIndexArrayOperatorUpdateTransformer, here we provide an
+    implementation in terms of Catalyst primitives. The idea is to accept the simpler single index
+    operator assignment syntax for Jax arrays, to subsequently transform it under the hood into the
+    set of 'at' and 'divide' calls that Autograph supports. E.g.:
+        target[i] /= x -> target = target.at[i].divide(x)
+
+    .. note::
+        For this feature to work, 'converter.Feature.LISTS' had to be added to the
+        TOP_LEVEL_OPTIONS and NESTED_LEVEL_OPTIONS conversion options of our own Catalyst
+        Autograph transformer. If you create a new transformer and want to support this feature,
+        make sure you enable such option there as well.
+    """
+
+    # Apply the 'at...divide' transformation only to Jax arrays.
+    # Otherwise, fallback to Python's default syntax.
+    if isinstance(target, DynamicJaxprTracer):
+        target = target.at[i].divide(x)
+    else:
+        target[i] /= x
+
+    return target
+
+
+def update_item_with_pow(target, i, x):
+    """An implementation of the 'update_item_with_pow' function from operator_update. The interface
+    is defined in operator_update.SingleIndexArrayOperatorUpdateTransformer, here we provide an
+    implementation in terms of Catalyst primitives. The idea is to accept the simpler single index
+    operator assignment syntax for Jax arrays, to subsequently transform it under the hood into the
+    set of 'at' and 'power' calls that Autograph supports. E.g.:
+        target[i] **= x -> target = target.at[i].power(x)
+
+    .. note::
+        For this feature to work, 'converter.Feature.LISTS' had to be added to the
+        TOP_LEVEL_OPTIONS and NESTED_LEVEL_OPTIONS conversion options of our own Catalyst
+        Autograph transformer. If you create a new transformer and want to support this feature,
+        make sure you enable such option there as well.
+    """
+
+    # Apply the 'at...power' transformation only to Jax arrays.
+    # Otherwise, fallback to Python's default syntax.
+    if isinstance(target, DynamicJaxprTracer):
+        target = target.at[i].power(x)
+    else:
+        target[i] **= x
+
+    return target
+
+
 class CRange:
     """Catalyst range object.
 

frontend/catalyst/autograph/operator_update.py

@@ -0,0 +1,74 @@
+# Copyright 2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Converter for array element operator assignment."""
+
+import gast
+from malt.core import converter
+from malt.pyct import templates
+
+
+# The methods from this class should be migrated to the SliceTransformer class in DiastaticMalt
+class SingleIndexArrayOperatorUpdateTransformer(converter.Base):
+    """Converts array element operator assignment statements into calls to update_item_with_{op},
+    where op is one of the following:
+
+    - `add` corresponding to `+=`
+    - `sub` to `-=`
+    - `mult` to `*=`
+    - `div` to `/=`
+    - `pow` to `**=`
+    """
+
+    def _process_single_update(self, target, op, value):
+        if not isinstance(target, gast.Subscript):
+            return None
+        s = target.slice
+        if isinstance(s, (gast.Tuple, gast.Slice)):
+            return None
+        if not isinstance(op, (gast.Mult, gast.Add, gast.Sub, gast.Div, gast.Pow)):
+            return None
+
+        template = f"""
+            target = ag__.update_item_with_{type(op).__name__.lower()}(target, i, x)
+        """
+
+        return templates.replace(template, target=target.value, i=target.slice, x=value)
+
+    def visit_AugAssign(self, node):
+        """The AugAssign node is replaced with a call to ag__.update_item_with_{op}
+        when its target is a single index array subscript and its op is an arithmetic
+        operator (i.e. Add, Sub, Mult, Div, or Pow), otherwise the node is left as is.
+
+        Example:
+            `x[i] += y` is replaced with `x = ag__.update_item_with(x, i, y)`
+            `x[i] ^= y` remains unchanged
+        """
+        node = self.generic_visit(node)
+        replacement = self._process_single_update(node.target, node.op, node.value)
+        if replacement is not None:
+            return replacement
+        return node
+
+
+def transform(node, ctx):
+    """Replace an AugAssign node with a call to ag__.update_item_with_{op}
+    when the its target is a single index array subscript and its op is an arithmetic
+    operator (i.e. Add, Sub, Mult, Div, or Pow), otherwise the node is left as is.
+
+    Example:
+        `x[i] += y` is replaced with `x = ag__.update_item_with(x, i, y)`
+        `x[i] ^= y` remains unchanged
+    """
+    return SingleIndexArrayOperatorUpdateTransformer(ctx).visit(node)

frontend/catalyst/autograph/transformer.py

@@ -29,7 +29,7 @@
 from malt.impl.api import PyToPy
 
 import catalyst
-from catalyst.autograph import ag_primitives
+from catalyst.autograph import ag_primitives, operator_update
 from catalyst.utils.exceptions import AutoGraphError
 
 
@@ -111,6 +111,21 @@ def get_cached_function(self, fn):
 
         return new_fn
 
+    def transform_ast(self, node, ctx):
+        """Overload of PyToPy.transform_ast from DiastaticMalt
+
+        .. note::
+            Once the operator_update interface has been migrated to the
+            DiastaticMalt project, this overload can be deleted."""
+        # The operator_update transform would be more correct if placed with
+        # slices.transform in PyToPy.transform_ast in DiastaticMalt rather than
+        # at the beginning of the transformation. operator_update.transform
+        # should come after the unsupported features check and intial analysis,
+        # but it fails if it does not come before variables.transform.
+        node = operator_update.transform(node, ctx)
+        node = super().transform_ast(node, ctx)
+        return node
+
 
 def run_autograph(fn):
     """Decorator that converts the given function into graph form."""