fix[next]: `concat_where` crashes on neighbor-list fields by tehrengruber · Pull Request #2837 · GridTools/gt4py

tehrengruber · 2026-08-27T09:51:40Z

Problem

Any concat_where whose branches contain a local (neighbor-list) field crashes during type inference with a bare AssertionError, e.g.:

@gtx.field_operator
def testee(a: VField, b: VField) -> EField:
    t = concat_where(Edge < 2, a(E2V), b(E2V))
    return neighbor_sum(t, axis=E2VDim)

File ".../iterator/type_system/type_synthesizer.py", line 290, in deduce_return_type
    domain.dims, type_info.extract_dims(type_info.promote(tb, fb))
File ".../type_system/type_info.py", line 591, in promote
    assert all(isinstance(dtype, ts.ScalarType) for dtype in extracted_dtypes)
AssertionError

The combination was previously untested (test_concat_where.py had no local-dimension coverage), so the latent assert never fired in CI.

Root cause

The GTIR concat_where type synthesizer derives the result dimensions via type_info.promote(tb, fb). Besides merging the dims, promote also promotes the branches' dtypes and asserts they are all ScalarType — but at the GTIR level a local field is FieldType(dims=[Edge], dtype=ListType(...)).

Fix

Promote only the dims — promote_dims(domain.dims, extract_dims(tb), extract_dims(fb)) — and never build the intermediate promoted field type. Dtype equality of the two branches is already checked separately right above. The two forms are otherwise equivalent: promote_dims computes the canonically ordered union of its inputs, so pairwise-then-merge and the flat three-way merge agree on both results and error cases (fuzz-checked over dim combinations including LOCAL and staggered dims).

The FOAST-level where/concat_where deduction gets the same formulation for consistency; there it is behaviorally neutral, since FOAST-level local fields keep a scalar dtype (the local dim lives in dims).

This is a pragmatic solution to unblock #2833.

Tests

test_with_local_field (minimal repro) and test_with_tuples_of_local_fields (tuple branches, the shape that originally exposed the bug in #2833) in test_concat_where.py. Extracted from #2833, where the fix and test were previously bundled; they are independent of the tuple-comprehension feature.

Disclaimer: This PR and its description were written largely with the help of AI. Code and description look alright.

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.

havogt

Review of the current head (3342813). Everything below was checked against the branch; where I ran something I say so.

The fix works. I reproduced the crash on the merge base (ee1bb4f6a) and confirmed it is gone at this head with correct numerics. test_concat_where.py is 62 passed on gtfn, mypy src/ is clean, and pre-commit passes on the changed files.

One thing that has no code anchor, so it goes here: the PR description no longer matches the diff. The description says the fix is promote_dims(domain.dims, extract_dims(tb), extract_dims(fb)) in type_synthesizer.py, but this head does not touch that file at all — the fix widens type_info.promote to accept ListType. An earlier head of this branch did implement the described approach, so the body appears to have gone stale in a force-push. Two consequences: anyone reviewing from the description reviews the wrong patch, and the "fuzz-checked over dim combinations" claim no longer applies to anything in the diff. Worth updating before merge.

Inline comments below, roughly strongest first.

havogt · 2026-08-28T07:39:08Z

 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


Widening promote is a large blast radius for this fix — it is a frontend-wide API with four call sites.

The only caller that needs the widening, type_synthesizer.py:283-292, already asserts tb_dtype == fb_dtype one line earlier and then uses the result solely for extract_dims(...). So the promoted dtype this change makes possible is discarded immediately afterwards.

The approach described in the PR body — promote only the dims, at the ITIR call site — is a smaller diff and leaves this signature alone. Was there a reason it was abandoned? If it was, it would be useful to record that, because the body still recommends it.

havogt · 2026-08-28T07:39:08Z

-        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))


This drops a real invariant.

The removed assert all(isinstance(dtype, ts.ScalarType) for dtype in extracted_dtypes) enforced exactly the property the new docstring above asserts is true — that the frontend never passes a list here. With the assert gone, nothing enforces it any more, so a genuine frontend bug producing FieldType(dtype=ListType) now flows through silently instead of failing loudly at this line.

If the widening stays, consider keeping the precondition somewhere: either assert it at the frontend call sites, or keep a narrower assert here that admits ListType only when the caller opted in.

havogt · 2026-08-28T07:39:08Z

        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).")


Small wording regression. "Could not promote scalars of different dtype" -> "Could not promote dtypes of different type".

The dominant case reaching this line is still two scalars of differing dtype in a FOAST binop, surfaced as the __cause__ of the DSLError at type_deduction.py:665. The new phrasing is both vaguer and slightly ungrammatical ("dtypes of different type"). Suggest "Could not promote types of different dtype (not implemented).", or keeping the scalar wording and adding a separate branch for lists.

havogt · 2026-08-28T07:39:08Z

+
+@pytest.mark.uses_unstructured_shift
+@pytest.mark.uses_sparse_fields
+def test_with_local_field(unstructured_case, static_domains: bool):


The fix looks incomplete for an adjacent case that is squarely inside this PR's title.

