Skip to content

fix[next]: support module-qualified offsets in unstructured shifts - #2724

Closed
havogt wants to merge 1 commit into
GridTools:mainfrom
havogt:fix_module_qualified_unstructured_shift
Closed

fix[next]: support module-qualified offsets in unstructured shifts#2724
havogt wants to merge 1 commit into
GridTools:mainfrom
havogt:fix_module_qualified_unstructured_shift

Conversation

@havogt

@havogt havogt commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Description

field(mod.Off) raised FieldOperatorLoweringError: Unexpected shift arguments!, while neighbor_sum(..., axis=mod.LocalDim) worked. The field(Off) arm of _visit_shift matched foast.Name only, so a module-qualified offset (a foast.Attribute) fell through to case _. Extended to also match foast.Attribute.

Only the unstructured arm was affected. The Cartesian arm matches on left=foast.LocatedNode(type=ts.DimensionType(dim=...)), which is node-shape agnostic and carries the Dimension in the type — so field(mod.Dim + 1) already worked. (It was broken the same way until #2667, which changed that pattern from foast.Name to the type-driven form as a side effect of the staggered-shift work.)

The asymmetry is that ts.OffsetType carries only source/target, not the connectivity name, so the unstructured lowering has to fall back on the source-level identifier and therefore must match node shape.

Tests

Added to test_import_from_mod.py, all across the backend matrix:

  • test_import_dims_module_cartesian_shifta(cases.IDim + 1)
  • test_import_dims_module_staggered_shifta(cases.IHalfDim + 0.5)
  • test_import_offset_module_unstructured_shiftneighbor_sum(a(cases.V2E), axis=cases.V2EDim)

The first two pass without the fix; the third is the regression test.

Known remaining gap (not fixed here)

Because the tag comes from the identifier, an offset renamed on import lowers to the wrong tag:

from next_tests.integration_tests.cases import V2E as W2E

@gtx.field_operator
def testee(a: cases.EField) -> cases.VField:
    return neighbor_sum(a(W2E), axis=V2EDim)
embedded: [1 3]
gtfn:     KeyError: "Offset 'W2E' not found in offset provider."

Embedded resolves through FieldOffset.value (fbuiltins.py:485-488), the compiled path through the AST identifier. This dates back to the first FOAST shift lowering (#625, 2022-01-26) and became observable when embedded remap landed (#1309). Fixing it properly means carrying the name in OffsetType / ConnectivityType (cf. the existing TODO(havogt) on OffsetType), which would also make this arm type-driven and let the shape match go away entirely. Left for a follow-up since it touches OffsetType equality.

Requirements

  • All fixes and/or new features come with corresponding tests.
  • Important design decisions have been documented in the appropriate ADR.

`field(mod.Off)` raised `FieldOperatorLoweringError: Unexpected shift
arguments!` because the `field(Off)` arm of `_visit_shift` matched
`foast.Name` only. Extend it to `foast.Attribute`.

`ts.OffsetType` carries no name, so the lowering takes the offset tag
from the source-level identifier. The Cartesian arm does not have this
problem: it matches on `ts.DimensionType`, which carries the `Dimension`
itself, so it is node-shape agnostic already.

Also add tests for module-qualified Cartesian and staggered shifts,
which work but were uncovered.

Known remaining gap: an offset renamed on import (`from mod import V2E
as W2E`) still lowers to the wrong tag, since `attr`/`id` is the
source-level name rather than `FieldOffset.value`. Embedded resolves via
`FieldOffset.value` and disagrees. Fixing that means putting the name in
`OffsetType`/`ConnectivityType`.
@havogt

havogt commented Jul 30, 2026

Copy link
Copy Markdown
Contributor Author

Closing in favor of #2730

@havogt havogt closed this Jul 30, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant