Repository: GridTools/gt4py
Pull request #1155 (2023-02-07 14:39:33)

Extend ITIR type inference to collect the types of all subexpressions
Type inferrence on ITIR was previously only returning the type of a single expression. This PR introduces extends the inference to collect the type of all subexpressions and introduces a new function `_infer_all` returning a mapping from the (Python) id of a node to its type. As we start using this function we will make it public.

tehrengruber: Blocked by #1182

Patch:
diff --git a/setup.cfg b/setup.cfg
index 786c3ac0..9e727282 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -101,7 +101,7 @@ cuda117 =
 cuda11x =
     cupy-cuda11x
 dace =
-    dace>=0.14.1,<0.15
+    dace>=0.14.2,<0.15
     sympy
 formatting =
     clang-format>=9.0
diff --git a/src/gt4py/cartesian/backend/dace_backend.py b/src/gt4py/cartesian/backend/dace_backend.py
index 29753271..7a9a697c 100644
--- a/src/gt4py/cartesian/backend/dace_backend.py
+++ b/src/gt4py/cartesian/backend/dace_backend.py
@@ -469,16 +469,19 @@ class DaCeComputationCodegen:
 
     def generate_tmp_allocs(self, sdfg):
 
-        global_fmt = "dace_handle.__{sdfg_id}_{name} = allocate(allocator, gt::meta::lazy::id<{dtype}>(), {size})();"
+        global_fmt = (
+            "__{sdfg_id}_{name} = allocate(allocator, gt::meta::lazy::id<{dtype}>(), {size})();"
+        )
         threadlocal_fmt = (
-            "{name} = dace_handle.__{sdfg_id}_{name} + omp_get_thread_num() * ({local_size});"
+            "__{sdfg_id}_{name} = __full__{sdfg_id}_{name} + omp_get_thread_num() * ({local_size});"
         )
         res = []
         for array_sdfg, name, array in sdfg.arrays_recursive():
             if array.transient and array.lifetime == dace.AllocationLifetime.Persistent:
                 if array.storage != dace.StorageType.CPU_ThreadLocal:
