Repository: GridTools/gt4py
Pull request #1826 (2025-01-27 13:11:01)

refactor[next]: fixed point transformation pass infrastructure
Executing transformations until a fixed point is reached, i.e., not transformation is applicable anymore is a common pattern encountered in many passes. The `CollapseTuple` pass already contained some infrastructure / utilities for this purpose. Since we now want to use the same approach in `ConstantFolding` and `FieldOpFusion`, they are extracted into a common base class here.

This is essentially a step into the direction of a more general pass manager that allows composition of transformations. It is very hard to design something on the scratch board that fulfills all needs we have in the various transformations. The idea is to extend the capabilities of this class step-by-step to cover more and more transformation use-cases such that more passes can use this infrastructure until we can eventually evaluate how dissimilar passes can be composed, e.g. CollapseTuple and ConstantFolding.

Edit @tehrengruber


Patch:
diff --git a/.github/workflows/deploy-release.yml b/.github/workflows/deploy-release.yml
index b519e008..7a7505ca 100644
--- a/.github/workflows/deploy-release.yml
+++ b/.github/workflows/deploy-release.yml
@@ -42,7 +42,7 @@ jobs:
       id-token: write
     steps:
     - name: Download wheel
-      uses: actions/download-artifact@v3
+      uses: actions/download-artifact@v4
       with:
         name: gt4py-dist
         path: dist
@@ -60,7 +60,7 @@ jobs:
       id-token: write
     steps:
     - name: Download wheel
-      uses: actions/download-artifact@v3
+      uses: actions/download-artifact@v4
       with:
         name: gt4py-dist
         path: dist
diff --git a/src/gt4py/next/iterator/transforms/collapse_tuple.py b/src/gt4py/next/iterator/transforms/collapse_tuple.py
index 0a0cf6d3..6db58f37 100644
--- a/src/gt4py/next/iterator/transforms/collapse_tuple.py
+++ b/src/gt4py/next/iterator/transforms/collapse_tuple.py
@@ -23,6 +23,7 @@ from gt4py.next.iterator.ir_utils import (
     ir_makers as im,
     misc as ir_misc,
 )
+from gt4py.next.iterator.transforms import fixed_point_transformation
 from gt4py.next.iterator.transforms.inline_lambdas import InlineLambdas, inline_lambda
 from gt4py.next.iterator.type_system import inference as itir_type_inference
 from gt4py.next.type_system import type_info, type_specifications as ts
@@ -86,8 +87,10 @@ def _is_trivial_or_tuple_thereof_expr(node: ir.Node) -> bool:
 #  go through all available transformation and apply them. However the final result here still
 #  reads a little convoluted and is also different to how we write other transformations. We
 #  should revisit the pattern here and try to find a more general mechanism.
