Fix size check in CollapseTuple pass (#1405)

GridTools · Jan 3, 2024 · 7a9489f · 7a9489f
1 parent 100bc7f
commit 7a9489f
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/src/gt4py/next/iterator/transforms/collapse_tuple.py b/src/gt4py/next/iterator/transforms/collapse_tuple.py
@@ -41,6 +41,9 @@ def _get_tuple_size(elem: ir.Node, node_types: Optional[dict] = None) -> int | t
         ):
             return UnknownLength
 
+    if not type_.dtype.has_known_length:
+        return UnknownLength
+
     return len(type_.dtype)
 
 

diff --git a/src/gt4py/next/iterator/type_inference.py b/src/gt4py/next/iterator/type_inference.py
@@ -77,6 +77,12 @@ def __iter__(self) -> abc.Iterator[Type]:
             raise ValueError(f"Can not iterate over partially defined tuple '{self}'.")
         yield from self.others
 
+    @property
+    def has_known_length(self):
+        return isinstance(self.others, EmptyTuple) or (
+            isinstance(self.others, Tuple) and self.others.has_known_length
+        )
+
     def __len__(self) -> int:
         return sum(1 for _ in self)