reverted some edits for collapse_tuple

GridTools · Nov 3, 2023 · 1e2de16 · 1e2de16
1 parent 3c463a6
commit 1e2de16
Showing 1 changed file with 2 additions and 11 deletions.
diff --git a/src/gt4py/next/iterator/transforms/collapse_tuple.py b/src/gt4py/next/iterator/transforms/collapse_tuple.py
@@ -39,8 +39,6 @@ class CollapseTuple(eve.NodeTranslator):
     collapse_make_tuple_tuple_get: bool
     collapse_tuple_get_make_tuple: bool
 
-    _node_types: dict[int, type_inference.Type]
-
     @classmethod
     def apply(
         cls,
@@ -57,13 +55,8 @@ def apply(
         If `ignore_tuple_size`, apply the transformation even if length of the inner tuple
         is greater than the length of the outer tuple.
         """
-        node_types = it_type_inference.infer_all(node)
-
         return cls(
-            ignore_tuple_size,
-            collapse_make_tuple_tuple_get,
-            collapse_tuple_get_make_tuple,
-            node_types,
+            ignore_tuple_size, collapse_make_tuple_tuple_get, collapse_tuple_get_make_tuple
         ).visit(node)
 
     def visit_FunCall(self, node: ir.FunCall, **kwargs) -> ir.Node:
@@ -86,9 +79,7 @@ def visit_FunCall(self, node: ir.FunCall, **kwargs) -> ir.Node:
                     # tuple argument differs, just continue with the rest of the tree
                     return self.generic_visit(node)
 
-            if self.ignore_tuple_size or _get_tuple_size(self._node_types[id(first_expr)]) == len(
-                node.args
-            ):
+            if self.ignore_tuple_size or _get_tuple_size(first_expr) == len(node.args):
                 return first_expr
         if (
             self.collapse_tuple_get_make_tuple