-@dataclasses.dataclass(frozen=True)
-class CollapseTuple(eve.PreserveLocationVisitor, eve.NodeTranslator):
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class CollapseTuple(
+    fixed_point_transformation.FixedPointTransformation, eve.PreserveLocationVisitor
+):
     """
     Simplifies `make_tuple`, `tuple_get` calls.
 
@@ -98,7 +101,7 @@ class CollapseTuple(eve.PreserveLocationVisitor, eve.NodeTranslator):
     # TODO(tehrengruber): This Flag mechanism is a little low level. What we actually want
     #   is something like a pass manager, where for each pattern we have a corresponding
     #   transformation, etc.
-    class Flag(enum.Flag):
+    class Transformation(enum.Flag):
         #: `make_tuple(tuple_get(0, t), tuple_get(1, t), ..., tuple_get(N-1,t))` -> `t`
         COLLAPSE_MAKE_TUPLE_TUPLE_GET = enum.auto()
         #: `tuple_get(i, make_tuple(e_0, e_1, ..., e_i, ..., e_N))` -> `e_i`
@@ -137,12 +140,12 @@ class CollapseTuple(eve.PreserveLocationVisitor, eve.NodeTranslator):
         INLINE_TRIVIAL_LET = enum.auto()
 
         @classmethod
-        def all(self) -> CollapseTuple.Flag:
+        def all(self) -> CollapseTuple.Transformation:
             return functools.reduce(operator.or_, self.__members__.values())
 
     uids: eve_utils.UIDGenerator
     ignore_tuple_size: bool
-    flags: Flag = Flag.all()  # noqa: RUF009 [function-call-in-dataclass-default-argument]
+    enabled_transformations: Transformation = Transformation.all()  # noqa: RUF009 [function-call-in-dataclass-default-argument]
 
     PRESERVED_ANNEX_ATTRS = ("type",)
 
@@ -155,8 +158,8 @@ class CollapseTuple(eve.PreserveLocationVisitor, eve.NodeTranslator):
         remove_letified_make_tuple_elements: bool = True,
         offset_provider_type: Optional[common.OffsetProviderType] = None,
         within_stencil: Optional[bool] = None,
-        # manually passing flags is mostly for allowing separate testing of the modes
-        flags: Optional[Flag] = None,
+        # manually passing enabled transformations is mostly for allowing separate testing of the modes
+        enabled_transformations: Optional[Transformation] = None,
         # allow sym references without a symbol declaration, mostly for testing
         allow_undeclared_symbols: bool = False,
         uids: Optional[eve_utils.UIDGenerator] = None,
@@ -174,7 +177,7 @@ class CollapseTuple(eve.PreserveLocationVisitor, eve.NodeTranslator):
                 to remove left-overs from `LETIFY_MAKE_TUPLE_ELEMENTS` transformation.
                 `(λ(_tuple_el_1, _tuple_el_2) → {_tuple_el_1, _tuple_el_2})(1, 2)` -> {1, 2}`
         """
-        flags = flags or cls.flags
+        enabled_transformations = enabled_transformations or cls.enabled_transformations
         offset_provider_type = offset_provider_type or {}
         uids = uids or eve_utils.UIDGenerator()
 
@@ -194,7 +197,7 @@ class CollapseTuple(eve.PreserveLocationVisitor, eve.NodeTranslator):
 
         new_node = cls(
             ignore_tuple_size=ignore_tuple_size,
-            flags=flags,
+            enabled_transformations=enabled_transformations,
             uids=uids,
         ).visit(node, within_stencil=within_stencil)
 
@@ -210,45 +213,17 @@ class CollapseTuple(eve.PreserveLocationVisitor, eve.NodeTranslator):
 
         return new_node
 
-    def visit_FunCall(self, node: ir.FunCall, **kwargs) -> ir.Node:
+    def visit(self, node, **kwargs):
         if cpm.is_call_to(node, "as_fieldop"):
             kwargs = {**kwargs, "within_stencil": True}
 
-        node = self.generic_visit(node, **kwargs)
-        return self.fp_transform(node, **kwargs)
-
-    def fp_transform(self, node: ir.Node, **kwargs) -> ir.Node:
-        while True:
-            new_node = self.transform(node, **kwargs)
-            if new_node is None:
-                break
-            assert new_node != node
-            node = new_node
-        return node
-
-    def transform(self, node: ir.Node, **kwargs) -> Optional[ir.Node]:
-        if not isinstance(node, ir.FunCall):
-            return None
-
-        for transformation in self.Flag:
-            if self.flags & transformation:
-                assert isinstance(transformation.name, str)
-                method = getattr(self, f"transform_{transformation.name.lower()}")
-                result = method(node, **kwargs)
-                if result is not None:
-                    assert (
-                        result is not node
-                    )  # transformation should have returned None, since nothing changed
-                    itir_type_inference.reinfer(result)
-                    return result
-        return None
+        return super().visit(node, **kwargs)
 
     def transform_collapse_make_tuple_tuple_get(
         self, node: ir.FunCall, **kwargs
     ) -> Optional[ir.Node]:
