From e5b676acbb663b05bad412d8db90a4a46bc977f2 Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Thu, 27 Aug 2026 09:48:26 +0000 Subject: [PATCH 01/12] fix[next]: 'concat_where' crashes on local (neighbor-list) fields The GTIR 'concat_where' type synthesizer derives the result dims via 'type_info.promote(tb, fb)', which asserted that the promoted dtypes are 'ScalarType'. A local field carries a 'ListType' dtype at the GTIR level, so any 'concat_where' with a local-field branch failed with a bare 'AssertionError' during type inference -- and it did so even though both branches have the identical dtype, since the assertion is on the dtype's kind, not on the operands differing. Let 'promote' handle 'ListType' the same way it handles 'ScalarType': both promote only between equal types, so the two cases collapse into a single check. The docstring records that a 'ListType' only ever reaches 'promote' from the ITIR level, because the frontend represents the same concept as a field with a local dimension in 'dims' and a scalar dtype. The combination was previously untested ('test_concat_where.py' had no local-dimension coverage), so the latent assert never fired in CI. --- .../ffront/foast_passes/type_deduction.py | 8 ++- src/gt4py/next/type_system/type_info.py | 34 ++++++--- .../ffront_tests/test_concat_where.py | 72 ++++++++++++++++++- 3 files changed, 100 insertions(+), 14 deletions(-) diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index c9f51ad080..6d25c3edb3 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 diff --git a/src/gt4py/next/type_system/type_info.py b/src/gt4py/next/type_system/type_info.py index ac8467a5f6..1dd404fde3 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,20 @@ 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). + A `ListType` is the dtype of a local (neighbor-list) field and only ever reaches this + function from the ITIR level. The frontend represents the same concept as a field with + a local dimension in `dims` and a scalar dtype, so it never passes a list here. Lists + promote exactly like scalars, i.e. only between equal types. + >>> 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 +583,27 @@ 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 """ - if not always_field and all(isinstance(type_, ts.ScalarType) for type_ in types): + # Lists are only reached from the ITIR level (see above) and behave like scalars here: + # both promote only between equal types. + if not always_field and all(isinstance(type_, (ts.ScalarType, ts.ListType)) 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] + raise ValueError("Could not promote dtypes of different type (not implemented).") + if isinstance(types[0], ts.ScalarType) and types[0].shape is not None: raise NotImplementedError("Shape promotion not implemented.") return types[0] 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/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..f46c166e98 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 ( + E2V, + E2VDim, + Edge, + IDim, + JDim, + KDim, + 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,64 @@ def ref(interior0, boundary0, interior1, boundary1): cases.verify_with_default_data(cartesian_case, testee, ref) +@pytest.mark.uses_unstructured_shift +@pytest.mark.uses_sparse_fields +def test_with_local_field(unstructured_case, static_domains: bool): + @gtx.field_operator(static_domains=static_domains) + def testee(a: cases.VField, b: cases.VField) -> cases.EField: + t = concat_where(Edge < 2, a(E2V), b(E2V)) + return neighbor_sum(t, axis=E2VDim) + + e2v_table = unstructured_case.offset_provider["E2V"].asnumpy() + edge_mask = np.arange(unstructured_case.default_sizes[Edge]) < 2 + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a, b: np.sum( + np.where(edge_mask[:, np.newaxis], a[e2v_table], b[e2v_table]), + axis=1, + initial=0, + where=e2v_table != common._DEFAULT_SKIP_VALUE, + ), + ) + + +@pytest.mark.uses_tuple_returns +@pytest.mark.uses_unstructured_shift +@pytest.mark.uses_sparse_fields +def test_with_tuples_of_local_fields(unstructured_case, static_domains: bool): + @gtx.field_operator(static_domains=static_domains) + def testee( + a: cases.VField, + b: cases.VField, + c: cases.VField, + d: cases.VField, + ) -> tuple[cases.EField, cases.EField]: + t = concat_where(Edge < 2, (a(E2V), b(E2V)), (c(E2V), d(E2V))) + return neighbor_sum(t[0], axis=E2VDim), neighbor_sum(t[1], axis=E2VDim) + + e2v_table = unstructured_case.offset_provider["E2V"].asnumpy() + edge_mask = np.arange(unstructured_case.default_sizes[Edge]) < 2 + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a, b, c, d: ( + np.sum( + np.where(edge_mask[:, np.newaxis], a[e2v_table], c[e2v_table]), + axis=1, + initial=0, + where=e2v_table != common._DEFAULT_SKIP_VALUE, + ), + np.sum( + np.where(edge_mask[:, np.newaxis], b[e2v_table], d[e2v_table]), + axis=1, + initial=0, + where=e2v_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: From 8643e6d67e8d3b63198bf7a0a2c082724be99d90 Mon Sep 17 00:00:00 2001 From: Till Ehrengruber Date: Thu, 27 Aug 2026 19:58:03 +0000 Subject: [PATCH 02/12] Add 'concat_where' test with nested tuple branches (moved from #2833) --- .../ffront_tests/test_concat_where.py | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) 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 f46c166e98..6ed7b244c4 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 @@ -400,6 +400,51 @@ 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_sparse_fields def test_with_local_field(unstructured_case, static_domains: bool): From 06218ce44ee6d7528983acc8803fa4673e9b1b27 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Fri, 11 Sep 2026 13:07:41 +0200 Subject: [PATCH 03/12] fix[next]: promote list dtypes and use promote's result in concat_where `type_info.promote` accepts `ListType`s only together with other lists: element types must be equal, and an `offset_type` of `None` (a list from `make_const_list`) is compatible with any other offset, matching `map_list`. The ITIR `concat_where` type synthesizer and the frontend `where`/`concat_where` deduction now take both dims and dtype from `promote`'s result instead of checking dtype equality by hand and calling `promote` only for its dims. The frontend re-raises `promote`'s `ValueError` as a located `DSLError`, as the binary-operator deduction does. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh --- .../ffront/foast_passes/type_deduction.py | 16 ++++++---- .../iterator/type_system/type_synthesizer.py | 15 +++------ src/gt4py/next/type_system/type_info.py | 32 ++++++++++++++----- .../ffront_tests/test_func_to_foast.py | 5 ++- .../ffront_tests/test_type_deduction.py | 2 +- .../type_system_tests/test_type_info.py | 14 ++++++++ 6 files changed, 56 insertions(+), 28 deletions(-) diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index 6d25c3edb3..f8f739aa62 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -1024,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/iterator/type_system/type_synthesizer.py b/src/gt4py/next/iterator/type_system/type_synthesizer.py index a3c98a96fb..3f838181f4 100644 --- a/src/gt4py/next/iterator/type_system/type_synthesizer.py +++ b/src/gt4py/next/iterator/type_system/type_synthesizer.py @@ -279,18 +279,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 1dd404fde3..8cdb5270a9 100644 --- a/src/gt4py/next/type_system/type_info.py +++ b/src/gt4py/next/type_system/type_info.py @@ -565,10 +565,8 @@ def promote( the individual order of the dimensions of each argument (see :func:`common.promote_dims` for more details). - A `ListType` is the dtype of a local (neighbor-list) field and only ever reaches this - function from the ITIR level. The frontend represents the same concept as a field with - a local dimension in `dims` and a scalar dtype, so it never passes a list here. Lists - promote exactly like scalars, i.e. only between equal types. + `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"]) @@ -591,15 +589,33 @@ def promote( ... 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 """ - # Lists are only reached from the ITIR level (see above) and behave like scalars here: - # both promote only between equal types. - if not always_field and all(isinstance(type_, (ts.ScalarType, ts.ListType)) for type_ in types): + 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 dtypes of different type (not implemented).") + raise ValueError("Could not promote scalars of different dtype (not implemented).") 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)) dtype = promote(*(extract_dtype(type_) for type_ in types)) 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) From 8ead66c451c593badd04fd6efb74c051e69f14f4 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Fri, 11 Sep 2026 14:32:21 +0200 Subject: [PATCH 04/12] fix[next]: concat_where with a local field and a scalar branch The lowering wraps the non-local branch of a `concat_where` whose other branch contains a local field in `make_const_list`, as `where` already does, so a scalar branch reaches the type synthesizer as a list. `transform_to_as_fieldop` turns `concat_where` into an `if_` selecting between whole lists, and `if_` typed its result as its true branch, which drops the neighbor offset when the true branch is a constant list. For two list branches `if_` now returns their promotion. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh --- src/gt4py/next/ffront/foast_to_gtir.py | 6 +++++ .../iterator/type_system/type_synthesizer.py | 2 ++ .../ffront_tests/test_concat_where.py | 24 +++++++++++++++++++ 3 files changed, 32 insertions(+) diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index 480a05812f..01bd12d4f9 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -443,6 +443,12 @@ 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) + true_type, false_type = node.args[1].type, node.args[2].type + if not isinstance(node.type, ts.TupleType) and any( + type_info.contains_local_field(t) for t in (true_type, false_type) + ): + true_branch = promote_to_list(true_type)(true_branch) + false_branch = promote_to_list(false_type)(false_branch) return im.concat_where(domain, true_branch, false_branch) def _visit_broadcast(self, node: foast.Call, **kwargs: Any) -> itir.FunCall: diff --git a/src/gt4py/next/iterator/type_system/type_synthesizer.py b/src/gt4py/next/iterator/type_system/type_synthesizer.py index 3f838181f4..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 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 6ed7b244c4..15b0c60f10 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 @@ -503,6 +503,30 @@ def testee( ) +@pytest.mark.uses_tuple_returns +@pytest.mark.uses_unstructured_shift +@pytest.mark.uses_sparse_fields +@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.VField) -> tuple[cases.EField, cases.EField]: + return ( + neighbor_sum(concat_where(Edge < 2, a(E2V), 3), axis=E2VDim), + neighbor_sum(concat_where(Edge < 2, 3, a(E2V)), axis=E2VDim), + ) + + e2v_table = unstructured_case.offset_provider["E2V"].asnumpy() + edge_mask = np.arange(unstructured_case.default_sizes[Edge]) < 2 + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a: ( + np.sum(np.where(edge_mask[:, np.newaxis], a[e2v_table], 3), axis=1), + np.sum(np.where(edge_mask[:, np.newaxis], 3, a[e2v_table]), axis=1), + ), + ) + + 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: From 3f3ac05d1e7aa48dd2f8e9d3d991f3526754c4dc Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Fri, 11 Sep 2026 14:58:13 +0200 Subject: [PATCH 05/12] test[next]: mark concat_where list-output tests for dace The local-field `concat_where` tests borrowed `uses_sparse_fields` to be xfailed on dace, although they pass no sparse field. What dace lacks is `concat_where` with a list result, so they get a dedicated marker for that. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh --- pyproject.toml | 1 + tests/next_tests/definitions.py | 2 ++ .../feature_tests/ffront_tests/test_concat_where.py | 6 +++--- 3 files changed, 6 insertions(+), 3 deletions(-) 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/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 15b0c60f10..88127a7f52 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 @@ -446,7 +446,7 @@ def testee( @pytest.mark.uses_unstructured_shift -@pytest.mark.uses_sparse_fields +@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.VField, b: cases.VField) -> cases.EField: @@ -469,7 +469,7 @@ def testee(a: cases.VField, b: cases.VField) -> cases.EField: @pytest.mark.uses_tuple_returns @pytest.mark.uses_unstructured_shift -@pytest.mark.uses_sparse_fields +@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( @@ -505,7 +505,7 @@ def testee( @pytest.mark.uses_tuple_returns @pytest.mark.uses_unstructured_shift -@pytest.mark.uses_sparse_fields +@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) From 24ee9491642e7858c7664a62a783914d21e549be Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Fri, 11 Sep 2026 15:22:28 +0200 Subject: [PATCH 06/12] test[next]: cover skip values in the local-field concat_where tests The local-field tests use `V2E` instead of `E2V`, which has skip values in `skip_value_mesh`, so the skip-value mask in their references takes effect. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh --- .../ffront_tests/test_concat_where.py | 72 +++++++++++-------- 1 file changed, 41 insertions(+), 31 deletions(-) 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 88127a7f52..0692ec7bc0 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 @@ -9,12 +9,12 @@ import numpy as np import pytest from next_tests.integration_tests.cases import ( - E2V, - E2VDim, - Edge, IDim, JDim, KDim, + V2E, + V2EDim, + Vertex, cartesian_case, unstructured_case, ) @@ -449,20 +449,20 @@ def testee( @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.VField, b: cases.VField) -> cases.EField: - t = concat_where(Edge < 2, a(E2V), b(E2V)) - return neighbor_sum(t, axis=E2VDim) + 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) - e2v_table = unstructured_case.offset_provider["E2V"].asnumpy() - edge_mask = np.arange(unstructured_case.default_sizes[Edge]) < 2 + 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(edge_mask[:, np.newaxis], a[e2v_table], b[e2v_table]), + np.where(vertex_mask[:, np.newaxis], a[v2e_table], b[v2e_table]), axis=1, initial=0, - where=e2v_table != common._DEFAULT_SKIP_VALUE, + where=v2e_table != common._DEFAULT_SKIP_VALUE, ), ) @@ -473,31 +473,31 @@ def testee(a: cases.VField, b: cases.VField) -> cases.EField: def test_with_tuples_of_local_fields(unstructured_case, static_domains: bool): @gtx.field_operator(static_domains=static_domains) def testee( - a: cases.VField, - b: cases.VField, - c: cases.VField, - d: cases.VField, - ) -> tuple[cases.EField, cases.EField]: - t = concat_where(Edge < 2, (a(E2V), b(E2V)), (c(E2V), d(E2V))) - return neighbor_sum(t[0], axis=E2VDim), neighbor_sum(t[1], axis=E2VDim) - - e2v_table = unstructured_case.offset_provider["E2V"].asnumpy() - edge_mask = np.arange(unstructured_case.default_sizes[Edge]) < 2 + 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(edge_mask[:, np.newaxis], a[e2v_table], c[e2v_table]), + np.where(vertex_mask[:, np.newaxis], a[v2e_table], c[v2e_table]), axis=1, initial=0, - where=e2v_table != common._DEFAULT_SKIP_VALUE, + where=v2e_table != common._DEFAULT_SKIP_VALUE, ), np.sum( - np.where(edge_mask[:, np.newaxis], b[e2v_table], d[e2v_table]), + np.where(vertex_mask[:, np.newaxis], b[v2e_table], d[v2e_table]), axis=1, initial=0, - where=e2v_table != common._DEFAULT_SKIP_VALUE, + where=v2e_table != common._DEFAULT_SKIP_VALUE, ), ), ) @@ -509,20 +509,30 @@ def testee( @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.VField) -> tuple[cases.EField, cases.EField]: + def testee(a: cases.EField) -> tuple[cases.VField, cases.VField]: return ( - neighbor_sum(concat_where(Edge < 2, a(E2V), 3), axis=E2VDim), - neighbor_sum(concat_where(Edge < 2, 3, a(E2V)), axis=E2VDim), + neighbor_sum(concat_where(Vertex < 2, a(V2E), 3), axis=V2EDim), + neighbor_sum(concat_where(Vertex < 2, 3, a(V2E)), axis=V2EDim), ) - e2v_table = unstructured_case.offset_provider["E2V"].asnumpy() - edge_mask = np.arange(unstructured_case.default_sizes[Edge]) < 2 + 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(edge_mask[:, np.newaxis], a[e2v_table], 3), axis=1), - np.sum(np.where(edge_mask[:, np.newaxis], 3, a[e2v_table]), axis=1), + 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, + ), ), ) From ad788d947b5742802e53d01a4306792bb6ee9597 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Fri, 11 Sep 2026 16:19:32 +0200 Subject: [PATCH 07/12] fix[next]: promote tuple leaves of concat_where to lists In a tuple-valued concat_where, each leaf whose counterpart in the other branch is a local field is wrapped in make_const_list, as for non-tuple branches. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh --- src/gt4py/next/ffront/foast_to_gtir.py | 23 ++++++++++--- .../ffront_tests/test_concat_where.py | 34 +++++++++++++++++++ 2 files changed, 52 insertions(+), 5 deletions(-) diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index 01bd12d4f9..d37648d7e5 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -444,11 +444,24 @@ def _visit_concat_where(self, node: foast.Call, **kwargs: Any) -> itir.FunCall: # lambda functions (with the results domain depending on the caller / args) domain, true_branch, false_branch = self.visit(node.args, **kwargs) true_type, false_type = node.args[1].type, node.args[2].type - if not isinstance(node.type, ts.TupleType) and any( - type_info.contains_local_field(t) for t in (true_type, false_type) - ): - true_branch = promote_to_list(true_type)(true_branch) - false_branch = promote_to_list(false_type)(false_branch) + + def promote_leaf(expr: itir.Expr, arg_types: tuple[ts.TypeSpec, ts.TypeSpec]) -> itir.Expr: + own_type, other_type = arg_types + if type_info.contains_local_field(other_type): + return promote_to_list(own_type)(expr) + return expr + + if type_info.contains_local_field(node.type): + if isinstance(node.type, (ts.TupleType, ts.NamedCollectionType)): + true_branch = lowering_utils.process_elements( + promote_leaf, true_branch, node.type, arg_types=(true_type, false_type) + ) + false_branch = lowering_utils.process_elements( + promote_leaf, false_branch, node.type, arg_types=(false_type, true_type) + ) + else: + true_branch = promote_leaf(true_branch, (true_type, false_type)) + false_branch = promote_leaf(false_branch, (false_type, true_type)) return im.concat_where(domain, true_branch, false_branch) def _visit_broadcast(self, node: foast.Call, **kwargs: Any) -> itir.FunCall: 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 0692ec7bc0..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 @@ -537,6 +537,40 @@ def testee(a: cases.EField) -> tuple[cases.VField, cases.VField]: ) +@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: From 39bf5cd778987d4b3c7a35676aa6c6e7d15dea07 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Fri, 11 Sep 2026 16:46:13 +0200 Subject: [PATCH 08/12] fix[next]: lower where and concat_where per element for named collections concat_where builds one concat_where per tuple element through process_elements, as where does, promoting a non-local element to a list next to a local one. Both take the per-element path for named collections too, so a named collection holding a local field is no longer mapped or promoted as a whole. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh --- src/gt4py/next/ffront/foast_to_gtir.py | 38 +++--- .../ffront_tests/test_named_collections.py | 122 +++++++++++++++++- 2 files changed, 137 insertions(+), 23 deletions(-) diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index d37648d7e5..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,26 +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) - true_type, false_type = node.args[1].type, node.args[2].type - - def promote_leaf(expr: itir.Expr, arg_types: tuple[ts.TypeSpec, ts.TypeSpec]) -> itir.Expr: - own_type, other_type = arg_types - if type_info.contains_local_field(other_type): - return promote_to_list(own_type)(expr) - return expr - - if type_info.contains_local_field(node.type): - if isinstance(node.type, (ts.TupleType, ts.NamedCollectionType)): - true_branch = lowering_utils.process_elements( - promote_leaf, true_branch, node.type, arg_types=(true_type, false_type) - ) - false_branch = lowering_utils.process_elements( - promote_leaf, false_branch, node.type, arg_types=(false_type, true_type) - ) - else: - true_branch = promote_leaf(true_branch, (true_type, false_type)) - false_branch = promote_leaf(false_branch, (false_type, true_type)) - 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/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..a597348c68 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,77 @@ def testee( ) +@pytest.mark.uses_tuple_returns +@pytest.mark.uses_tuple_args +def test_where(cartesian_case): + @gtx.field_operator + def testee( + i: cases.IField, + interior: NamedTupleNamedCollection, + boundary: NamedTupleNamedCollection, + ) -> NamedTupleNamedCollection: + return where(i == 0, boundary, interior) + + 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)() + + ref = tuple( + np.where( + i.asnumpy()[:, np.newaxis] == 0, + tuple(vel.asnumpy() for vel in (boundary.u, boundary.v)), + tuple(vel.asnumpy() for vel in (interior.u, interior.v)), + ) + ) + + cases.verify(cartesian_case, testee, i, interior, boundary, out=out, ref=ref) + + +@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 +613,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), + ), + ) From 95de2bfaa1c180226a99fa3cb62e4cc5966da76b Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Fri, 11 Sep 2026 17:03:52 +0200 Subject: [PATCH 09/12] fix[next]: keep the named-collection type in embedded where and concat_where Embedded where and concat_where rebuilt a NamedTuple as a plain tuple and passed a dataclass collection on to the field dispatch. Both now extract the collection, select per element and construct the same collection type again. Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh --- src/gt4py/next/ffront/fbuiltins.py | 4 ++ .../ffront_tests/test_named_collections.py | 45 ++++++++++++------- 2 files changed, 32 insertions(+), 17 deletions(-) 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/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 a597348c68..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 @@ -431,32 +431,43 @@ 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): - @gtx.field_operator - def testee( - i: cases.IField, - interior: NamedTupleNamedCollection, - boundary: NamedTupleNamedCollection, - ) -> NamedTupleNamedCollection: - return where(i == 0, boundary, interior) - +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)() - ref = tuple( - np.where( - i.asnumpy()[:, np.newaxis] == 0, - tuple(vel.asnumpy() for vel in (boundary.u, boundary.v)), - tuple(vel.asnumpy() for vel in (interior.u, interior.v)), - ) + 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()), + ), ) - cases.verify(cartesian_case, testee, i, interior, boundary, out=out, ref=ref) - @pytest.mark.uses_tuple_returns @pytest.mark.uses_unstructured_shift From 6685fc425ccdff7420cb63db77c7dcc2c8dc76b3 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 14 Sep 2026 08:04:21 +0200 Subject: [PATCH 10/12] style[next]: keep the where and concat_where lowering within ruff's line length Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh --- src/gt4py/next/ffront/foast_to_gtir.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index e840589f2a..68c8267b62 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -415,8 +415,9 @@ 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). + # to keep the IR simpler if not isinstance(node.type, (ts.TupleType, ts.NamedCollectionType)): - return self._lower_and_map("if_", *node.args) # to keep the IR simpler + return self._lower_and_map("if_", *node.args) cond_ = self.visit(node.args[0]) cond_symref_name = f"__cond_{itir.lenient_ir_fingerprinter(cond_)}" @@ -452,8 +453,9 @@ def create_concat_where( return im.concat_where(domain, true_, false_) branch_types = (node.args[1].type, node.args[2].type) + # to keep the IR simpler 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 create_concat_where(true_branch, false_branch, branch_types) return lowering_utils.process_elements( create_concat_where, (true_branch, false_branch), node.type, arg_types=branch_types ) From ee44aae8ed9e987eb8bcf55943e4d090167941b2 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 14 Sep 2026 11:02:44 +0200 Subject: [PATCH 11/12] fix[next]: reject mismatched named-collection branches in embedded where and concat_where Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh --- src/gt4py/next/ffront/fbuiltins.py | 8 ++++- .../unit_tests/ffront_tests/test_fbuiltins.py | 31 ++++++++++++++++++- 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/src/gt4py/next/ffront/fbuiltins.py b/src/gt4py/next/ffront/fbuiltins.py index dfd4ece6c7..c4f93b13cf 100644 --- a/src/gt4py/next/ffront/fbuiltins.py +++ b/src/gt4py/next/ffront/fbuiltins.py @@ -243,7 +243,13 @@ 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): + if isinstance(true_field, named_collections.CUSTOM_NAMED_COLLECTION_TYPES) or isinstance( + false_field, named_collections.CUSTOM_NAMED_COLLECTION_TYPES + ): + if type(true_field) is not type(false_field): + raise ValueError( + f"Either both or none can be a named collection of the same type in '{true_field=}' and '{false_field=}'." + ) 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 ) diff --git a/tests/next_tests/unit_tests/ffront_tests/test_fbuiltins.py b/tests/next_tests/unit_tests/ffront_tests/test_fbuiltins.py index 907703ba92..f9d852e730 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_fbuiltins.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_fbuiltins.py @@ -6,12 +6,13 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +import dataclasses import typing import numpy as np import pytest -from gt4py.next import common +from gt4py.next import common, constructors from gt4py.next.ffront import fbuiltins from gt4py.next.type_system import type_specifications as ts @@ -19,6 +20,20 @@ # values inside the domain of every unary math builtin (0.5 is invalid for `arccosh`) _SAFE_INPUT = {"arccosh": 2.0} +IDim = common.Dimension("IDim") + + +@dataclasses.dataclass +class _Pair: + a: common.Field[[IDim], np.float64] + b: common.Field[[IDim], np.float64] + + +@dataclasses.dataclass +class _OtherPair: + a: common.Field[[IDim], np.float64] + b: common.Field[[IDim], np.float64] + @pytest.mark.parametrize("tuple_spelling", [typing.Tuple, tuple, typing.Tuple[int], tuple[int]]) def test_type_conversion_helper_accepts_both_tuple_spellings(tuple_spelling): @@ -50,3 +65,17 @@ def test_unary_math_builtin_scalar_preserves_dtype(name, dtype): @pytest.mark.parametrize("name", fbuiltins.UNARY_MATH_FP_PREDICATE_BUILTIN_NAMES) def test_unary_math_predicate_builtin_scalar_returns_bool(name, dtype): assert isinstance(getattr(fbuiltins, name)(dtype(0.5)), (bool, np.bool_)) + + +@pytest.mark.parametrize( + "true_type, false_type", [(float, _Pair), (_Pair, float), (_Pair, _OtherPair)] +) +def test_where_rejects_mismatched_named_collections(true_type, false_type): + def make(type_): + if type_ is float: + return 1.0 + return type_(*(constructors.as_field([IDim], np.arange(3.0)) for _ in range(2))) + + mask = constructors.as_field([IDim], np.array([True, False, True])) + with pytest.raises(ValueError, match="named collection of the same type"): + fbuiltins.where(mask, make(true_type), make(false_type)) From 581525c795bef39974be9e2561249391f07a2bcd Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 14 Sep 2026 11:02:44 +0200 Subject: [PATCH 12/12] test[next]: cover promoting lists with different element types Claude-Session: https://claude.ai/code/session_01VR1cyTQ4wysovMBAwPBAWh --- .../next_tests/unit_tests/type_system_tests/test_type_info.py | 3 +++ 1 file changed, 3 insertions(+) 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 dbe8e58ba8..14a90409bd 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 @@ -466,6 +466,7 @@ def test_needs_value_extraction(type_spec: ts.TypeSpec, expected: bool): def test_promote_lists(): float64 = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) + int32 = ts.ScalarType(kind=ts.ScalarKind.INT32) V2EDim = Dimension("V2E", kind=DimensionKind.LOCAL) C2EDim = Dimension("C2E", kind=DimensionKind.LOCAL) const_list = ts.ListType(element_type=float64, offset_type=None) @@ -474,5 +475,7 @@ def test_promote_lists(): 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="different element type"): + type_info.promote(v2e_list, ts.ListType(element_type=int32, offset_type=V2EDim)) with pytest.raises(ValueError, match="non-lists"): type_info.promote(v2e_list, float64)