From 7a9489f73ddddd6aff219fc3890bed23e791a9a8 Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Thu, 4 Jan 2024 00:47:33 +0100 Subject: [PATCH] Fix size check in CollapseTuple pass (#1405) --- src/gt4py/next/iterator/transforms/collapse_tuple.py | 3 +++ src/gt4py/next/iterator/type_inference.py | 6 ++++++ 2 files changed, 9 insertions(+) diff --git a/src/gt4py/next/iterator/transforms/collapse_tuple.py b/src/gt4py/next/iterator/transforms/collapse_tuple.py index 7d710fc919..30457f2246 100644 --- 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 index 2375118cd1..68627cfd89 100644 --- 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)