Repository: GridTools/gt4py
Pull request #1033 (2022-11-07 10:10:42)

Functional: implicit broadcast 0d field
Use scalar as input and perform implicit broadcasting to return field of empty dimensions
```python
@field_operator
def fieldop_implicit_broadcast(
    zero_dim_inp: Field[[], float], inp: Field[[IDim], float], scalar: float
) -> Field[[IDim], float]:
    return inp + zero_dim_inp * scalar

fieldop_implicit_broadcast(1.0, inp, 1.0)
```





Patch:
diff --git a/setup.cfg b/setup.cfg
index cca97301..5028dc43 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_passes.type_deduction]
-ignore_errors = True
-
 [mypy-functional.ffront.foast_to_itir]
 ignore_errors = True
 
diff --git a/src/functional/ffront/common_types.py b/src/functional/ffront/common_types.py
index 3a9f707b..505ea1b5 100644
--- a/src/functional/ffront/common_types.py
+++ b/src/functional/ffront/common_types.py
@@ -165,7 +165,7 @@ class FieldType(DataType, CallableType):
 class FunctionType(SymbolType, CallableType):
     args: list[DataType | DeferredSymbolType]
     kwargs: dict[str, DataType | DeferredSymbolType]
-    returns: DataType | DeferredSymbolType | VoidType
+    returns: DataType | DeferredSymbolType | VoidType | SymbolType
 
     def __str__(self):
         arg_strs = [str(arg) for arg in self.args]
diff --git a/src/functional/ffront/foast_passes/type_deduction.py b/src/functional/ffront/foast_passes/type_deduction.py
index 5cff76de..c5481a4b 100644
--- a/src/functional/ffront/foast_passes/type_deduction.py
+++ b/src/functional/ffront/foast_passes/type_deduction.py
@@ -17,7 +17,6 @@ import functional.ffront.field_operator_ast as foast
 from eve import NodeTranslator, NodeVisitor, traits
 from functional.common import DimensionKind, GTSyntaxError, GTTypeError
 from functional.ffront import common_types as ct, fbuiltins, type_info
-from functional.ffront.common_types import TupleType
 from functional.ffront.foast_passes.utils import compute_assign_indices
 from functional.ffront.symbol_makers import make_symbol_type_from_value
 
@@ -50,8 +49,8 @@ def boolified_type(symbol_type: ct.SymbolType) -> ct.ScalarType | ct.FieldType:
 
 
 def construct_tuple_type(
-    true_branch_types: list[ct.TupleType],
-    false_branch_types: list[ct.TupleType],
+    true_branch_types: list,
+    false_branch_types: list,
     mask_type: ct.FieldType,
 ) -> list:
     """
@@ -80,7 +79,7 @@ def construct_tuple_type(
 
 
 def promote_to_mask_type(
-    mask_type: ct.FieldType, input_type: ct.FieldType | ct.ScalarType | ct.TupleType
+    mask_type: ct.FieldType, input_type: ct.FieldType | ct.ScalarType
 ) -> ct.FieldType:
     """
     Promote mask type with the input type.
@@ -106,7 +105,7 @@ def promote_to_mask_type(
         item in input_type.dims for item in mask_type.dims
     ):
         return_dtype = input_type.dtype if isinstance(input_type, ct.FieldType) else input_type
-        return type_info.promote(input_type, ct.FieldType(dims=mask_type.dims, dtype=return_dtype))
+        return type_info.promote(input_type, ct.FieldType(dims=mask_type.dims, dtype=return_dtype))  # type: ignore
     else:
         return input_type
 
@@ -261,10 +260,10 @@ class FieldOperatorTypeDeduction(traits.VisitorWithSymbolTableTrait, NodeTransla
         TargetType = list[foast.Starred | foast.Symbol]
         values = self.visit(node.value, **kwargs)
 
-        if isinstance(values.type, TupleType):
+        if isinstance(values.type, ct.TupleType):
             num_elts: int = len(values.type.types)
             targets: TargetType = node.targets
-            indices: tuple[tuple[int] | int] = compute_assign_indices(targets, num_elts)
+            indices: list[tuple[int, int] | int] = compute_assign_indices(targets, num_elts)
 
             if not any(isinstance(i, tuple) for i in indices) and len(indices) != num_elts:
                 raise FieldOperatorTypeDeductionError.from_foast_node(
@@ -272,15 +271,17 @@ class FieldOperatorTypeDeduction(traits.VisitorWithSymbolTableTrait, NodeTransla
                 )
 
             new_targets: TargetType = []
+            new_type: ct.TupleType | ct.DataType
             for i, index in enumerate(indices):
                 old_target = targets[i]
 
                 if isinstance(index, tuple):
                     lower, upper = index
-                    new_type = TupleType(types=[t for t in values.type.types[lower:upper]])
+                    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
+                            id=old_target.id.id, location=old_target.location, type=new_type  # type: ignore
                         ),
                         type=new_type,
                         location=old_target.location,
@@ -658,8 +659,9 @@ class FieldOperatorTypeDeduction(traits.VisitorWithSymbolTableTrait, NodeTransla
 
     def _visit_where(self, node: foast.Call, **kwargs) -> foast.Call:
         mask_type = cast(ct.FieldType, node.args[0].type)
-        true_branch_type = cast(ct.FieldType, node.args[1].type)
-        false_branch_type = cast(ct.FieldType, node.args[2].type)
+        true_branch_type = node.args[1].type
+        false_branch_type = node.args[2].type
+        return_type: ct.TupleType | ct.FieldType
         if not type_info.is_logical(mask_type):
             raise FieldOperatorTypeDeductionError.from_foast_node(
                 node,
@@ -671,10 +673,10 @@ class FieldOperatorTypeDeduction(traits.VisitorWithSymbolTableTrait, NodeTransla
             if isinstance(true_branch_type, ct.TupleType) and isinstance(
                 false_branch_type, ct.TupleType
             ):
-                true_branch_types = node.args[1].type.types
-                false_branch_types = node.args[2].type.types
                 return_type = ct.TupleType(
-                    types=construct_tuple_type(true_branch_types, false_branch_types, mask_type)
+                    types=construct_tuple_type(
+                        true_branch_type.types, false_branch_type.types, mask_type
+                    )
                 )
             elif isinstance(true_branch_type, ct.TupleType) or isinstance(
                 false_branch_type, ct.TupleType
@@ -685,7 +687,9 @@ class FieldOperatorTypeDeduction(traits.VisitorWithSymbolTableTrait, NodeTransla
                     f"{node.args[1].type} and {node.args[2].type}",
                 )
             else:
-                promoted_type = type_info.promote(true_branch_type, false_branch_type)
+                true_branch_fieldtype = cast(ct.FieldType, true_branch_type)
+                false_branch_fieldtype = cast(ct.FieldType, false_branch_type)
+                promoted_type = type_info.promote(true_branch_fieldtype, false_branch_fieldtype)
                 return_type = promote_to_mask_type(mask_type, promoted_type)
 
         except GTTypeError as ex:
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