Address review comments

GridTools · Nov 15, 2024 · 75695d9 · 75695d9
1 parent 16f143f
commit 75695d9
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/src/gt4py/next/iterator/transforms/collapse_tuple.py b/src/gt4py/next/iterator/transforms/collapse_tuple.py
@@ -289,9 +289,12 @@ def transform_inline_trivial_make_tuple(self, node: ir.FunCall, **kwargs) -> Opt
         return None
 
     def transform_propagate_to_if_on_tuples(self, node: ir.FunCall, **kwargs) -> Optional[ir.Node]:
-        # TODO(tehrengruber): This significantly increases the size of the tree. Skip transformation
-        #  in local-view for now. Revisit.
-        if not cpm.is_call_to(node, "if_") and not kwargs["within_stencil"]:
+        if kwargs["within_stencil"]:
+            # TODO(tehrengruber): This significantly increases the size of the tree. Skip transformation
+            #  in local-view for now. Revisit.
+            return None
+
+        if 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]`