-        if node.fun == ir.SymRef(id="make_tuple") and all(
-            isinstance(arg, ir.FunCall) and arg.fun == ir.SymRef(id="tuple_get")
-            for arg in node.args
+        if cpm.is_call_to(node, "make_tuple") and all(
+            cpm.is_call_to(arg, "tuple_get") for arg in node.args
         ):
             # `make_tuple(tuple_get(0, t), tuple_get(1, t), ..., tuple_get(N-1,t))` -> `t`
             assert isinstance(node.args[0], ir.FunCall)
@@ -275,10 +250,9 @@ class CollapseTuple(eve.PreserveLocationVisitor, eve.NodeTranslator):
         self, node: ir.FunCall, **kwargs
     ) -> Optional[ir.Node]:
         if (
-            node.fun == ir.SymRef(id="tuple_get")
-            and isinstance(node.args[1], ir.FunCall)
-            and node.args[1].fun == ir.SymRef(id="make_tuple")
+            cpm.is_call_to(node, "tuple_get")
             and isinstance(node.args[0], ir.Literal)
+            and cpm.is_call_to(node.args[1], "make_tuple")
         ):
             # `tuple_get(i, make_tuple(e_0, e_1, ..., e_i, ..., e_N))` -> `e_i`
             assert type_info.is_integer(node.args[0].type)
@@ -291,7 +265,7 @@ class CollapseTuple(eve.PreserveLocationVisitor, eve.NodeTranslator):
         return None
 
     def transform_propagate_tuple_get(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
-        if node.fun == ir.SymRef(id="tuple_get") and isinstance(node.args[0], ir.Literal):
+        if cpm.is_call_to(node, "tuple_get") and isinstance(node.args[0], ir.Literal):
             # TODO(tehrengruber): extend to general symbols as long as the tail call in the let
             #   does not capture
             # `tuple_get(i, let(...)(make_tuple()))` -> `let(...)(tuple_get(i, make_tuple()))`
@@ -314,8 +288,8 @@ class CollapseTuple(eve.PreserveLocationVisitor, eve.NodeTranslator):
                 )
         return None
 
-    def transform_letify_make_tuple_elements(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
-        if node.fun == ir.SymRef(id="make_tuple"):
+    def transform_letify_make_tuple_elements(self, node: ir.Node, **kwargs) -> Optional[ir.Node]:
+        if cpm.is_call_to(node, "make_tuple"):
             # `make_tuple(expr1, expr1)`
             # -> `let((_tuple_el_1, expr1), (_tuple_el_2, expr2))(make_tuple(_tuple_el_1, _tuple_el_2))`
             bound_vars: dict[ir.Sym, ir.Expr] = {}
@@ -334,7 +308,7 @@ class CollapseTuple(eve.PreserveLocationVisitor, eve.NodeTranslator):
                 )
         return None
 
