From a07a7d0814c3da618c95cc4e84c66c67f5aa2458 Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Tue, 2 Jul 2024 11:58:11 +0200 Subject: [PATCH] feat[next]: Check fencil/program args in ITIR type inference (#1565) When the new ITIR type inference was introduced this broke the code generation in icon4py. This PR adds a small assert and improves the docstring. --- src/gt4py/next/iterator/type_system/inference.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/gt4py/next/iterator/type_system/inference.py b/src/gt4py/next/iterator/type_system/inference.py index c8f80daace..e6bee51c95 100644 --- a/src/gt4py/next/iterator/type_system/inference.py +++ b/src/gt4py/next/iterator/type_system/inference.py @@ -331,6 +331,11 @@ def apply( allow_undeclared_symbols: Allow references to symbols that don't have a corresponding declaration. This is useful for testing or inference on partially inferred sub-nodes. + Preconditions: + + All parameters in :class:`itir.Program` and :class:`itir.FencilDefinition` must have a type + defined, as they are the starting point for type propagation. + Design decisions: - Lamba functions are monomorphic Builtin functions like ``plus`` are by design polymorphic and only their argument and return @@ -378,6 +383,12 @@ def apply( if not allow_undeclared_symbols: node = RemoveTypes().visit(node) + if isinstance(node, (itir.FencilDefinition, itir.Program)): + assert all(isinstance(param.type, ts.DataType) for param in node.params), ( + "All parameters in 'itir.Program' and 'itir.FencilDefinition' must have a type " + "defined, as they are the starting point for type propagation.", + ) + instance = cls( offset_provider=offset_provider, dimensions=(