+                    fmt = "dace_handle." + global_fmt
                     res.append(
-                        global_fmt.format(
+                        fmt.format(
                             name=name,
                             sdfg_id=array_sdfg.sdfg_id,
                             dtype=array.dtype.ctype,
@@ -486,7 +489,13 @@ class DaCeComputationCodegen:
                         )
                     )
                 else:
-                    fmts = [global_fmt, "#pragma omp parallel", "{{", threadlocal_fmt, "}}"]
+                    fmts = [
+                        "{dtype} *__full" + global_fmt,
+                        "#pragma omp parallel",
+                        "{{",
+                        threadlocal_fmt,
+                        "}}",
+                    ]
                     res.extend(
                         [
                             fmt.format(
diff --git a/src/gt4py/next/iterator/type_inference.py b/src/gt4py/next/iterator/type_inference.py
index 237fe334..5a7ea359 100644
--- a/src/gt4py/next/iterator/type_inference.py
+++ b/src/gt4py/next/iterator/type_inference.py
@@ -12,11 +12,13 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
+import dataclasses
 import typing
 from collections import abc
+from typing import Optional
 
 import gt4py.eve as eve
-from gt4py.next.common import DimensionKind
+from gt4py.next.common import Connectivity, Dimension, DimensionKind
 from gt4py.next.iterator import ir
 from gt4py.next.iterator.embedded import NeighborTableOffsetProvider, StridedNeighborOffsetProvider
 from gt4py.next.iterator.runtime import CartesianAxis
@@ -27,6 +29,15 @@ from gt4py.next.type_inference import Type, TypeVar, freshen, reindex_vars, unif
 
 T = typing.TypeVar("T", bound="Type")
 
+# list of nodes that have a type
+TYPED_IR_NODES: typing.Final = (
+    ir.Expr,
+    ir.FunctionDefinition,
+    ir.StencilClosure,
+    ir.FencilDefinition,
+    ir.Sym,
+)
+
 
 class EmptyTuple(Type):
     def __iter__(self) -> abc.Iterator[Type]:
@@ -233,7 +244,8 @@ class FencilDefinitionType(Type):
 
 
 class LetPolymorphic(Type):
-    """Wrapper for let-polymorphic types.
+    """
+    Wrapper for let-polymorphic types.
 
     Used for fencil-level function definitions.
     """
@@ -379,20 +391,33 @@ def _infer_shift_location_types(shift_args, offset_provider, constraints):
     return current_loc_in, current_loc_out
 
 
-class _TypeInferrer(eve.NodeTranslator):
-    """Visit the full iterator IR tree, convert nodes to respective types and generate constraints."""
+@dataclasses.dataclass
+class _TypeInferrer(eve.traits.VisitorWithSymbolTableTrait, eve.NodeTranslator):
+    """
+    Visit the full iterator IR tree, convert nodes to respective types and generate constraints.
 
-    def __init__(self, offset_provider):
-        self.offset_provider = offset_provider
+    Attributes:
+        collected_types: Mapping from the (Python) id of a node to its type.
+        constraints: Set of constraints, where a constraint is a pair of types that need to agree.
+            See `unify` for more information.
+    """
 
-    def visit_SymRef(
-        self, node: ir.SymRef, constraints: set[tuple[Type, Type]], symtypes: dict[str, Type]
-    ) -> Type:
-        if node.id in symtypes:
-            res = symtypes[node.id]
-            if isinstance(res, LetPolymorphic):
-                return freshen(res.dtype)
-            return res
+    offset_provider: Optional[dict[str, Connectivity | Dimension]]
+    collected_types: dict[int, Type] = dataclasses.field(default_factory=dict)
+    constraints: set[tuple[Type, Type]] = dataclasses.field(default_factory=set)
+
+    def visit(self, node, **kwargs) -> typing.Any:
+        result = super().visit(node, **kwargs)
+        if isinstance(node, TYPED_IR_NODES):
+            assert isinstance(result, Type)
+            self.collected_types[id(node)] = result
+
+        return result
+
+    def visit_Sym(self, node: ir.Sym, **kwargs) -> Type:
+        return TypeVar.fresh()
+
+    def visit_SymRef(self, node: ir.SymRef, *, symtable, **kwargs) -> Type:
         if node.id in BUILTIN_TYPES:
             return freshen(BUILTIN_TYPES[node.id])
         if node.id in (
@@ -406,52 +431,49 @@ class _TypeInferrer(eve.NodeTranslator):
             raise TypeError(
                 f"Builtin '{node.id}' is only supported as applied/called function by the type checker"
             )
+        if node.id in symtable:
+            res = self.collected_types[id(symtable[node.id])]
+            if isinstance(res, LetPolymorphic):
+                return freshen(res.dtype)
+            return res
         if node.id in (ir.BUILTINS - ir.TYPEBUILTINS):
             raise NotImplementedError(f"Missing type definition for builtin '{node.id}'")
 
         return TypeVar.fresh()
 
-    def visit_Literal(
-        self, node: ir.Literal, constraints: set[tuple[Type, Type]], symtypes: dict[str, Type]
-    ) -> Val:
+    def visit_Literal(self, node: ir.Literal, **kwargs) -> Val:
         return Val(kind=Value(), dtype=Primitive(name=node.type))
 
-    def visit_AxisLiteral(
-        self,
-        node: ir.AxisLiteral,
-        constraints: set[tuple[Type, Type]],
-        symtypes: dict[str, Type],
-    ) -> Val:
+    def visit_AxisLiteral(self, node: ir.AxisLiteral, **kwargs) -> Val:
         return Val(kind=Value(), dtype=AXIS_DTYPE, size=Scalar())
 
-    def visit_OffsetLiteral(
-        self,
-        node: ir.OffsetLiteral,
-        constraints: set[tuple[Type, Type]],
-        symtypes: dict[str, Type],
-    ) -> TypeVar:
+    def visit_OffsetLiteral(self, node: ir.OffsetLiteral, **kwargs) -> TypeVar:
         return TypeVar.fresh()
 
     def visit_Lambda(
-        self, node: ir.Lambda, constraints: set[tuple[Type, Type]], symtypes: dict[str, Type]
+        self,
+        node: ir.Lambda,
+        **kwargs,
     ) -> FunctionType:
-        ptypes = {p.id: TypeVar.fresh() for p in node.params}
-        ret = self.visit(node.expr, constraints=constraints, symtypes=symtypes | ptypes)
+        ptypes = {p.id: self.visit(p, **kwargs) for p in node.params}
+        ret = self.visit(node.expr, **kwargs)
         return FunctionType(args=Tuple.from_elems(*(ptypes[p.id] for p in node.params)), ret=ret)
 
     def visit_FunCall(
-        self, node: ir.FunCall, constraints: set[tuple[Type, Type]], symtypes: dict[str, Type]
+        self,
+        node: ir.FunCall,
+        **kwargs,
     ) -> Type:
         if isinstance(node.fun, ir.SymRef):
             if node.fun.id == "make_tuple":
                 # Calls to make_tuple are handled as being part of the grammar,
                 # not as function calls
-                argtypes = self.visit(node.args, constraints=constraints, symtypes=symtypes)
+                argtypes = self.visit(node.args, **kwargs)
                 kind = TypeVar.fresh()
                 size = TypeVar.fresh()
                 dtype = Tuple.from_elems(*(TypeVar.fresh() for _ in argtypes))
                 for d, a in zip(dtype, argtypes):
-                    constraints.add((Val(kind=kind, dtype=d, size=size), a))
+                    self.constraints.add((Val(kind=kind, dtype=d, size=size), a))
                 return Val(kind=kind, dtype=dtype, size=size)
             if node.fun.id == "tuple_get":
                 # Calls to tuple_get are handled as being part of the grammar,
@@ -461,7 +483,7 @@ class _TypeInferrer(eve.NodeTranslator):
                 if not isinstance(node.args[0], ir.Literal) or node.args[0].type != "int":
                     raise TypeError("The first argument to tuple_get must be a literal int")
                 idx = int(node.args[0].value)
-                tup = self.visit(node.args[1], constraints=constraints, symtypes=symtypes)
+                tup = self.visit(node.args[1], **kwargs)
                 kind = TypeVar.fresh()
                 elem = TypeVar.fresh()
                 size = TypeVar.fresh()
@@ -475,19 +497,19 @@ class _TypeInferrer(eve.NodeTranslator):
                     dtype=dtype,
                     size=size,
                 )
-                constraints.add((tup, val))
+                self.constraints.add((tup, val))
                 return Val(kind=kind, dtype=elem, size=size)
             if node.fun.id == "cast_":
                 if len(node.args) != 2:
                     raise TypeError("cast_ requires exactly two arguments.")
-                val_arg_type = self.visit(node.args[0], constraints=constraints, symtypes=symtypes)
+                val_arg_type = self.visit(node.args[0], **kwargs)
                 type_arg = node.args[1]
                 if not isinstance(type_arg, ir.SymRef) or type_arg.id not in ir.TYPEBUILTINS:
                     raise TypeError("The second argument to `cast_` must be a type literal.")
 
                 size = TypeVar.fresh()
 
-                constraints.add(
+                self.constraints.add(
                     (
                         val_arg_type,
                         Val(
@@ -507,7 +529,7 @@ class _TypeInferrer(eve.NodeTranslator):
                 # Calls to shift are handled as being part of the grammar, not
                 # as function calls, as the type depends on the offset provider
                 current_loc_in, current_loc_out = _infer_shift_location_types(
-                    node.args, self.offset_provider, constraints
+                    node.args, self.offset_provider, self.constraints
                 )
                 defined_loc = TypeVar.fresh()
                 dtype_ = TypeVar.fresh()
@@ -532,53 +554,53 @@ class _TypeInferrer(eve.NodeTranslator):
                 )
             if node.fun.id.endswith("domain"):
                 for arg in node.args:
-                    constraints.add(
+                    self.constraints.add(
                         (
                             Val(kind=Value(), dtype=NAMED_RANGE_DTYPE, size=Scalar()),
-                            self.visit(arg, constraints=constraints, symtypes=symtypes),
+                            self.visit(arg, **kwargs),
                         )
                     )
                 return Val(kind=Value(), dtype=DOMAIN_DTYPE, size=Scalar())
 
-        fun = self.visit(node.fun, constraints=constraints, symtypes=symtypes)
-        args = Tuple.from_elems(*self.visit(node.args, constraints=constraints, symtypes=symtypes))
+        fun = self.visit(node.fun, **kwargs)
+        args = Tuple.from_elems(*self.visit(node.args, **kwargs))
         ret = TypeVar.fresh()
-        constraints.add((fun, FunctionType(args=args, ret=ret)))
+        self.constraints.add((fun, FunctionType(args=args, ret=ret)))
         return ret
 
     def visit_FunctionDefinition(
         self,
         node: ir.FunctionDefinition,
-        constraints: set[tuple[Type, Type]],
-        symtypes: dict[str, Type],
-    ) -> FunctionDefinitionType:
-        if node.id in symtypes:
-            raise TypeError(f"Multiple definitions of symbol {node.id}")
-
-        fun = self.visit(
-            ir.Lambda(params=node.params, expr=node.expr),
-            constraints=constraints,
-            symtypes=symtypes,
-        )
-        constraints.add((fun, FunctionType()))
-        return FunctionDefinitionType(name=node.id, fun=fun)
+        **kwargs,
+    ) -> LetPolymorphic:
+        fun = ir.Lambda(params=node.params, expr=node.expr)
+
+        # Since functions defined in a function definition are let-polymorphic we don't want
+        # their parameters to inherit the constraints of the arguments in a call to them. A simple
+        # way to do this is to run the type inference on the function itself and reindex its type
+        # vars when referencing the function, i.e. in a `SymRef`.
+        collected_types = _infer_all(fun, offset_provider=self.offset_provider, reindex=False)
+        fun_type = LetPolymorphic(dtype=collected_types.pop(id(fun)))
+        assert not set(self.collected_types.keys()) & set(collected_types.keys())
+        self.collected_types = {**self.collected_types, **collected_types}
+
+        return fun_type
 
     def visit_StencilClosure(
         self,
         node: ir.StencilClosure,
-        constraints: set[tuple[Type, Type]],
-        symtypes: dict[str, Type],
+        **kwargs,
     ) -> Closure:
-        domain = self.visit(node.domain, constraints=constraints, symtypes=symtypes)
-        stencil = self.visit(node.stencil, constraints=constraints, symtypes=symtypes)
-        output = self.visit(node.output, constraints=constraints, symtypes=symtypes)
-        inputs = Tuple.from_elems(
-            *self.visit(node.inputs, constraints=constraints, symtypes=symtypes)
-        )
+        domain = self.visit(node.domain, **kwargs)
+        stencil = self.visit(node.stencil, **kwargs)
+        output = self.visit(node.output, **kwargs)
+        inputs = Tuple.from_elems(*self.visit(node.inputs, **kwargs))
         output_dtype = TypeVar.fresh()
         output_loc = TypeVar.fresh()
-        constraints.add((domain, Val(kind=Value(), dtype=Primitive(name="domain"), size=Scalar())))
-        constraints.add(
+        self.constraints.add(
+            (domain, Val(kind=Value(), dtype=Primitive(name="domain"), size=Scalar()))
+        )
+        self.constraints.add(
             (
                 output,
                 Val(
@@ -590,7 +612,7 @@ class _TypeInferrer(eve.NodeTranslator):
                 ),
             )
         )
-        constraints.add(
+        self.constraints.add(
             (
                 stencil,
                 FunctionType(args=inputs, ret=Val(kind=Value(), dtype=output_dtype, size=Column())),
@@ -601,47 +623,59 @@ class _TypeInferrer(eve.NodeTranslator):
     def visit_FencilDefinition(
         self,
         node: ir.FencilDefinition,
-        constraints: set[tuple[Type, Type]],
-        symtypes: dict[str, Type],
+        **kwargs,
     ) -> FencilDefinitionType:
         ftypes = []
-        fmap = dict[str, LetPolymorphic]()
         # Note: functions have to be ordered according to Lisp/Scheme `let*`
         # statements; that is, functions can only reference other functions
         # that are defined before
-        for f in node.function_definitions:
-            c = set[tuple[Type, Type]]()
-            ftype: FunctionDefinitionType = self.visit(f, constraints=c, symtypes=symtypes | fmap)
-            ftype = typing.cast(FunctionDefinitionType, unify(ftype, c))
+        for fun_def in node.function_definitions:
+            fun_type: LetPolymorphic = self.visit(fun_def, **kwargs)
+            ftype = FunctionDefinitionType(name=fun_def.id, fun=fun_type.dtype)
             ftypes.append(ftype)
-            fmap[ftype.name] = LetPolymorphic(dtype=ftype.fun)
 
-        params = {p.id: TypeVar.fresh() for p in node.params}
-        self.visit(node.closures, constraints=constraints, symtypes=symtypes | fmap | params)
+        params = [self.visit(p, **kwargs) for p in node.params]
+        self.visit(node.closures, **kwargs)
         return FencilDefinitionType(
             name=node.id,
             fundefs=Tuple.from_elems(*ftypes),
-            params=Tuple.from_elems(*params.values()),
+            params=Tuple.from_elems(*params),
         )
 
 
+def _infer_all(
+    node: ir.Node,
+    offset_provider: Optional[dict[str, Connectivity | Dimension]] = None,
+    reindex: bool = True,
+) -> dict[int, Type]:
+    """
+    Infer the types of the child expressions of a given iterator IR expression.
+
+    The result is a dictionary mapping the (Python) id of child nodes to their type.
+    """
+    # Collect preliminary types of all nodes and constraints on them
+    inferrer = _TypeInferrer(offset_provider=offset_provider)
+    inferrer.visit(node)
+
+    # Ensure dict order is pre-order of the tree
+    collected_types = dict(reversed(inferrer.collected_types.items()))
+
+    # Compute the most general type that satisfies all constraints
+    unified_types = unify(list(collected_types.values()), inferrer.constraints)
+
+    if reindex:
+        unified_types = reindex_vars(list(unified_types))
+
+    return {id_: unified_type for id_, unified_type in zip(collected_types.keys(), unified_types)}
+
+
 def infer(
     expr: ir.Node,
-    symtypes: typing.Optional[dict[str, Type]] = None,
     offset_provider: typing.Optional[dict[str, typing.Any]] = None,
 ) -> Type:
     """Infer the type of the given iterator IR expression."""
-    if symtypes is None:
-        symtypes = dict()
-
-    # Collect constraints
-    constraints = set[tuple[Type, Type]]()
-    dtype = _TypeInferrer(offset_provider=offset_provider).visit(
-        expr, constraints=constraints, symtypes=symtypes
-    )
-    # Compute the most general type that satisfies all constraints
-    unified = unify(dtype, constraints)
-    return reindex_vars(unified)
+    inferred_types = _infer_all(expr, offset_provider)
+    return inferred_types[id(expr)]
 
 
 class PrettyPrinter(eve.NodeTranslator):
diff --git a/src/gt4py/next/type_inference.py b/src/gt4py/next/type_inference.py
index adbd8d03..b6a74998 100644
--- a/src/gt4py/next/type_inference.py
+++ b/src/gt4py/next/type_inference.py
@@ -93,24 +93,50 @@ class _TypeVarReindexer(eve.NodeTranslator):
         return node.__class__(**new_values)
 
 
-def freshen(dtype: T) -> T:
+@typing.overload
+def freshen(dtypes: list[T]) -> list[T]:
+    ...
+
+
+@typing.overload
+def freshen(dtypes: T) -> T:
+    ...
+
+
+def freshen(dtypes: list[T] | T) -> list[T] | T:
     """Re-instantiate `dtype` with fresh type variables."""
+    if not isinstance(dtypes, list):
+        assert isinstance(dtypes, Type)
+        return freshen([dtypes])[0]
 
     def indexer(index_map: dict[int, int]) -> int:
         return TypeVar.fresh_index()
 
     index_map = dict[int, int]()
-    return _TypeVarReindexer(indexer).visit(dtype, index_map=index_map)
+    return [_TypeVarReindexer(indexer).visit(dtype, index_map=index_map) for dtype in dtypes]
+
+
+@typing.overload
+def reindex_vars(dtypes: list[T]) -> list[T]:
+    ...
+
 
+@typing.overload
+def reindex_vars(dtypes: T) -> T:
+    ...
 
-def reindex_vars(dtype: T) -> T:
+
+def reindex_vars(dtypes: list[T] | T) -> list[T] | T:
     """Reindex all type variables, to have nice indices starting at zero."""
+    if not isinstance(dtypes, list):
+        assert isinstance(dtypes, Type)
+        return reindex_vars([dtypes])[0]
 
     def indexer(index_map: dict[int, int]) -> int:
         return len(index_map)
 
     index_map = dict[int, int]()
-    return _TypeVarReindexer(indexer).visit(dtype, index_map=index_map)
+    return [_TypeVarReindexer(indexer).visit(dtype, index_map=index_map) for dtype in dtypes]
 
 
 class _FreeVariables(eve.NodeVisitor):
@@ -129,11 +155,13 @@ def _free_variables(x: Type) -> set[TypeVar]:
 
 
 class _Dedup(eve.NodeTranslator):
-    """Deduplicate nodes that have the same value but a different `id`."""
+    """Deduplicate type nodes that have the same value but a different `id`."""
 
-    def visit(self, node: T, *, memo: dict[T, T]) -> T:  # type: ignore[override]
-        node = super().visit(node, memo=memo)
-        return memo.setdefault(node, node)
+    def visit(self, node, *, memo: dict[T, T]) -> typing.Any:  # type: ignore[override]
+        if isinstance(node, Type):
+            node = super().visit(node, memo=memo)
+            return memo.setdefault(node, node)
+        return node
 
 
 class _Renamer:
@@ -194,16 +222,10 @@ class _Renamer:
         except KeyError:
             return
 
-        follow_up_renames = list[tuple[Type, Type]]()
         for node, field in nodes:
             # Default case: just update a field value of the node
             self._update_node(node, field, replacement)
 
-        # Handle follow-up renames
-        for s, d in follow_up_renames:
-            self.register(d)
-            self.rename(s, d)
-
 
 class _Box(Type):
     """Simple value holder, used for wrapping root nodes of a type tree.
@@ -221,20 +243,21 @@ class _Unifier:
     type variable renaming.
     """
 
-    def __init__(self, dtype: Type, constraints: set[tuple[Type, Type]]) -> None:
+    def __init__(self, dtypes: list[Type], constraints: set[tuple[Type, Type]]) -> None:
         # Wrap the original `dtype` and all `constraints` to make sure they have a parent node and
         # thus the root nodes are correctly handled by the renamer
-        self._dtype = _Box(value=dtype)
+        self._dtypes = [_Box(value=dtype) for dtype in dtypes]
         self._constraints = [(_Box(value=s), _Box(value=t)) for s, t in constraints]
 
         # Create a renamer and register `dtype` and all `constraints` types
         self._renamer = _Renamer()
-        self._renamer.register(self._dtype)
+        for dtype in self._dtypes:
+            self._renamer.register(dtype)
         for s, t in self._constraints:
             self._renamer.register(s)
             self._renamer.register(t)
 
-    def unify(self) -> Type:
+    def unify(self) -> list[Type]:
         """Run the unification."""
         while self._constraints:
             constraint = self._constraints.pop()
@@ -247,7 +270,7 @@ class _Unifier:
                     f"Can not satisfy constraint: {constraint[0].value} ≡ {constraint[1].value}"
                 )
 
-        return self._dtype.value
+        return [dtype.value for dtype in self._dtypes]  # unbox
 
     def _rename(self, x: Type, y: Type) -> None:
         """Type renaming/replacement."""
@@ -297,13 +320,26 @@ class _Unifier:
         return False
 
 
-def unify(dtype: Type, constraints: set[tuple[Type, Type]]) -> Type:
+@typing.overload
+def unify(dtypes: list[Type], constraints: set[tuple[Type, Type]]) -> list[Type]:
+    ...
+
+
+@typing.overload
+def unify(dtypes: Type, constraints: set[tuple[Type, Type]]) -> Type:
+    ...
+
+
+def unify(dtypes: list[Type] | Type, constraints: set[tuple[Type, Type]]) -> list[Type] | Type:
     """Unify all given constraints."""
-    # Deduplicate nodes, this can speed up later things a bit
-    memo: dict = {}
-    dtype = _Dedup().visit(dtype, memo=memo)
+    if isinstance(dtypes, Type):
+        return unify([dtypes], constraints)[0]
+
+    # Deduplicate type nodes, this can speed up later things a bit
+    memo = dict[Type, Type]()
+    dtypes = [_Dedup().visit(dtype, memo=memo) for dtype in dtypes]
     constraints = {_Dedup().visit(c, memo=memo) for c in constraints}
     del memo
 
-    unifier = _Unifier(dtype, constraints)
+    unifier = _Unifier(dtypes, constraints)
     return unifier.unify()


Test patch:
diff --git a/tests/next_tests/iterator_tests/test_type_inference.py b/tests/next_tests/iterator_tests/test_type_inference.py
index ee6bf711..502685bf 100644
--- a/tests/next_tests/iterator_tests/test_type_inference.py
+++ b/tests/next_tests/iterator_tests/test_type_inference.py
@@ -561,9 +561,8 @@ def test_partial_shift_with_unstructured_offset_provider():
 
 def test_function_definition():
     testee = ir.FunctionDefinition(id="f", params=[ir.Sym(id="x")], expr=ir.SymRef(id="x"))
-    expected = ti.FunctionDefinitionType(
-        name="f",
-        fun=ti.FunctionType(
+    expected = ti.LetPolymorphic(
+        dtype=ti.FunctionType(
             args=ti.Tuple.from_elems(
                 ti.TypeVar(idx=0),
             ),
@@ -572,7 +571,7 @@ def test_function_definition():
     )
     inferred = ti.infer(testee)
     assert inferred == expected
-    assert ti.pformat(inferred) == "f :: (T₀) → T₀"
+    assert ti.pformat(inferred.dtype) == "(T₀) → T₀"
 
 
 CARTESIAN_DOMAIN = ir.FunCall(