-    def transform_inline_trivial_make_tuple(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
+    def transform_inline_trivial_make_tuple(self, node: ir.Node, **kwargs) -> Optional[ir.Node]:
         if cpm.is_let(node):
             # `let(tup, make_tuple(trivial_expr1, trivial_expr2))(foo(tup))`
             #  -> `foo(make_tuple(trivial_expr1, trivial_expr2))`
@@ -349,7 +323,7 @@ class CollapseTuple(eve.PreserveLocationVisitor, eve.NodeTranslator):
             #  in local-view for now. Revisit.
             return None
 
-        if not cpm.is_call_to(node, "if_"):
+        if isinstance(node, ir.FunCall) and not cpm.is_call_to(node, "if_"):
             # TODO(tehrengruber): Only inline if type of branch value is a tuple.
             # Examples:
             # `(if cond then {1, 2} else {3, 4})[0]` -> `if cond then {1, 2}[0] else {3, 4}[0]`
@@ -391,7 +365,7 @@ class CollapseTuple(eve.PreserveLocationVisitor, eve.NodeTranslator):
         # `if True then {2, 1} else {4, 3}`. The examples in the comments below all refer to this
         # tuple reordering example here.
 
-        if cpm.is_call_to(node, "if_"):
+        if not isinstance(node, ir.FunCall) or cpm.is_call_to(node, "if_"):
             return None
 
         # The first argument that is eligible also transforms all remaining args (They will be
diff --git a/src/gt4py/next/iterator/transforms/fixed_point_transformation.py b/src/gt4py/next/iterator/transforms/fixed_point_transformation.py
new file mode 100644
index 00000000..be34af84
--- /dev/null
+++ b/src/gt4py/next/iterator/transforms/fixed_point_transformation.py
@@ -0,0 +1,67 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2024, ETH Zurich
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+import dataclasses
+import enum
+from typing import ClassVar, Optional, Type
+
+from gt4py import eve
+from gt4py.next.iterator import ir
+from gt4py.next.iterator.type_system import inference as itir_type_inference
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class FixedPointTransformation(eve.NodeTranslator):
+    """
+    Transformation pass that transforms until no transformation is applicable anymore.
+    """
+
+    #: Enum of all transformation (names). The transformations need to be defined as methods
+    #: named `transform_<NAME>`.
+    Transformation: ClassVar[Type[enum.Flag]]
+
+    #: All transformations enabled in this instance, e.g. `Transformation.T1 & Transformation.T2`.
+    #: Usually the default value is chosen to be all transformations.
+    enabled_transformations: enum.Flag
+
+    def visit(self, node, **kwargs):
+        node = super().visit(node, **kwargs)
+        return self.fp_transform(node, **kwargs) if isinstance(node, ir.Node) else node
+
+    def fp_transform(self, node: ir.Node, **kwargs) -> ir.Node:
+        """
+        Transform node until a fixed point is reached, e.g. no transformation is applicable anymore.
+        """
+        while True:
+            new_node = self.transform(node, **kwargs)
+            if new_node is None:
+                break
+            assert new_node != node
+            node = new_node
+        return node
+
+    def transform(self, node: ir.Node, **kwargs) -> Optional[ir.Node]:
+        """
+        Transform node once.
+
+        Execute transformations until one is applicable. As soon as a transformation occured
+        the function will return the transformed node. Note that the transformation itself
+        may call other transformations on child nodes again.
+        """
+        for transformation in self.Transformation:
+            if self.enabled_transformations & transformation:
+                assert isinstance(transformation.name, str)
+                method = getattr(self, f"transform_{transformation.name.lower()}")
+                result = method(node, **kwargs)
+                if result is not None:
+                    assert (
+                        result is not node
+                    )  # transformation should have returned None, since nothing changed
+                    itir_type_inference.reinfer(result)
+                    return result
+        return None
diff --git a/src/gt4py/next/iterator/transforms/pass_manager.py b/src/gt4py/next/iterator/transforms/pass_manager.py
index 6906f81e..0a798484 100644
--- a/src/gt4py/next/iterator/transforms/pass_manager.py
+++ b/src/gt4py/next/iterator/transforms/pass_manager.py
@@ -76,7 +76,7 @@ def apply_common_transforms(
     # required in order to get rid of expressions without a domain (e.g. when a tuple element is never accessed)
     ir = CollapseTuple.apply(
         ir,
-        flags=~CollapseTuple.Flag.PROPAGATE_TO_IF_ON_TUPLES,
+        enabled_transformations=~CollapseTuple.Transformation.PROPAGATE_TO_IF_ON_TUPLES,
         uids=collapse_tuple_uids,
         offset_provider_type=offset_provider_type,
     )  # type: ignore[assignment]  # always an itir.Program
@@ -98,7 +98,7 @@ def apply_common_transforms(
         # is constant-folded the surrounding tuple_get calls can be removed.
         inlined = CollapseTuple.apply(
             inlined,
-            flags=~CollapseTuple.Flag.PROPAGATE_TO_IF_ON_TUPLES,
+            enabled_transformations=~CollapseTuple.Transformation.PROPAGATE_TO_IF_ON_TUPLES,
             uids=collapse_tuple_uids,
             offset_provider_type=offset_provider_type,
         )  # type: ignore[assignment]  # always an itir.Program
@@ -136,7 +136,7 @@ def apply_common_transforms(
             ir,
             ignore_tuple_size=True,
             uids=collapse_tuple_uids,
-            flags=~CollapseTuple.Flag.PROPAGATE_TO_IF_ON_TUPLES,
+            enabled_transformations=~CollapseTuple.Transformation.PROPAGATE_TO_IF_ON_TUPLES,
             offset_provider_type=offset_provider_type,
         )  # type: ignore[assignment]  # always an itir.Program
 
@@ -176,7 +176,7 @@ def apply_fieldview_transforms(
     ir = InlineLambdas.apply(ir, opcount_preserving=True, force_inline_lambda_args=True)
     ir = CollapseTuple.apply(
         ir,
-        flags=~CollapseTuple.Flag.PROPAGATE_TO_IF_ON_TUPLES,
+        enabled_transformations=~CollapseTuple.Transformation.PROPAGATE_TO_IF_ON_TUPLES,
         offset_provider_type=common.offset_provider_to_type(offset_provider),
     )  # type: ignore[assignment] # type is still `itir.Program`
     ir = inline_dynamic_shifts.InlineDynamicShifts.apply(


Test patch:
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py
index 938b9985..916ae4e5 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_collapse_tuple.py
@@ -19,7 +19,7 @@ def test_simple_make_tuple_tuple_get():
     actual = CollapseTuple.apply(
         testee,
         remove_letified_make_tuple_elements=False,
-        flags=CollapseTuple.Flag.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
+        enabled_transformations=CollapseTuple.Transformation.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -37,7 +37,7 @@ def test_nested_make_tuple_tuple_get():
     actual = CollapseTuple.apply(
         testee,
         remove_letified_make_tuple_elements=False,
-        flags=CollapseTuple.Flag.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
+        enabled_transformations=CollapseTuple.Transformation.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -53,7 +53,7 @@ def test_different_tuples_make_tuple_tuple_get():
     actual = CollapseTuple.apply(
         testee,
         remove_letified_make_tuple_elements=False,
-        flags=CollapseTuple.Flag.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
+        enabled_transformations=CollapseTuple.Transformation.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -67,7 +67,7 @@ def test_incompatible_order_make_tuple_tuple_get():
     actual = CollapseTuple.apply(
         testee,
         remove_letified_make_tuple_elements=False,
-        flags=CollapseTuple.Flag.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
+        enabled_transformations=CollapseTuple.Transformation.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -79,7 +79,7 @@ def test_incompatible_size_make_tuple_tuple_get():
     actual = CollapseTuple.apply(
         testee,
         remove_letified_make_tuple_elements=False,
-        flags=CollapseTuple.Flag.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
+        enabled_transformations=CollapseTuple.Transformation.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -91,7 +91,7 @@ def test_merged_with_smaller_outer_size_make_tuple_tuple_get():
     actual = CollapseTuple.apply(
         testee,
         ignore_tuple_size=True,
-        flags=CollapseTuple.Flag.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
+        enabled_transformations=CollapseTuple.Transformation.COLLAPSE_MAKE_TUPLE_TUPLE_GET,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -104,7 +104,7 @@ def test_simple_tuple_get_make_tuple():
     actual = CollapseTuple.apply(
         testee,
         remove_letified_make_tuple_elements=False,
-        flags=CollapseTuple.Flag.COLLAPSE_TUPLE_GET_MAKE_TUPLE,
+        enabled_transformations=CollapseTuple.Transformation.COLLAPSE_TUPLE_GET_MAKE_TUPLE,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -117,7 +117,7 @@ def test_propagate_tuple_get():
     actual = CollapseTuple.apply(
         testee,
         remove_letified_make_tuple_elements=False,
-        flags=CollapseTuple.Flag.PROPAGATE_TUPLE_GET,
+        enabled_transformations=CollapseTuple.Transformation.PROPAGATE_TUPLE_GET,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -135,7 +135,7 @@ def test_letify_make_tuple_elements():
     actual = CollapseTuple.apply(
         testee,
         remove_letified_make_tuple_elements=False,
-        flags=CollapseTuple.Flag.LETIFY_MAKE_TUPLE_ELEMENTS,
+        enabled_transformations=CollapseTuple.Transformation.LETIFY_MAKE_TUPLE_ELEMENTS,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -149,7 +149,7 @@ def test_letify_make_tuple_with_trivial_elements():
     actual = CollapseTuple.apply(
         testee,
         remove_letified_make_tuple_elements=False,
-        flags=CollapseTuple.Flag.LETIFY_MAKE_TUPLE_ELEMENTS,
+        enabled_transformations=CollapseTuple.Transformation.LETIFY_MAKE_TUPLE_ELEMENTS,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -163,7 +163,7 @@ def test_inline_trivial_make_tuple():
     actual = CollapseTuple.apply(
         testee,
         remove_letified_make_tuple_elements=False,
-        flags=CollapseTuple.Flag.INLINE_TRIVIAL_MAKE_TUPLE,
+        enabled_transformations=CollapseTuple.Transformation.INLINE_TRIVIAL_MAKE_TUPLE,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -182,7 +182,7 @@ def test_propagate_to_if_on_tuples():
     actual = CollapseTuple.apply(
         testee,
         remove_letified_make_tuple_elements=False,
-        flags=CollapseTuple.Flag.PROPAGATE_TO_IF_ON_TUPLES,
+        enabled_transformations=CollapseTuple.Transformation.PROPAGATE_TO_IF_ON_TUPLES,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -199,8 +199,8 @@ def test_propagate_to_if_on_tuples_with_let():
     actual = CollapseTuple.apply(
         testee,
         remove_letified_make_tuple_elements=True,
-        flags=CollapseTuple.Flag.PROPAGATE_TO_IF_ON_TUPLES
-        | CollapseTuple.Flag.LETIFY_MAKE_TUPLE_ELEMENTS,
+        enabled_transformations=CollapseTuple.Transformation.PROPAGATE_TO_IF_ON_TUPLES
+        | CollapseTuple.Transformation.LETIFY_MAKE_TUPLE_ELEMENTS,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -213,7 +213,7 @@ def test_propagate_nested_lift():
     actual = CollapseTuple.apply(
         testee,
         remove_letified_make_tuple_elements=False,
-        flags=CollapseTuple.Flag.PROPAGATE_NESTED_LET,
+        enabled_transformations=CollapseTuple.Transformation.PROPAGATE_NESTED_LET,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -249,7 +249,7 @@ def test_if_make_tuple_reorder_cps():
     expected = im.if_(True, im.make_tuple(2, 1), im.make_tuple(4, 3))
     actual = CollapseTuple.apply(
         testee,
-        flags=~CollapseTuple.Flag.PROPAGATE_TO_IF_ON_TUPLES,
+        enabled_transformations=~CollapseTuple.Transformation.PROPAGATE_TO_IF_ON_TUPLES,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -275,7 +275,7 @@ def test_nested_if_make_tuple_reorder_cps():
     )
     actual = CollapseTuple.apply(
         testee,
-        flags=~CollapseTuple.Flag.PROPAGATE_TO_IF_ON_TUPLES,
+        enabled_transformations=~CollapseTuple.Transformation.PROPAGATE_TO_IF_ON_TUPLES,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -291,7 +291,7 @@ def test_if_make_tuple_reorder_cps_nested():
     expected = im.if_(True, im.make_tuple(2, 1, 1), im.make_tuple(4, 3, 3))
     actual = CollapseTuple.apply(
         testee,
-        flags=~CollapseTuple.Flag.PROPAGATE_TO_IF_ON_TUPLES,
+        enabled_transformations=~CollapseTuple.Transformation.PROPAGATE_TO_IF_ON_TUPLES,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )
@@ -306,7 +306,7 @@ def test_if_make_tuple_reorder_cps_external():
     expected = im.if_(True, im.make_tuple(external_ref, 2, 1), im.make_tuple(external_ref, 4, 3))
     actual = CollapseTuple.apply(
         testee,
-        flags=~CollapseTuple.Flag.PROPAGATE_TO_IF_ON_TUPLES,
+        enabled_transformations=~CollapseTuple.Transformation.PROPAGATE_TO_IF_ON_TUPLES,
         allow_undeclared_symbols=True,
         within_stencil=False,
     )