fix[next]: lower unstructured shifts with the offset's own tag by egparedes · Pull Request #2898 · GridTools/gt4py

egparedes · 2026-09-18T06:05:00Z

Lowering emitted the Python variable name an offset was bound to as the IR shift tag, because ts.OffsetType did not carry the tag. Embedded execution keys on FieldOffset.value, so the same program needed a different offset provider depending on how it was run — confirmed by running it on v1.2.2:

MyOff = FieldOffset("TAGNAME", ...)
embedded:  {"TAGNAME": conn} OK ; {"MyOff": conn}   -> KeyError 'TAGNAME'
compiled:  {"MyOff": conn}   OK ; {"TAGNAME": conn} -> KeyError 'MyOff'

ts.OffsetType now carries tag, and lowering emits it.

Lowering also matches shift arguments on their type instead of their node shape, so module-qualified offsets work. a(mod.V2E), a(mod.E2V[0]) and a(mod.Koff[1]) used to fail on every compiled backend with Unexpected shift arguments! (a foast.Attribute is not a foast.Name), while embedded ran them — the same embedded/compiled divergence. Only tagged offsets are matched (tag=str()), and field(Off) requires two targets, so an untagged (Dim + 1)[0] or a bare Cartesian a(Koff) still gets the lowering error rather than an assertion or a silently wrong neighbors.

This supersedes #2730 by @havogt, which carried the same fix under the field name name; the type-driven matching and one of its tests are taken from there (co-authored). tag is kept over name because it matches common.Tag / FieldOffset.value, and name is easily confused with the Python variable name — which is exactly what the bug conflated.

Why `tag` is `Optional`

A Cartesian shift written Dim + offset has no tag and needs none — it lowers to an itir.CartesianOffset carrying both dimensions, with no provider lookup. type_deduction builds an OffsetType for exactly that case (:711), so a required field would break it. Subscripting (Off[1]) drops the local dimension but propagates the tag, which is the offset's identity.

⚠️ Behaviour change

For a declaration whose tag differs from the variable it is bound to, compiled backends previously required offset_provider={"MyOff": conn} and now require {"TAGNAME": conn}. That divergence from embedded execution is the bug being fixed, but it is user-visible.

ICON4Py is unaffected: all 16 FieldOffset declarations have tag == variable name, there are no Koff[...] subscripts in model code, and as_offset never consults the tag.

No CHANGELOG.md entry — that file is only ever touched by release PRs.

Tests

The regression test grows from one cell (a(Off[1]) on GTFN_CPU) to {shift, neighbor_sum} × {tag ≠ variable name, tag ≠ local dimension name} across the whole backend matrix, plus two lowering unit tests in test_foast_to_gtir.py that assert the emitted OffsetLiteral directly. Both unit tests were verified to fail with the fix reverted.

