Repository: GridTools/gt4py
Pull request #1706 (2024-10-24 08:54:52)

feat[next]: Allow partial type inference on ITIR


Patch:
diff --git a/src/gt4py/next/iterator/transforms/collapse_tuple.py b/src/gt4py/next/iterator/transforms/collapse_tuple.py
index 40d98208..b61fb2ba 100644
--- a/src/gt4py/next/iterator/transforms/collapse_tuple.py
+++ b/src/gt4py/next/iterator/transforms/collapse_tuple.py
@@ -192,8 +192,13 @@ class CollapseTuple(eve.PreserveLocationVisitor, eve.NodeTranslator):
                     # tuple argument differs, just continue with the rest of the tree
                     return None
 
-            assert self.ignore_tuple_size or isinstance(first_expr.type, ts.TupleType)
-            if self.ignore_tuple_size or len(first_expr.type.types) == len(node.args):  # type: ignore[union-attr] # ensured by assert above
+            assert self.ignore_tuple_size or isinstance(
+                first_expr.type, (ts.TupleType, ts.DeferredType)
+            )
+            if self.ignore_tuple_size or (
+                isinstance(first_expr.type, ts.TupleType)
+                and len(first_expr.type.types) == len(node.args)
+            ):
                 return first_expr
         return None
 
diff --git a/src/gt4py/next/iterator/type_system/inference.py b/src/gt4py/next/iterator/type_system/inference.py
index 4640aa11..edcb9b54 100644
--- a/src/gt4py/next/iterator/type_system/inference.py
+++ b/src/gt4py/next/iterator/type_system/inference.py
@@ -84,7 +84,7 @@ def _is_compatible_type(type_a: ts.TypeSpec, type_b: ts.TypeSpec):
             is_compatible &= _is_compatible_type(arg_a, arg_b)
         is_compatible &= _is_compatible_type(type_a.returns, type_b.returns)
     else:
-        is_compatible &= type_a == type_b
+        is_compatible &= type_info.is_concretizable(type_a, type_b)
 
     return is_compatible
 
@@ -435,7 +435,7 @@ class ITIRTypeInference(eve.NodeTranslator):
         result = super().visit(node, **kwargs)
         if isinstance(node, itir.Node):
             if isinstance(result, ts.TypeSpec):
-                if node.type:
+                if node.type and not isinstance(node.type, ts.DeferredType):
                     assert _is_compatible_type(node.type, result)
                 node.type = result
             elif isinstance(result, ObservableTypeSynthesizer) or result is None:
@@ -511,17 +511,18 @@ class ITIRTypeInference(eve.NodeTranslator):
                 path,
                 node.expr.type,
             )
-            assert isinstance(target_type, ts.FieldType)
-            assert isinstance(expr_type, ts.FieldType)
+            assert isinstance(target_type, (ts.FieldType, ts.DeferredType))
+            assert isinstance(expr_type, (ts.FieldType, ts.DeferredType))
             # TODO(tehrengruber): The lowering emits domains that always have the horizontal domain
             #  first. Since the expr inherits the ordering from the domain this can lead to a mismatch
             #  between the target and expr (e.g. when the target has dimension K, Vertex). We should
             #  probably just change the behaviour of the lowering. Until then we do this more
             #  complicated comparison.
-            assert (
-                set(expr_type.dims) == set(target_type.dims)
-                and target_type.dtype == expr_type.dtype
-            )
+            if isinstance(target_type, ts.FieldType) and isinstance(expr_type, ts.FieldType):
+                assert (
+                    set(expr_type.dims).issubset(set(target_type.dims))
+                    and target_type.dtype == expr_type.dtype
+                )
 
     # TODO(tehrengruber): Remove after new ITIR format with apply_stencil is used everywhere
     def visit_StencilClosure(self, node: itir.StencilClosure, *, ctx) -> it_ts.StencilClosureType:
@@ -623,6 +624,8 @@ class ITIRTypeInference(eve.NodeTranslator):
             self.visit(tuple_, ctx=ctx)  # ensure tuple is typed
             assert isinstance(index_literal, itir.Literal)
             index = int(index_literal.value)
