diff --git a/pyproject.toml b/pyproject.toml index e51bac0d90..252a8e2858 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -310,6 +310,7 @@ markers = [ 'uses_max_over: tests that use the max_over builtin', 'uses_mesh_with_skip_values: tests that use a mesh with skip values', 'uses_concat_where: tests that use the concat_where builtin', + 'uses_concat_where_with_list_output: tests that use concat_where with a neighbor-list result', 'embedded_concat_where_infinite_domain: tests with concat_where resulting in an infinite domain', 'embedded_concat_where_non_contiguous_domain: tests with concat_where on non-contiguous domains', 'uses_program_metrics: tests that require backend support for program metrics', diff --git a/src/gt4py/next/ffront/fbuiltins.py b/src/gt4py/next/ffront/fbuiltins.py index 602221fbcc..dfd4ece6c7 100644 --- a/src/gt4py/next/ffront/fbuiltins.py +++ b/src/gt4py/next/ffront/fbuiltins.py @@ -243,6 +243,10 @@ def __call__( ) -> Tuple: ... def __call__(self, cond: CondT, true_field: FieldT1, false_field: FieldT2) -> _R: # type: ignore[misc] # supposedly this signature does not accept all the possible args allowed by the overloads ?? + if isinstance(true_field, named_collections.CUSTOM_NAMED_COLLECTION_TYPES): + return named_collections.tree_map_named_collection(lambda t, f: self(cond, t, f))( # type: ignore[return-value] # `NamedCollection` is not `_R` + true_field, false_field + ) if isinstance(true_field, tuple) or isinstance(false_field, tuple): if not (isinstance(true_field, tuple) and isinstance(false_field, tuple)): raise ValueError( diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index c9f51ad080..f8f739aa62 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -871,8 +871,12 @@ def _visit_math_built_in(self, node: foast.Call, **kwargs: Any) -> foast.Call: ) elif func_name in fbuiltins.BINARY_MATH_NUMBER_BUILTIN_NAMES: try: - return_type = type_info.promote( - *((cast(ts.FieldType | ts.ScalarType, arg.type)) for arg in node.args) + return_type = cast( + # a `ListType` only occurs at the ITIR level, never in the frontend + ts.FieldType | ts.ScalarType, + type_info.promote( + *((cast(ts.FieldType | ts.ScalarType, arg.type)) for arg in node.args) + ), ) except ValueError as ex: raise errors.DSLError(node.location, error_msg_preamble) from ex @@ -1020,15 +1024,17 @@ def _deduce_where_return_type( def deduce_return_type( tb: ts.FieldType | ts.ScalarType, fb: ts.FieldType | ts.ScalarType ) -> ts.FieldType: - if (t_dtype := type_info.extract_dtype(tb)) != (f_dtype := type_info.extract_dtype(fb)): + try: + promoted = type_info.promote(tb, fb) + except ValueError as ex: raise errors.DSLError( location, - f"Field arguments to '{func_name}' must be of same dtype, got '{t_dtype}' != " - f"'{f_dtype}'.", - ) - return_dims = promote_dims(cond_dims, type_info.extract_dims(type_info.promote(tb, fb))) - return_type = ts.FieldType(dims=return_dims, dtype=t_dtype) - return return_type + f"Could not promote '{tb}' and '{fb}' to common type in call to '{func_name}'.", + ) from ex + return ts.FieldType( + dims=promote_dims(cond_dims, type_info.extract_dims(promoted)), + dtype=type_info.extract_dtype(promoted), + ) return deduce_return_type(true_branch, false_branch) diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index 480a05812f..e840589f2a 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -415,8 +415,8 @@ def _visit_where(self, node: foast.Call, **kwargs: Any) -> itir.FunCall: # TODO(tehrengruber): For tuples we expand the tuple structure via `process_elements` # instead of emitting `tree_map_tuple` so mixed field types are supported, # e.g. (local field, regular field). - if not isinstance(node.type, ts.TupleType): # to keep the IR simpler - return self._lower_and_map("if_", *node.args) + if not isinstance(node.type, (ts.TupleType, ts.NamedCollectionType)): + return self._lower_and_map("if_", *node.args) # to keep the IR simpler cond_ = self.visit(node.args[0]) cond_symref_name = f"__cond_{itir.lenient_ir_fingerprinter(cond_)}" @@ -443,7 +443,20 @@ def _visit_concat_where(self, node: foast.Call, **kwargs: Any) -> itir.FunCall: # TODO(tehrengruber): Use `tree_map_tuple` when the domain inference is able to handle # lambda functions (with the results domain depending on the caller / args) domain, true_branch, false_branch = self.visit(node.args, **kwargs) - return im.concat_where(domain, true_branch, false_branch) + + def create_concat_where( + true_: itir.Expr, false_: itir.Expr, arg_types: tuple[ts.TypeSpec, ts.TypeSpec] + ) -> itir.FunCall: + if any(type_info.contains_local_field(t) for t in arg_types): + true_, false_ = (promote_to_list(t)(e) for t, e in zip(arg_types, (true_, false_))) + return im.concat_where(domain, true_, false_) + + branch_types = (node.args[1].type, node.args[2].type) + if not isinstance(node.type, (ts.TupleType, ts.NamedCollectionType)): + return create_concat_where(true_branch, false_branch, branch_types) # to keep the IR simpler + return lowering_utils.process_elements( + create_concat_where, (true_branch, false_branch), node.type, arg_types=branch_types + ) def _visit_broadcast(self, node: foast.Call, **kwargs: Any) -> itir.FunCall: return im.call("broadcast")(*self.visit(node.args, **kwargs)) diff --git a/src/gt4py/next/iterator/type_system/type_synthesizer.py b/src/gt4py/next/iterator/type_system/type_synthesizer.py index a3c98a96fb..f7de560aaa 100644 --- a/src/gt4py/next/iterator/type_system/type_synthesizer.py +++ b/src/gt4py/next/iterator/type_system/type_synthesizer.py @@ -216,6 +216,8 @@ def if_( # want this, but for roundtrip it is totally fine. # assert true_branch == false_branch # noqa: ERA001 + if isinstance(true_branch, ts.ListType) and isinstance(false_branch, ts.ListType): + return type_info.promote(true_branch, false_branch) return true_branch @@ -279,18 +281,11 @@ def deduce_return_type(tb: ts.FieldType | ts.ScalarType, fb: ts.FieldType | ts.S if any(isinstance(b, ts.DeferredType) for b in [tb, fb]): return ts.DeferredType(constraint=ts.FieldType) - tb_dtype, fb_dtype = (type_info.extract_dtype(b) for b in [tb, fb]) - - assert tb_dtype == fb_dtype, ( - f"Field arguments to 'concat_where' must be of same dtype, got '{tb_dtype}' != '{fb_dtype}'." - ) - dtype = tb_dtype - - return_dims = common.promote_dims( - domain.dims, type_info.extract_dims(type_info.promote(tb, fb)) + promoted = type_info.promote(tb, fb) + return ts.FieldType( + dims=common.promote_dims(domain.dims, type_info.extract_dims(promoted)), + dtype=type_info.extract_dtype(promoted), ) - return_type = ts.FieldType(dims=return_dims, dtype=dtype) - return return_type return deduce_return_type(true_field, false_field) diff --git a/src/gt4py/next/type_system/type_info.py b/src/gt4py/next/type_system/type_info.py index ac8467a5f6..8cdb5270a9 100644 --- a/src/gt4py/next/type_system/type_info.py +++ b/src/gt4py/next/type_system/type_info.py @@ -9,7 +9,7 @@ import functools import types from collections.abc import Callable, Collection, Iterable, Iterator -from typing import Any, Final, Literal, Sequence, Type, TypeGuard, TypeVar, cast, overload +from typing import Any, Final, Literal, Sequence, Type, TypeGuard, TypeVar, overload import numpy as np @@ -556,15 +556,18 @@ def is_concretizable(symbol_type: ts.TypeSpec, to_type: ts.TypeSpec) -> bool: def promote( - *types: ts.FieldType | ts.ScalarType, always_field: bool = False -) -> ts.FieldType | ts.ScalarType: + *types: ts.FieldType | ts.ScalarType | ts.ListType, always_field: bool = False +) -> ts.FieldType | ts.ScalarType | ts.ListType: """ - Promote a set of field or scalar types to a common type. + Promote a set of field, scalar or list types to a common type. The resulting type is defined on all dimensions of the arguments, respecting the individual order of the dimensions of each argument (see :func:`common.promote_dims` for more details). + `ListType`s promote only with each other: their element types must be equal, and an + `offset_type` of `None` (a list from `make_const_list`) is compatible with any other. + >>> dtype = ts.ScalarType(kind=ts.ScalarKind.INT64) >>> I, J, K = (common.Dimension(value=dim) for dim in ["I", "J", "K"]) >>> promoted: ts.FieldType = promote( @@ -578,18 +581,45 @@ def promote( ... ) >>> promoted.dims == [I, J, K] and promoted.dtype == dtype True + + >>> V2E = common.Dimension(value="V2E", kind=common.DimensionKind.LOCAL) + >>> list_dtype = ts.ListType(element_type=dtype, offset_type=V2E) + >>> promote( + ... ts.FieldType(dims=[I], dtype=list_dtype), + ... ts.FieldType(dims=[I, J], dtype=list_dtype), + ... ) == ts.FieldType(dims=[I, J], dtype=list_dtype) + True + + >>> const_list_dtype = ts.ListType(element_type=dtype, offset_type=None) + >>> promote( + ... ts.FieldType(dims=[I], dtype=const_list_dtype), + ... ts.FieldType(dims=[I], dtype=list_dtype), + ... ) == ts.FieldType(dims=[I], dtype=list_dtype) + True """ if not always_field and all(isinstance(type_, ts.ScalarType) for type_ in types): if not all(type_ == types[0] for type_ in types): raise ValueError("Could not promote scalars of different dtype (not implemented).") - if not all(type_.shape is None for type_ in types): # type: ignore[union-attr] + if isinstance(types[0], ts.ScalarType) and types[0].shape is not None: raise NotImplementedError("Shape promotion not implemented.") return types[0] + elif not always_field and any(isinstance(type_, ts.ListType) for type_ in types): + lists = [type_ for type_ in types if isinstance(type_, ts.ListType)] + if len(lists) != len(types): + raise ValueError("Could not promote lists together with non-lists.") + if not all(list_.element_type == lists[0].element_type for list_ in lists): + raise ValueError("Could not promote lists of different element type (not implemented).") + offset_types = [list_.offset_type for list_ in lists if list_.offset_type is not None] + if not all(offset_type == offset_types[0] for offset_type in offset_types): + raise ValueError("Could not promote lists over different offsets.") + return ts.ListType( + element_type=lists[0].element_type, + offset_type=offset_types[0] if offset_types else None, + ) elif all(isinstance(type_, (ts.ScalarType, ts.FieldType)) for type_ in types): dims = common.promote_dims(*(extract_dims(type_) for type_ in types)) - extracted_dtypes = [extract_dtype(type_) for type_ in types] - assert all(isinstance(dtype, ts.ScalarType) for dtype in extracted_dtypes) - dtype = cast(ts.ScalarType, promote(*extracted_dtypes)) # type: ignore[arg-type] # checked is `ScalarType` + dtype = promote(*(extract_dtype(type_) for type_ in types)) + assert isinstance(dtype, (ts.ScalarType, ts.ListType)) return ts.FieldType(dims=dims, dtype=dtype) raise TypeError("Expected a 'FieldType' or 'ScalarType'.") diff --git a/tests/next_tests/definitions.py b/tests/next_tests/definitions.py index e23d7d9319..8c2b71fddf 100644 --- a/tests/next_tests/definitions.py +++ b/tests/next_tests/definitions.py @@ -125,6 +125,7 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum): USES_PROGRAM_METRICS = "uses_program_metrics" USES_SCALAR_IN_DOMAIN_AND_FO = "uses_scalar_in_domain_and_fo" USES_CONCAT_WHERE = "uses_concat_where" +USES_CONCAT_WHERE_WITH_LIST_OUTPUT = "uses_concat_where_with_list_output" EMBEDDED_CONCAT_WHERE_INFINITE_DOMAIN = "embedded_concat_where_infinite_domain" EMBEDDED_CONCAT_WHERE_NON_CONTIGUOUS_DOMAIN = "embedded_concat_where_non_contiguous_domain" USES_PROGRAM_WITH_SLICED_OUT_ARGUMENTS = "uses_program_with_sliced_out_arguments" @@ -166,6 +167,7 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum): + [ (USES_CAN_DEREF, XFAIL, UNSUPPORTED_MESSAGE), (USES_COMPOSITE_SHIFTS, XFAIL, UNSUPPORTED_MESSAGE), + (USES_CONCAT_WHERE_WITH_LIST_OUTPUT, XFAIL, UNSUPPORTED_MESSAGE), (USES_LIFT, XFAIL, UNSUPPORTED_MESSAGE), (USES_REDUCE_WITH_LAMBDA, XFAIL, UNSUPPORTED_MESSAGE), (USES_SCAN_IN_STENCIL, XFAIL, BINDINGS_UNSUPPORTED_MESSAGE), diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py index 577f5a520e..b0e599848a 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_concat_where.py @@ -8,13 +8,23 @@ import numpy as np import pytest -from next_tests.integration_tests.cases import IDim, JDim, KDim, cartesian_case +from next_tests.integration_tests.cases import ( + IDim, + JDim, + KDim, + V2E, + V2EDim, + Vertex, + cartesian_case, + unstructured_case, +) from gt4py import next as gtx -from gt4py.next import broadcast +from gt4py.next import broadcast, common, neighbor_sum from gt4py.next.ffront.experimental import concat_where from next_tests.integration_tests import cases from next_tests.integration_tests.cases_utils import ( exec_alloc_descriptor, + mesh_descriptor, ) pytestmark = pytest.mark.uses_concat_where @@ -390,6 +400,177 @@ def ref(interior0, boundary0, interior1, boundary1): cases.verify_with_default_data(cartesian_case, testee, ref) +@pytest.mark.uses_tuple_returns +def test_with_nested_tuples(cartesian_case, static_domains: bool): + @gtx.field_operator(static_domains=static_domains) + def testee( + interior0: cases.IJKField, + boundary0: cases.IJField, + interior1: cases.IJKField, + boundary1: cases.IJField, + interior2: cases.IJKField, + boundary2: cases.IJField, + ) -> tuple[cases.IJKField, tuple[cases.IJKField, cases.IJKField]]: + return concat_where( + KDim == 0, + (boundary0, (boundary1, boundary2)), + (interior0, (interior1, interior2)), + ) + + interiors = tuple(cases.allocate(cartesian_case, testee, f"interior{i}")() for i in range(3)) + boundaries = tuple(cases.allocate(cartesian_case, testee, f"boundary{i}")() for i in range(3)) + out = cases.allocate(cartesian_case, testee, cases.RETURN)() + + k = np.arange(0, cartesian_case.default_sizes[KDim]) + refs = tuple( + np.where( + k[np.newaxis, np.newaxis, :] == 0, + boundary.asnumpy()[:, :, np.newaxis], + interior.asnumpy(), + ) + for boundary, interior in zip(boundaries, interiors) + ) + + cases.verify( + cartesian_case, + testee, + interiors[0], + boundaries[0], + interiors[1], + boundaries[1], + interiors[2], + boundaries[2], + out=out, + ref=(refs[0], (refs[1], refs[2])), + ) + + +@pytest.mark.uses_unstructured_shift +@pytest.mark.uses_concat_where_with_list_output +def test_with_local_field(unstructured_case, static_domains: bool): + @gtx.field_operator(static_domains=static_domains) + def testee(a: cases.EField, b: cases.EField) -> cases.VField: + t = concat_where(Vertex < 2, a(V2E), b(V2E)) + return neighbor_sum(t, axis=V2EDim) + + v2e_table = unstructured_case.offset_provider["V2E"].asnumpy() + vertex_mask = np.arange(unstructured_case.default_sizes[Vertex]) < 2 + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a, b: np.sum( + np.where(vertex_mask[:, np.newaxis], a[v2e_table], b[v2e_table]), + axis=1, + initial=0, + where=v2e_table != common._DEFAULT_SKIP_VALUE, + ), + ) + + +@pytest.mark.uses_tuple_returns +@pytest.mark.uses_unstructured_shift +@pytest.mark.uses_concat_where_with_list_output +def test_with_tuples_of_local_fields(unstructured_case, static_domains: bool): + @gtx.field_operator(static_domains=static_domains) + def testee( + a: cases.EField, + b: cases.EField, + c: cases.EField, + d: cases.EField, + ) -> tuple[cases.VField, cases.VField]: + t = concat_where(Vertex < 2, (a(V2E), b(V2E)), (c(V2E), d(V2E))) + return neighbor_sum(t[0], axis=V2EDim), neighbor_sum(t[1], axis=V2EDim) + + v2e_table = unstructured_case.offset_provider["V2E"].asnumpy() + vertex_mask = np.arange(unstructured_case.default_sizes[Vertex]) < 2 + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a, b, c, d: ( + np.sum( + np.where(vertex_mask[:, np.newaxis], a[v2e_table], c[v2e_table]), + axis=1, + initial=0, + where=v2e_table != common._DEFAULT_SKIP_VALUE, + ), + np.sum( + np.where(vertex_mask[:, np.newaxis], b[v2e_table], d[v2e_table]), + axis=1, + initial=0, + where=v2e_table != common._DEFAULT_SKIP_VALUE, + ), + ), + ) + + +@pytest.mark.uses_tuple_returns +@pytest.mark.uses_unstructured_shift +@pytest.mark.uses_concat_where_with_list_output +@pytest.mark.embedded_concat_where_infinite_domain +def test_with_local_field_and_scalar(unstructured_case, static_domains: bool): + @gtx.field_operator(static_domains=static_domains) + def testee(a: cases.EField) -> tuple[cases.VField, cases.VField]: + return ( + neighbor_sum(concat_where(Vertex < 2, a(V2E), 3), axis=V2EDim), + neighbor_sum(concat_where(Vertex < 2, 3, a(V2E)), axis=V2EDim), + ) + + v2e_table = unstructured_case.offset_provider["V2E"].asnumpy() + vertex_mask = np.arange(unstructured_case.default_sizes[Vertex]) < 2 + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a: ( + np.sum( + np.where(vertex_mask[:, np.newaxis], a[v2e_table], 3), + axis=1, + initial=0, + where=v2e_table != common._DEFAULT_SKIP_VALUE, + ), + np.sum( + np.where(vertex_mask[:, np.newaxis], 3, a[v2e_table]), + axis=1, + initial=0, + where=v2e_table != common._DEFAULT_SKIP_VALUE, + ), + ), + ) + + +@pytest.mark.uses_tuple_returns +@pytest.mark.uses_unstructured_shift +@pytest.mark.uses_concat_where_with_list_output +@pytest.mark.embedded_concat_where_infinite_domain +def test_with_tuples_of_local_and_nonlocal_leaves(unstructured_case, static_domains: bool): + @gtx.field_operator(static_domains=static_domains) + def testee( + a: cases.EField, b: cases.EField, c: cases.VField + ) -> tuple[cases.VField, cases.VField]: + t = concat_where(Vertex < 2, (a(V2E), c), (3, b(V2E))) + return neighbor_sum(t[0], axis=V2EDim), neighbor_sum(t[1], axis=V2EDim) + + v2e_table = unstructured_case.offset_provider["V2E"].asnumpy() + vertex_mask = np.arange(unstructured_case.default_sizes[Vertex]) < 2 + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a, b, c: ( + np.sum( + np.where(vertex_mask[:, np.newaxis], a[v2e_table], 3), + axis=1, + initial=0, + where=v2e_table != common._DEFAULT_SKIP_VALUE, + ), + np.sum( + np.where(vertex_mask[:, np.newaxis], c[:, np.newaxis], b[v2e_table]), + axis=1, + initial=0, + where=v2e_table != common._DEFAULT_SKIP_VALUE, + ), + ), + ) + + def test_nested_conditions_with_empty_branches(cartesian_case, static_domains: bool): @gtx.field_operator(static_domains=static_domains) def testee(interior: cases.IField, boundary: cases.IField, N: gtx.int32) -> cases.IField: diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_named_collections.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_named_collections.py index 0e3a749b3a..484c530913 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_named_collections.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_named_collections.py @@ -13,15 +13,26 @@ from typing import NamedTuple import gt4py.next as gtx +from gt4py.next import common, neighbor_sum from gt4py.next.ffront import decorator from gt4py.next.ffront.fbuiltins import where from gt4py.next.ffront.experimental import concat_where import dataclasses from next_tests.integration_tests import cases -from next_tests.integration_tests.cases import IDim, JDim, KDim, cartesian_case +from next_tests.integration_tests.cases import ( + IDim, + JDim, + KDim, + V2E, + V2EDim, + Vertex, + cartesian_case, + unstructured_case, +) from next_tests.integration_tests.cases_utils import ( exec_alloc_descriptor, + mesh_descriptor, ) # TODO(havogt): Since currently direct field_operator calls and program calls take different code paths the tests are duplicated. @@ -39,6 +50,11 @@ class NamedTupleNamedCollection(NamedTuple): v: gtx.Field[[IDim, JDim], gtx.float32] +class LocalFieldNamedCollection(NamedTuple): + neighbors: gtx.Field[[Vertex, V2EDim], np.int32] + center: gtx.Field[[Vertex], np.int32] + + @gtx.field_operator def constructed_outside_named_tuple( vel: NamedTupleNamedCollection, @@ -415,6 +431,88 @@ def testee( ) +@gtx.field_operator +def where_named_tuple( + i: cases.IField, interior: NamedTupleNamedCollection, boundary: NamedTupleNamedCollection +) -> NamedTupleNamedCollection: + return where(i == 0, boundary, interior) + + +@gtx.field_operator +def where_dataclass( + i: cases.IField, interior: DataclassNamedCollection, boundary: DataclassNamedCollection +) -> DataclassNamedCollection: + return where(i == 0, boundary, interior) + + +@pytest.mark.parametrize("testee", [where_named_tuple, where_dataclass]) +@pytest.mark.uses_tuple_returns +@pytest.mark.uses_tuple_args +def test_where(cartesian_case, testee): + i = cases.allocate(cartesian_case, testee, "i", strategy=cases.IndexInitializer())() + interior = cases.allocate(cartesian_case, testee, "interior")() + boundary = cases.allocate(cartesian_case, testee, "boundary")() + out = cases.allocate(cartesian_case, testee, cases.RETURN)() + + is_boundary = i.asnumpy()[:, np.newaxis] == 0 + cases.verify( + cartesian_case, + testee, + i, + interior, + boundary, + out=out, + ref=out.__class__( + u=np.where(is_boundary, boundary.u.asnumpy(), interior.u.asnumpy()), + v=np.where(is_boundary, boundary.v.asnumpy(), interior.v.asnumpy()), + ), + ) + + +@pytest.mark.uses_tuple_returns +@pytest.mark.uses_unstructured_shift +def test_where_with_local_fields(unstructured_case): + @gtx.field_operator + def testee( + mask: cases.VBoolField, a: cases.EField, b: cases.EField, c: cases.VField, d: cases.VField + ) -> tuple[cases.VField, cases.VField]: + t = where( + mask, + LocalFieldNamedCollection(neighbors=a(V2E), center=c), + LocalFieldNamedCollection(neighbors=b(V2E), center=d), + ) + return neighbor_sum(t.neighbors, axis=V2EDim), t.center + + v2e_table = unstructured_case.offset_provider["V2E"].asnumpy() + mask = unstructured_case.as_field( + [Vertex], np.random.choice(a=[False, True], size=unstructured_case.default_sizes[Vertex]) + ) + a, b, c, d = (cases.allocate(unstructured_case, testee, name)() for name in "abcd") + out = cases.allocate(unstructured_case, testee, cases.RETURN)() + + cases.verify( + unstructured_case, + testee, + mask, + a, + b, + c, + d, + out=out, + ref=( + np.sum( + np.where( + mask.asnumpy()[:, np.newaxis], a.asnumpy()[v2e_table], b.asnumpy()[v2e_table] + ), + axis=1, + initial=0, + where=v2e_table != common._DEFAULT_SKIP_VALUE, + ), + np.where(mask.asnumpy(), c.asnumpy(), d.asnumpy()), + ), + ) + + @pytest.mark.uses_tuple_returns @pytest.mark.uses_tuple_args def test_where_nested(cartesian_case): @@ -526,3 +624,36 @@ def testee( out=out, ref=refs, ) + + +@pytest.mark.uses_tuple_returns +@pytest.mark.uses_concat_where +@pytest.mark.uses_unstructured_shift +@pytest.mark.uses_concat_where_with_list_output +def test_concat_where_with_local_fields(unstructured_case): + @gtx.field_operator + def testee( + a: cases.EField, b: cases.EField, c: cases.VField, d: cases.VField + ) -> tuple[cases.VField, cases.VField]: + t = concat_where( + Vertex < 2, + LocalFieldNamedCollection(neighbors=a(V2E), center=c), + LocalFieldNamedCollection(neighbors=b(V2E), center=d), + ) + return neighbor_sum(t.neighbors, axis=V2EDim), t.center + + v2e_table = unstructured_case.offset_provider["V2E"].asnumpy() + vertex_mask = np.arange(unstructured_case.default_sizes[Vertex]) < 2 + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a, b, c, d: ( + np.sum( + np.where(vertex_mask[:, np.newaxis], a[v2e_table], b[v2e_table]), + axis=1, + initial=0, + where=v2e_table != common._DEFAULT_SKIP_VALUE, + ), + np.where(vertex_mask, c, d), + ), + ) diff --git a/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py b/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py index a75e7bb031..d72efaaa22 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_func_to_foast.py @@ -286,7 +286,10 @@ def conditional_wrong_arg_type( with pytest.raises( errors.DSLError, - match="Field arguments to 'where' must be of same dtype, got 'float32' != 'float64'.", + match=re.escape( + "Could not promote 'Field[[TDim], float32]' and 'Field[[TDim], float64]' " + "to common type in call to 'where'." + ), ): _ = FieldOperatorParser.apply_to_function(conditional_wrong_arg_type) diff --git a/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py b/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py index 22bd1a7a9e..b6ae674312 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_type_deduction.py @@ -239,7 +239,7 @@ def domain_comparison(a: Field[[TDim], float], b: Field[[TDim], float]): with pytest.raises( errors.DSLError, - match="Field arguments to 'concat_where' must be of same dtype, got 'float64' != 'int32'.", + match="Could not promote 'float64' and 'int32' to common type in call to 'concat_where'.", ): _ = FieldOperatorParser.apply_to_function(domain_comparison) diff --git a/tests/next_tests/unit_tests/type_system_tests/test_type_info.py b/tests/next_tests/unit_tests/type_system_tests/test_type_info.py index 35c3d2eba1..dbe8e58ba8 100644 --- a/tests/next_tests/unit_tests/type_system_tests/test_type_info.py +++ b/tests/next_tests/unit_tests/type_system_tests/test_type_info.py @@ -462,3 +462,17 @@ def test_return_type( ) def test_needs_value_extraction(type_spec: ts.TypeSpec, expected: bool): assert type_info.needs_value_extraction(type_spec) is expected + + +def test_promote_lists(): + float64 = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + V2EDim = Dimension("V2E", kind=DimensionKind.LOCAL) + C2EDim = Dimension("C2E", kind=DimensionKind.LOCAL) + const_list = ts.ListType(element_type=float64, offset_type=None) + v2e_list = ts.ListType(element_type=float64, offset_type=V2EDim) + + assert type_info.promote(const_list, v2e_list) == v2e_list + with pytest.raises(ValueError, match="different offsets"): + type_info.promote(v2e_list, ts.ListType(element_type=float64, offset_type=C2EDim)) + with pytest.raises(ValueError, match="non-lists"): + type_info.promote(v2e_list, float64)