For the type-driven matching: test_import_from_mod.py gains a module-qualified neighbor_sum(a(cases.V2E)) (from #2730) and a(cases.E2V[0]) across the backend matrix (20 compiled-backend failures before the fix), and test_foast_to_gtir.py gains two error-path tests ((TDim + 1)[0], bare inp(TOff)), each verified to fail with its guard removed.

Each Case holds exactly one connectivity on purpose: DaCe walks every offset-provider entry while building the SDFG and looks a connectivity up by its local dimension's name, so a second non-conforming entry fails a program that never uses it.

The remaining failures are marked per backend, from measurement rather than assumption:

	roundtrip	roundtrip.gtir	gtfn	embedded	dace
shift, tag ≠ varname	✅	✅	✅	✅	✅
reduction, tag ≠ varname	✅	✅	✅	✅	✅
shift, tag ≠ local dim	✅	✅	✅	✅	xfail
reduction, tag ≠ local dim	✅	xfail	xfail	xfail	xfail

Note roundtrip passes the reduction case while roundtrip.gtir does not: roundtrip.default runs apply_common_transforms, so the reduction unrolls keyed by the offset tag, whereas roundtrip.gtir runs only the fieldview transforms and reaches iterator/embedded.py keyed on the local dimension.

Both remaining constraints are one underlying issue: those paths resolve a connectivity through the local dimension's name rather than the offset's identity. Fixing it needs a back-pointer from the local dimension to its connectivity, which is a later step in this stack.

GPU and JAX cells are marked by shared-code-path reasoning, not measurement — they skip locally. xfail_strict is on, so if any is wrong CI fails loudly rather than passing silently.

Verification

At head (d151516): pytest tests/next_tests -m "not uses_dace" → 4495 passed / 0 failed; pytest tests/next_tests/regression_tests -m uses_dace → 6 passed / 0 failed; test_import_from_mod.py on all CPU backends incl. DaCe → 36 passed; mypy src/, tach check, pre-commit run clean. The full -m uses_dace run (1477 passed / 0 failed) was done on the first commit.

Context

First PR of a stack implementing egparedes/connectivities-as-types, an alternative to #2844. This PR stands alone — it is a bugfix that is correct regardless of whether the rest of the stack lands.

Lowering emitted the *Python variable name* an offset was bound to as the IR shift tag, because `ts.OffsetType` did not carry the tag. Embedded execution keys on `FieldOffset.value`, so the same program needed a different offset provider depending on how it was run: MyOff = FieldOffset("TAGNAME", ...) embedded: {"TAGNAME": conn} OK ; {"MyOff": conn} -> KeyError 'TAGNAME' compiled: {"MyOff": conn} OK ; {"TAGNAME": conn} -> KeyError 'MyOff' `ts.OffsetType` now carries `tag`. It is `Optional`, not required: a Cartesian shift built from `Dim + offset` has no tag and needs none, since it lowers to a `CartesianOffset` carrying both dimensions with no provider lookup. Subscripting (`Off[1]`) drops the local dimension but propagates the tag, which is the offset's identity. The regression test grows from one cell -- `a(Off[1])` on gtfn -- to the cross product of {shift, neighbor_sum} x {tag != variable name, tag != local dimension name} over the whole backend matrix. Each `Case` holds exactly one connectivity on purpose: DaCe walks every offset-provider entry while building the SDFG and looks a connectivity up by its *local dimension's* name, so a second, non-conforming entry fails a program that does not use it. The cells that still fail are marked, per backend, from measurement: * `uses_offset_tag_differing_from_local_dim` -- DaCe only; the gtfn shift path was fixed in #1789. * `uses_offset_tag_differing_from_local_dim_in_reduction` -- embedded, gtfn, DaCe and the lower-level `iterator/embedded.py` execution. Notably *not* the roundtrip backend, which passes; only `roundtrip.gtir` fails. Both remaining constraints are the same one: those paths resolve a connectivity through the local dimension's name rather than the offset's identity. Fixing that needs a back-pointer from the local dimension to its connectivity, which is a separate change.

`field(mod.Off)` and `field(mod.Off[idx])` failed lowering with "Unexpected shift arguments!": both shift arms of `_visit_shift` matched a `foast.Name`, and a module-qualified offset is a `foast.Attribute`. Embedded execution was unaffected, so again the execution paths disagreed. Now that `ts.OffsetType` carries the tag, both arms match on the type instead of the node shape, like the `Dim + idx` arm already does. Matching on `tag=str()` rather than asserting it keeps an untagged `Dim + idx` offset out: `(Dim + 1)[0]` still gets the lowering error instead of an assertion. The `field(Off)` arm now requires two targets in its pattern, so a bare Cartesian `FieldOffset` (`field(Koff)`) also gets the lowering error instead of an assertion. That arm moves below `as_offset(...)`, whose type is its first argument's `OffsetType`. Taken from #2730, which carried the same tag under the name `name`. Co-authored-by: Hannes Vogt <vogt@hey.com>

Copilot

Copilot review overview

🟢 Approval recommended

The implementation consistently propagates provider tags and includes comprehensive unit and backend-matrix regression coverage.

Review effort: Balanced
Findings: None

What changed in this PR

Fixes offset lowering to use FieldOffset tags consistently across embedded and compiled execution.

Changes:

Adds tags to OffsetType and preserves them through type deduction.
Uses offset types rather than AST shape during shift lowering.
Expands backend-matrix regression coverage and expected-failure markers.

File	Description
`src/gt4py/next/type_system/type_specifications.py`	Adds optional offset tags.
`src/gt4py/next/ffront/fbuiltins.py`	Propagates `FieldOffset.value` as the tag.
`src/gt4py/next/ffront/foast_passes/type_deduction.py`	Preserves tags through subscripting.
`src/gt4py/next/ffront/foast_to_gtir.py`	Lowers shifts using type-carried tags.
`tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py`	Tests lowering and invalid offset forms.
`tests/next_tests/regression_tests/ffront_tests/test_offset_dimensions_names.py`	Tests differing connectivity names across backends.
`tests/next_tests/integration_tests/feature_tests/ffront_tests/test_import_from_mod.py`	Covers module-qualified offsets.
`tests/next_tests/definitions.py`	Records backend-specific expected failures.
`pyproject.toml`	Registers the new test markers.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

havogt · 2026-09-23T09:22:57Z

+def test_subscripted_untagged_offset_is_a_lowering_error():
+    """Only a tagged offset lowers as `Off[idx]`; the untagged `Dim + idx` one has no key."""
+
+    def foo(inp: gtx.Field[[TDim], float64]):
+        return inp((TDim + 1)[0])
+
+    parsed = FieldOperatorParser.apply_to_function(foo)
+
+    with pytest.raises(FieldOperatorLoweringError, match="Unexpected shift arguments"):
+        FieldOperatorLowering.apply(parsed)
+
+
+def test_bare_cartesian_offset_is_a_lowering_error():
+    """Only an offset with a local dimension lowers as `field(Off)`, i.e. to `neighbors`."""
+
+    def foo(inp: gtx.Field[[TDim], float64]):
+        return inp(TOff)
+
+    parsed = FieldOperatorParser.apply_to_function(foo)
+
+    with pytest.raises(FieldOperatorLoweringError, match="Unexpected shift arguments"):
+        FieldOperatorLowering.apply(parsed)
+
+


I guess they should produce proper DSL Errors with a clear message on what the user did wrong?!

Agreed — both are things a user writes, so they should be located DSLErrors, not an internal lowering error. Fixed in 25f7a32: they are now rejected in type deduction, and the two tests moved to test_diagnostic_messages.py, which pins the rendered text.

Cannot shift by the Cartesian offset 'IOff' without an index. 6 | return a(IOff) | ^^^^ Hint: Give the displacement, e.g. 'IOff[1]'.

Cannot index a dimension shift. 6 | return a((IDim + 1)[0]) | ^^^^^^^^ Note: A shift written as 'Dim + offset' already contains its displacement, unlike a 'FieldOffset', which is indexed to choose one. Hint: Write the displacement directly, e.g. 'IDim + 1'.

The first message quotes the offset as written (IOff, or mod.Koff when module-qualified) rather than its tag 'Ioff' — they can differ, which is the whole point of this PR, and the user's spelling is the one in their file.

One subtlety the existing test_as_offset caught: as_offset(Off, f) has exactly the same type as a bare Off (tagged, one target), so the check matches a bare offset reference rather than the type alone. FieldOperatorLoweringError stays in _visit_shift as the internal guard for anything that slips through.

Local runs on 25f7a32: tests/next_tests -m "not uses_dace" 4495 passed / 0 failed (same counts as before the change, so nothing valid became an error); -m uses_dace 1485 passed / 0 failed; mypy src/ and pre-commit clean.

@havogt

`a(Koff)` (a Cartesian `FieldOffset` without an index) and `a((Dim + 1)[0])` (an indexed dimension shift) are programs a user can write, but they reached lowering and died there with `FieldOperatorLoweringError: Unexpected shift arguments!`, which carries no source location and no explanation. Reject both in type deduction with a located `DSLError` and a hint. The Cartesian-offset message quotes the offset as written (`Koff`, or `mod.Koff`), not its tag, since they can differ and the user's spelling is the one in their file. The lowering error stays as the internal guard it is. Note `as_offset(Off, f)` has the same type as a bare `Off` -- tagged, one target -- so the check matches a bare offset reference, not the type alone. Reported by @havogt in review of #2898.

havogt

lgtm

…ation (#33) Updates the *connectivities as types* note (#32) to the design as implemented in the GridTools/gt4py stack: GridTools/gt4py#2898 → #2899 → #2907 → #2908 → #2909 → #2910 → #2911 → #2912 (ADR 0028: dimensions as nominal types; ADR 0029: connectivities as types). ## What changed in the note - **TL;DR / status**: an implementation-status callout. - **Concepts and sketch** rewritten to the real code: - no `DimensionBaseIndex`, and `LocalDimensionIndex` subclasses `DimensionIndex` - `NeighborConnectivity` is not a `Connectivity` - `offset_tag` names a connectivity in the IR - a connectivity can adopt an existing local dimension, or share one (flattened sparse offsets like ICON4Py's `C2CE`) - `MultiDimensionIndex` is a tuple subclass - `as_offset(KDim, field)` - **Identity rules**: what `resolve` memoizes (so redefined declarations are found), the interactive-`__main__` fallback, the narrow `copyreg` hook for `Staggered[D]`, fingerprinting, the corrected injective mangling (the original scheme was not injective), and a new rule on `offset_tag`. - **Binding model**: providers are normalized to tags at the entry points, and tables are checked against declarations once per compiled variant, including the same-structure requirement for tables over one shared local dimension. - **Effect per layer / what it deletes**: updated. `ts.OffsetType`, `iterator.runtime.offset` and the string-keyed providers below the entry points are *kept*. - **Open questions**: each answered as implemented. Naming convergence with `dependent-local-dimensions` is still open. - **Staging** → **Implementation**: the PR table and a list of where the implementation departs from the original proposal. - The research appendix is marked as a historical record of the pre-implementation tree. `status` stays `draft`: this was written with AI assistance and needs a human review.

…#34) Reworks `personal/egparedes/connectivities-as-types/` after the author's review: - **Current design only.** The status blockquotes, the "where the implementation departs" section, the "how the implementation answered them" list and the review-commit narration are gone. The `Implementation` section is a branch/PR pointer (`connectivities-as-types-8-typed-positions`, GridTools/gt4py#2898–#2912). - **Supersedes `shared/dimensions-as-types`.** The note is now self-contained on the dimension design (ADR 0028). The shared note gets a superseded banner and `status: superseded` (note: not in the AGENTS.md status vocabulary; the previous `proposed` was not either), and its index entry says so. It stays until this proposal moves to `shared/`. - **`Origin` → `Domain`.** `NeighborConnectivity[Domain, Codomain]`, `V2E.domain`. The claimed ADR 0019 origin of the term was wrong; `origin` also collides with the buffer origin (`__gt_origin__`). The gt4py branch still spells it `Origin`; the note flags the pending rename. - **What stays, and why.** A table explains `Connectivity`, `CartesianConnectivity`, `NeighborTable`, `NeighborConnectivityType` and `ts.OffsetType`, and flags the two follow-ups (slimming `NeighborConnectivityType`, renaming `OffsetType.source/target`). - **Declaration = static part of the table's type.** The Concepts section and the Binding model state the "one table per declaration per binding context" invariant explicitly, with the shared-local-dimension relaxation. - **`Staggered[D]` section.** Meaning (ADR 0026), runtime representation, embedded shift convention, IR tag grammar and serialization, gtfn/DaCe treatment. Appendices, the typing probe docstring and the index are updated to match.

…review Fixes found by an independent review of the whole proposal. Wrong claims: - `Staggered[Staggered[K]] = K` is semantically *right* — `Staggered` is an involution on the two cell classes, as the note itself states. Only its expressibility is in question, so the two would be incompatible nominal types. The previous "not the right equation either" argument conflated the type constructor with the index shift. - The extent invariant said the two ranges "differ in size by exactly one when bounded", which the conventions appendix disproves: SGRID's `low` and `high` are bounded with equal sizes. Interleaving gives at most one, and halo extension relaxes even that. - "Both assignments occur" was supported by two examples of the *same* assignment. Replaced the first with a vertex-indexed structured grid. - The degree bit-vector table silently assumed the declared axis indexes its 0-cells, contradicting "the types do not say which class has degree 0" two paragraphs earlier. The assumption is now stated. - Both checkers reject the `__add__` self-type at the definition site, not only mypy, and with different diagnostics, so it costs two separately-spelled suppressions. `staggered_probe.py` gains P4 covering the fourth row of the "four checks become static" table — the one row it did not reach — and both suppressions. The mypy wording quoted from `dimension-generic-fields` was stale. - Scan redesign never appeals to `kind`; "rather than from `kind`" was this note's inference presented as that note's position. - The surface-syntax citation moved from §6.4 to §6.5 for the placement map, and now says its §3.3 cross-link was added alongside this note, so it is a pointer and not independent corroboration. - UGRID *recommends* anticlockwise face nodes ("should"), not requires, which weakens the derivable-incidence-signs suggestion accordingly. SGRID does have absolute numbering via integer coordinate variables, so the `UnitRange` column is a correct relative decoding and "more expressive" is qualified. The `_FillValue` and F4-analogy readings are corrected. Rendering and accounting: - The Implementation table rows for PRs A and B spanned two source lines, so GFM truncated both. Back onto single lines. - "PR 1" was never defined; it is GridTools/gt4py#2898. - The concept-count parenthetical summed to 23, not 25; the appendix's "connectivity type classes" row was missing. - The A1-A10 accounting left two constraints unexplained: A9 also dissolves, A10 remains, re-expressed over `tag`. - The Sketch showed the PEP 695 `Staggered` as if it were the runtime form. - `ts.ShiftType` appeared in both "What it deletes" and "Kept"; the provider-key rule for `{V2E.tag: table}` was unparseable as written. - `typing_probe.py` is not error-clean and now says so. New open question 9 on `DimensionMeta.__eq__` versus the "equality is `is`" identity rule, which the note asserts in one place and overrides in another.

This was referenced Sep 18, 2026

feat[next]: a concrete Dimension is a class, identified by its qualified name #2899

Draft

fix[next]: lower unstructured shifts with the offset's own tag #2896

Closed

feat[next]: a concrete Dimension is a class, identified by its qualified name #2897

Closed

egparedes added this pull request to stack #2900 September 18, 2026 06:16

egparedes and others added 3 commits September 18, 2026 21:29

style[next]: trim comments on the offset tag

ef327f9

egparedes force-pushed the connectivities-as-types-1-shift-tag branch from 59ed409 to ef327f9 Compare September 18, 2026 19:31

egparedes marked this pull request as ready for review September 19, 2026 05:50

egparedes requested a balanced review from Copilot September 19, 2026 05:50

Copilot started reviewing on behalf of egparedes September 19, 2026 05:51 View session

Copilot AI reviewed Sep 19, 2026

View reviewed changes

egparedes requested a review from havogt September 21, 2026 20:38

egparedes mentioned this pull request Sep 22, 2026

personal/egparedes: connectivities as types — update to the implementation GridTools/gt4py_knowledge#33

Merged

havogt reviewed Sep 23, 2026

View reviewed changes

havogt approved these changes Sep 23, 2026

View reviewed changes

egparedes merged commit 7f54f21 into main Sep 23, 2026
32 checks passed

egparedes mentioned this pull request Sep 23, 2026

refactor[next]: carry the connectivity name in OffsetType #2730

Closed

1 of 2 tasks

This was referenced Sep 23, 2026

personal/egparedes: connectivities as types — address review comments GridTools/gt4py_knowledge#34

Merged

feat[next]: NeighborConnectivity[Domain, Codomain] declarations, NeighborTableType, and shared local dimensions on every backend #2907

Draft

Navigation Menu

fix[next]: lower unstructured shifts with the offset's own tag - #2898

fix[next]: lower unstructured shifts with the offset's own tag#2898
egparedes merged 4 commits into
mainfrom
connectivities-as-types-1-shift-tag

egparedes commented Sep 18, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

havogt Sep 23, 2026

Uh oh!

egparedes Sep 23, 2026

Uh oh!

havogt left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

egparedes commented Sep 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why tag is Optional

⚠️ Behaviour change

Tests

Verification

Context

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟢 Approval recommended

Uh oh!

havogt Sep 23, 2026

Choose a reason for hiding this comment

Uh oh!

egparedes Sep 23, 2026

Choose a reason for hiding this comment

Uh oh!

havogt left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

egparedes commented Sep 18, 2026 •

edited

Loading

Why `tag` is `Optional`