Repository: GridTools/gt4py
Pull request #1076 (2022-12-06 12:07:13)

Fix mypy foast to itir
Un-ignore `foast_to_itir` and fix all the reported mypy errors





Patch:
diff --git a/setup.cfg b/setup.cfg
index 5028dc43..35b16039 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -232,9 +232,6 @@ ignore_errors = True
 [mypy-functional.ffront.dialect_parser]
 ignore_errors = True
 
-[mypy-functional.ffront.foast_to_itir]
-ignore_errors = True
-
 [mypy-functional.ffront.func_to_foast]
 ignore_errors = True
 
diff --git a/src/functional/ffront/foast_passes/type_deduction.py b/src/functional/ffront/foast_passes/type_deduction.py
index c5481a4b..0deacc12 100644
--- a/src/functional/ffront/foast_passes/type_deduction.py
+++ b/src/functional/ffront/foast_passes/type_deduction.py
@@ -278,10 +278,11 @@ class FieldOperatorTypeDeduction(traits.VisitorWithSymbolTableTrait, NodeTransla
                 if isinstance(index, tuple):
                     lower, upper = index
                     new_type = ct.TupleType(types=[t for t in values.type.types[lower:upper]])
-                    # TODO(samkellerhals): possibly change definition of Starred node to not need nested id attribute.
                     new_target = foast.Starred(
                         id=foast.DataSymbol(
-                            id=old_target.id.id, location=old_target.location, type=new_type  # type: ignore
+                            id=self.visit(old_target.id).id,
+                            location=old_target.location,
+                            type=new_type,
                         ),
                         type=new_type,
                         location=old_target.location,
diff --git a/src/functional/ffront/foast_to_itir.py b/src/functional/ffront/foast_to_itir.py
index 964715a2..21418a01 100644
--- a/src/functional/ffront/foast_to_itir.py
+++ b/src/functional/ffront/foast_to_itir.py
@@ -45,9 +45,7 @@ class ITIRTypeKind(enum.Enum):
     ENCAPSULATED_ITERATOR = 2
 
 
-def iterator_type_kind(
-    symbol_type: ct.ScalarType | ct.FieldType | ct.TupleType,
-) -> ITIRTypeKind:
+def iterator_type_kind(symbol_type: ct.SymbolType) -> ITIRTypeKind:
     """
     Return the corresponding type kind (on iterator level) to a FOAST expression of the given symbol type.
 
@@ -100,7 +98,7 @@ def is_expr_with_iterator_type_kind(it_type_kind: ITIRTypeKind) -> Callable[[foa
     return predicate
 
 
-def to_value(node: foast.LocatedNode) -> Callable[[itir.Expr], itir.Expr]:
+def to_value(node: foast.Expr) -> Callable[[itir.Expr], itir.Expr]:
     """
     Either ``deref_`` or noop callable depending on the input node.
 
@@ -293,7 +291,7 @@ class FieldOperatorLowering(NodeTranslator):
             return im.make_tuple_(*elts)
         raise AssertionError("Unexpected `IteratorTypeKind`.")
 
-    def _lift_if_field(self, node: foast.LocatedNode) -> Callable[[itir.Expr], itir.Expr]:
+    def _lift_if_field(self, node: foast.Expr) -> Callable[[itir.FunCall], itir.FunCall]:
         if iterator_type_kind(node.type) is ITIRTypeKind.VALUE:
             return lambda x: x
         elif iterator_type_kind(node.type) is ITIRTypeKind.ITERATOR:
@@ -451,23 +449,24 @@ class FieldOperatorLowering(NodeTranslator):
 
         return lowered_arg
 
-    def _visit_math_built_in(self, node: foast.Call, **kwargs) -> itir.Expr:
+    def _visit_math_built_in(self, node: foast.Call, **kwargs) -> itir.FunCall:
         args = tuple(to_value(arg)(self.visit(arg, **kwargs)) for arg in node.args)
         return self._lift_if_field(node)(im.call_(self.visit(node.func, **kwargs))(*args))
 
     def _visit_neighbor_sum(self, node: foast.Call, **kwargs) -> itir.FunCall:
         return self._visit_reduce(node, **kwargs)
 
-    def _visit_type_constr(self, node: foast.Call, **kwargs) -> itir.Literal:
+    def _visit_type_constr(self, node: foast.Call, **kwargs) -> itir.FunCall:
         if isinstance(node.args[0], foast.Constant):
-            target_type = fbuiltins.BUILTINS[node.type.kind.name.lower()]
-            source_type = {**fbuiltins.BUILTINS, "string": str}[node.args[0].type.kind.name.lower()]
+            node_kind = self.visit(node.type).kind.name.lower()
+            target_type = fbuiltins.BUILTINS[node_kind]
+            source_type = {**fbuiltins.BUILTINS, "string": str}[node.args[0].type.__str__().lower()]
             if target_type is bool and source_type is not bool:
                 return im.literal_(str(bool(source_type(node.args[0].value))), node.func.id)
-            return im.literal_(str(node.args[0].value), node.type.kind.name.lower())
+            return im.literal_(str(node.args[0].value), node_kind)
         raise FieldOperatorLoweringError(f"Encountered a type cast, which is not supported: {node}")
 
-    def _make_literal(self, val: Any, type_: ct.ScalarType | ct.TupleType) -> itir.Literal:
+    def _make_literal(self, val: Any, type_: ct.SymbolType) -> itir.Literal:
         # TODO(tehrengruber): check constant of this type is supported in iterator ir
         if isinstance(type_, ct.TupleType):
             return im.make_tuple_(
diff --git a/src/functional/ffront/func_to_past.py b/src/functional/ffront/func_to_past.py
index f59dbf24..d131afd7 100644
--- a/src/functional/ffront/func_to_past.py
+++ b/src/functional/ffront/func_to_past.py
@@ -37,6 +37,8 @@ class ProgramParser(DialectParser[past.Program]):
     def _postprocess_dialect_ast(
         cls, output_node: past.Program, closure_vars: dict[str, Any], annotations: dict[str, Any]
     ) -> past.Program:
+        if "return" in annotations and not isinstance(None, annotations["return"]):
+            raise ProgramSyntaxError("Program should not have a return value!")
         output_node = ClosureVarTypeDeduction.apply(output_node, closure_vars)
         return ProgramTypeDeduction.apply(output_node)
 
diff --git a/src/functional/ffront/symbol_makers.py b/src/functional/ffront/symbol_makers.py
index 65f84464..f507d0b0 100644
--- a/src/functional/ffront/symbol_makers.py
+++ b/src/functional/ffront/symbol_makers.py
@@ -170,9 +170,8 @@ def make_symbol_type_from_value(value: Any) -> ct.SymbolType:
     elif isinstance(value, common.Dimension):
         symbol_type = ct.DimensionType(dim=value)
     elif isinstance(value, LocatedField):
-        dims = list(value.axes)
         dtype = make_symbol_type_from_typing(value.dtype.type)
-        symbol_type = ct.FieldType(dims=dims, dtype=dtype)
+        symbol_type = ct.FieldType(dims=list(value.axes), dtype=dtype)
     elif isinstance(value, tuple):
         # Since the elements of the tuple might be one of the special cases
         # above, we can not resort to generic `infer_type` but need to do it
diff --git a/src/functional/ffront/type_info.py b/src/functional/ffront/type_info.py
index fe1ac138..d0effa71 100644
--- a/src/functional/ffront/type_info.py
+++ b/src/functional/ffront/type_info.py
@@ -140,7 +140,7 @@ def is_floating_point(symbol_type: ct.SymbolType) -> bool:
     True
     >>> is_floating_point(ct.ScalarType(kind=ct.ScalarKind.INT32))
     False
-    >>> is_floating_point(ct.FieldType(dims=[], dtype=ct.ScalarType(kind=ct.ScalarKind.FLOAT32)))
+    >>> is_floating_point(ct.FieldType(dims=[Dimension(value="I")], dtype=ct.ScalarType(kind=ct.ScalarKind.FLOAT64)))
     True
     """
     return extract_dtype(symbol_type).kind in [
@@ -161,7 +161,7 @@ def is_integral(symbol_type: ct.SymbolType) -> bool:
     True
     >>> is_integral(ct.ScalarType(kind=ct.ScalarKind.FLOAT32))
     False
-    >>> is_integral(ct.FieldType(dims=[], dtype=ct.ScalarType(kind=ct.ScalarKind.INT)))
+    >>> is_integral(ct.FieldType(dims=[Dimension(value="I")], dtype=ct.ScalarType(kind=ct.ScalarKind.INT)))
     True
     """
     return extract_dtype(symbol_type).kind in [
@@ -171,6 +171,28 @@ def is_integral(symbol_type: ct.SymbolType) -> bool:
     ]
 
 
+def is_number(symbol_type: ct.SymbolType) -> bool:
+    """
+    Check if ``symbol_type`` is either intergral or float.
+
+    Examples:
+    ---------
+    >>> is_number(ct.ScalarType(kind=ct.ScalarKind.FLOAT64))
+    True
+    >>> is_number(ct.ScalarType(kind=ct.ScalarKind.INT32))
+    True
+    >>> is_number(ct.ScalarType(kind=ct.ScalarKind.BOOL))
+    False
+    >>> is_number(ct.FieldType(dims=[], dtype=ct.ScalarType(kind=ct.ScalarKind.INT)))
+    False
+    """
+    if not isinstance(symbol_type, ct.ScalarType):
+        return False
+    # TODO(nfarabullini): re-factor is_arithmetic such that it only checks for scalars
+    #  and the emtpy field pass in an another function
+    return is_arithmetic(symbol_type)
+
+
 def is_logical(symbol_type: ct.SymbolType) -> bool:
     return extract_dtype(symbol_type).kind is ct.ScalarKind.BOOL
 
@@ -286,6 +308,25 @@ def is_concretizable(symbol_type: ct.SymbolType, to_type: ct.SymbolType) -> bool
     return False
 
 
+def _is_zero_dim_field(field: ct.SymbolType) -> bool:
+    return isinstance(field, ct.FieldType) and field.dims != Ellipsis and len(field.dims) == 0
+
+
+def promote_zero_dims(
+    args: list[ct.SymbolType], function_type: ct.FieldOperatorType | ct.ProgramType
+) -> list:
+    """Promote arg types to zero dimensional fields if compatible and required by function signature."""
+    new_args = args.copy()
+    for arg_i, arg in enumerate(args):
+        def_type = function_type.definition.args[arg_i]
+        if _is_zero_dim_field(def_type) and is_number(arg):
+            if extract_dtype(def_type) == extract_dtype(arg):
+                new_args[arg_i] = def_type
+            else:
+                raise GTTypeError(f"{arg} is not compatible with {def_type}.")
+    return new_args
+
+
 def promote(*types: ct.FieldType | ct.ScalarType) -> ct.FieldType | ct.ScalarType:
     """
     Promote a set of field or scalar types to a common type.
@@ -499,7 +540,9 @@ def function_signature_incompatibilities(
 
 @function_signature_incompatibilities.register
 def function_signature_incompatibilities_func(
-    func_type: ct.FunctionType, args: list[ct.SymbolType], kwargs: dict[str, ct.SymbolType]
+    func_type: ct.FunctionType,
+    args: list[ct.SymbolType],
+    kwargs: dict[str, ct.SymbolType],
 ) -> Iterator[str]:
 
     # check positional arguments
@@ -528,7 +571,8 @@ def function_signature_incompatibilities_func(
 def function_signature_incompatibilities_fieldop(
     fieldop_type: ct.FieldOperatorType, args: list[ct.SymbolType], kwargs: dict[str, ct.SymbolType]
 ) -> Iterator[str]:
-    yield from function_signature_incompatibilities_func(fieldop_type.definition, args, kwargs)
+    new_args = promote_zero_dims(args, fieldop_type)
+    yield from function_signature_incompatibilities_func(fieldop_type.definition, new_args, kwargs)
 
 
 @function_signature_incompatibilities.register
@@ -584,7 +628,8 @@ def function_signature_incompatibilities_scanop(
 def function_signature_incompatibilities_program(
     program_type: ct.ProgramType, args: list[ct.SymbolType], kwargs: dict[str, ct.SymbolType]
 ) -> Iterator[str]:
-    yield from function_signature_incompatibilities_func(program_type.definition, args, kwargs)
+    new_args = promote_zero_dims(args, program_type)
+    yield from function_signature_incompatibilities_func(program_type.definition, new_args, kwargs)
 
 
 @function_signature_incompatibilities.register


Test patch:
diff --git a/tests/functional_tests/ffront_tests/test_execution.py b/tests/functional_tests/ffront_tests/test_execution.py
index 5d19b903..5da4d656 100644
--- a/tests/functional_tests/ffront_tests/test_execution.py
+++ b/tests/functional_tests/ffront_tests/test_execution.py
@@ -1125,7 +1125,7 @@ def test_docstring():
         return a
 
     @program
-    def test_docstring(a: Field[[IDim], float64]) -> Field[[IDim], float64]:
+    def test_docstring(a: Field[[IDim], float64]):
         """My docstring."""
         fieldop_with_docstring(a, out=a)
 
@@ -1281,6 +1281,44 @@ def test_undefined_symbols():
             return undefined_symbol
 
 
+def test_zero_dims_fields(fieldview_backend):
+    if fieldview_backend == gtfn_cpu.run_gtfn:
+        pytest.skip("Implicit broadcast are not supported yet.")
+
+    inp = np_as_located_field()(np.array(1.0))
+    out = np_as_located_field()(np.array(0.0))
+
+    @field_operator(backend=fieldview_backend)
+    def implicit_broadcast_scalar(inp: Field[[], float]):
+        return inp
+
+    implicit_broadcast_scalar(inp, out=out, offset_provider={})
+    assert np.allclose(out, np.array(1.0))
+
+
+def test_implicit_broadcast_mixed_dims(fieldview_backend):
+    if fieldview_backend == gtfn_cpu.run_gtfn:
+        pytest.skip("Implicit broadcast are not supported yet.")
+
+    input1 = np_as_located_field(IDim)(np.ones((10,)))
+    inp = np_as_located_field()(np.array(1.0))
+    out = np_as_located_field(IDim)(np.ones((10,)))
+
+    @field_operator(backend=fieldview_backend)
+    def fieldop_implicit_broadcast(
+        zero_dim_inp: Field[[], float], inp: Field[[IDim], float], scalar: float
+    ) -> Field[[IDim], float]:
+        return inp + zero_dim_inp * scalar
+
+    @field_operator(backend=fieldview_backend)
+    def fieldop_implicit_broadcast_2(inp: Field[[IDim], float]) -> Field[[IDim], float]:
+        fi = fieldop_implicit_broadcast(1.0, inp, 1.0)
+        return fi
+
+    fieldop_implicit_broadcast_2(input1, out=out, offset_provider={})
+    assert np.allclose(out, np.asarray(inp) * 2)
+
+
 def test_tuple_unpacking(fieldview_backend):
     size = 10
     inp = np_as_located_field(IDim)(np.ones((size)))
diff --git a/tests/functional_tests/ffront_tests/test_foast_parsing.py b/tests/functional_tests/ffront_tests/test_foast_parsing.py
index 22518a14..fb23e246 100644
--- a/tests/functional_tests/ffront_tests/test_foast_parsing.py
+++ b/tests/functional_tests/ffront_tests/test_foast_parsing.py
@@ -324,3 +324,18 @@ def test_empty_dims_type():
         match=r"Annotated return type does not match deduced return type",
     ):
         _ = FieldOperatorParser.apply_to_function(empty_dims)
+
+
+def test_zero_dims_ternary():
+    ADim = Dimension("ADim")
+
+    def zero_dims_ternary(
+        cond: Field[[], float64], a: Field[[ADim], float64], b: Field[[ADim], float64]
+    ):
+        return a if cond == 1 else b
+
+    msg = r"Could not deduce type"
+    with pytest.raises(FieldOperatorTypeDeductionError) as exc_info:
+        _ = FieldOperatorParser.apply_to_function(zero_dims_ternary)
+
+    assert re.search(msg, exc_info.value.args[0]) is not None