+            if isinstance(tuple_.type, ts.DeferredType):
+                return ts.DeferredType(constraint=None)
             assert isinstance(tuple_.type, ts.TupleType)
             return tuple_.type.types[index]
 
diff --git a/src/gt4py/next/iterator/type_system/type_synthesizer.py b/src/gt4py/next/iterator/type_system/type_synthesizer.py
index c836de13..c55cfd8d 100644
--- a/src/gt4py/next/iterator/type_system/type_synthesizer.py
+++ b/src/gt4py/next/iterator/type_system/type_synthesizer.py
@@ -291,6 +291,9 @@ def as_fieldop(
 
     @TypeSynthesizer
     def applied_as_fieldop(*fields) -> ts.FieldType | ts.DeferredType:
+        if any(isinstance(f, ts.DeferredType) for f in fields):
+            return ts.DeferredType(constraint=None)
+
         stencil_return = stencil(
             *(_convert_as_fieldop_input_to_iterator(domain, field) for field in fields),
             offset_provider=offset_provider,


Test patch:
diff --git a/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py b/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py
index 20a1d7e9..7b6214fb 100644
--- a/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py
+++ b/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py
@@ -103,6 +103,10 @@ def expression_test_cases():
         # tuple_get
         (im.tuple_get(0, im.make_tuple(im.ref("a", int_type), im.ref("b", bool_type))), int_type),
         (im.tuple_get(1, im.make_tuple(im.ref("a", int_type), im.ref("b", bool_type))), bool_type),
+        (
+            im.tuple_get(0, im.ref("t", ts.DeferredType(constraint=None))),
+            ts.DeferredType(constraint=None),
+        ),
         # neighbors
         (
             im.neighbors("E2V", im.ref("a", it_on_e_of_e_type)),
@@ -171,6 +175,12 @@ def expression_test_cases():
             )(im.ref("inp1", float_i_field), im.ref("inp2", float_i_field)),
             ts.TupleType(types=[float_i_field, float_i_field]),
         ),
+        (
+            im.as_fieldop(im.lambda_("x")(im.deref("x")))(
+                im.ref("inp", ts.DeferredType(constraint=None))
+            ),
+            ts.DeferredType(constraint=None),
+        ),
         # if in field-view scope
         (
             im.if_(
@@ -458,6 +468,33 @@ def test_program_tuple_setat_short_target():
     )
 
 
+def test_program_setat_without_domain():
+    cartesian_domain = im.call("cartesian_domain")(
+        im.call("named_range")(itir.AxisLiteral(value="IDim"), 0, 1)
+    )
+
+    testee = itir.Program(
+        id="f",
+        function_definitions=[],
+        params=[im.sym("inp", float_i_field), im.sym("out", float_i_field)],
+        declarations=[],
+        body=[
+            itir.SetAt(
+                expr=im.as_fieldop(im.lambda_("x")(im.deref("x")))("inp"),
+                domain=cartesian_domain,
+                target=im.ref("out", float_i_field),
+            )
+        ],
+    )
+
+    result = itir_type_inference.infer(testee, offset_provider={"Ioff": IDim})
+
+    assert (
+        isinstance(result.body[0].expr.type, ts.DeferredType)
+        and result.body[0].expr.type.constraint == ts.FieldType
+    )
+
+
 def test_if_stmt():
     cartesian_domain = im.call("cartesian_domain")(
         im.call("named_range")(itir.AxisLiteral(value="IDim"), 0, 1)
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 bcf8b726..720076c8 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
@@ -8,6 +8,7 @@
 
 from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.iterator.transforms.collapse_tuple import CollapseTuple
+from gt4py.next.type_system import type_specifications as ts
 
 
 def test_simple_make_tuple_tuple_get():
@@ -213,3 +214,11 @@ def test_if_on_tuples_with_let():
         testee, remove_letified_make_tuple_elements=False, allow_undeclared_symbols=True
     )
     assert actual == expected
+
+
+def test_tuple_get_on_untyped_ref():
+    # test pass gracefully handles untyped nodes.
+    testee = im.tuple_get(0, im.ref("val", ts.DeferredType(constraint=None)))
+
+    actual = CollapseTuple.apply(testee, allow_undeclared_symbols=True)
+    assert actual == testee