concat_where(Edge < 2, a(E2V), 0.0) — a local field against a scalar — still dies on a bare AssertionError at type_synthesizer.py:284, before promote is reached at all. FOAST accepts it (both branches promote to FieldType([Edge, E2VDim], float64)); at GTIR one branch is a ListType and the other a ScalarType, so the dtype-equality assert fires first. I ran this on this head to confirm.

Pre-existing rather than a regression, and note the promote_dims approach from the PR body would not have covered it either — promote(ListType, ScalarType) raises an uncaught ValueError. But it is the natural thing a user writes, so it is worth either covering here or calling out as explicitly out of scope.

havogt · 2026-08-28T07:39:08Z

+
+
+@pytest.mark.uses_unstructured_shift
+@pytest.mark.uses_sparse_fields


uses_sparse_fields is being used for its dace side effect rather than for its meaning — these tests have no sparse field arguments, whereas everywhere else the marker means exactly that (test_with_toy_connectivity.py:232, 281, 334).

The dace limitation actually being worked around is NotImplementedError: 'concat_where' with list output is not supported (dace/lowering/gtir_to_sdfg_concat_where.py:249), which I confirmed by running the repro on run_dace_cpu_noopt. A future reader hitting this sees the xfail reason "'uses_sparse_fields' tests not supported by dace", which will be misleading. A dedicated marker would be more honest.

Related coverage note: with uses_concat_where xfailing roundtrip/embedded and this marker xfailing all dace backends, gtfn is the sole guard for the new tests. I verified they do genuinely pin the bug — reverting both source hunks makes them fail on gtfn — but they pass on the embedded backends with the fix reverted, so the guard is narrower than it looks. It is cheap, at least: the gtfn failure is at type inference, before any C++ compile.

havogt · 2026-08-28T07:39:08Z

+            np.where(edge_mask[:, np.newaxis], a[e2v_table], b[e2v_table]),
+            axis=1,
+            initial=0,
+            where=e2v_table != common._DEFAULT_SKIP_VALUE,


This skip-value mask is dead code. E2V is constructed with skip_value=None in both mesh descriptors (cases_utils.py:300 and :395), so e2v_table != common._DEFAULT_SKIP_VALUE is all-True in every parametrization and initial=0 never applies.

So despite the uses_sparse_fields marker, these tests never exercise skip values on the local dimension. Either drop the where=/initial= as misleading, or switch to V2E, which does carry _DEFAULT_SKIP_VALUE in skip_value_mesh and would give real coverage.

(Same at :494 and :500.)

havogt · 2026-08-28T07:39:08Z



+@pytest.mark.uses_tuple_returns
+def test_with_nested_tuples(cartesian_case, static_domains: bool):


test_with_nested_tuples does not exercise this fix — it is cartesian with scalar dtypes, and I confirmed it passes with both source hunks reverted.

Harmless as coverage, but it is the one remaining place where the split from #2833 still shows in a fix[next] PR. Its commit message says "moved from #2833"; it might belong back there, or in its own PR.

tehrengruber commented Aug 27, 2026

View reviewed changes

Comment thread src/gt4py/next/iterator/type_system/type_synthesizer.py Outdated

tehrengruber force-pushed the fix-concat-where-local-fields branch from 6e70606 to 5215204 Compare August 27, 2026 19:45

Add 'concat_where' test with nested tuple branches (moved from #2833)

3342813

tehrengruber marked this pull request as ready for review August 27, 2026 20:19

tehrengruber requested a review from havogt August 27, 2026 21:45

tehrengruber changed the title ~~fix[next]: concat_where crashes on local (neighbor-list) fields~~ fix[next]: concat_where crashes on neighbor-list fields Aug 28, 2026

havogt reviewed Aug 28, 2026

View reviewed changes

Navigation Menu

Uh oh!

fix[next]: `concat_where` crashes on neighbor-list fields - #2837

fix[next]: `concat_where` crashes on neighbor-list fields#2837
tehrengruber wants to merge 2 commits into
mainfrom
fix-concat-where-local-fields

tehrengruber commented Aug 27, 2026 •

edited

Loading

Uh oh!

Uh oh!

havogt left a comment

Uh oh!

havogt Aug 28, 2026

Uh oh!

havogt Aug 28, 2026

Uh oh!

havogt Aug 28, 2026

Uh oh!

havogt Aug 28, 2026

Uh oh!

havogt Aug 28, 2026

Uh oh!

havogt Aug 28, 2026

Uh oh!

havogt Aug 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants



		@pytest.mark.uses_unstructured_shift
		@pytest.mark.uses_sparse_fields



		@pytest.mark.uses_tuple_returns
		def test_with_nested_tuples(cartesian_case, static_domains: bool):

Uh oh!

Conversation

tehrengruber commented Aug 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Root cause

Fix

Tests

Uh oh!

Uh oh!

havogt left a comment

Choose a reason for hiding this comment

Uh oh!

havogt Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

havogt Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

havogt Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

havogt Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

havogt Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

havogt Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

havogt Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

tehrengruber commented Aug 27, 2026 •

edited

Loading