fix[next]: Fix premap between same dimensions by havogt · Pull Request #2607 · GridTools/gt4py

havogt · 2026-05-27T11:17:22Z

Description

Fixes #1583: embedded premap produced wrong results / errors when a connectivity's source
dimension equals its target dimension while introducing a LOCAL neighbor dimension (e.g.
C2E2CO: domain (Cell, C2E2CO), codomain Cell). Such a connectivity was misclassified as
reshuffling (dimension-preserving) instead of remapping (dimension-introducing).

Root cause: the reshuffling-vs-remapping choice was driven by ConnectivityKind.ALTER_DIMS,
computed as a property of the connectivity alone. Whether premap changes dimensions is actually a
property of the (connectivity, field) pair, so any connectivity-only heuristic is wrong for some
field.

Changes:

Unify _reshuffling_premap and _remapping_premap into a single _gather_premap — the general
n-dimensional advanced-index pullback. The output domain is derived from the connectivity and
the field (the codomain is replaced by the connectivity's domain dims; dims already present are
narrowed in place), so the reshuffling/remapping distinction — and the bug — disappear. This also
lifts the previous "single connectivity only" restriction on dimension-introducing premaps.
Remove ConnectivityKind and all kind properties. Dispatch is now on connectivity type:
affine connectivities (CartesianConnectivity; cartesian shift / relocation) keep the no-copy
domain-only path (_domain_premap), everything else is a GatherConnectivity (the
connectivity-intrinsic "premap rearranges data" property) and goes through _gather_premap.
Collapse NeighborConnectivity into NeighborTable (there is no non-ndarray-backed neighbor
connectivity); is_neighbor_connectivity → is_neighbor_table (single structural check).

No behavior change for previously-supported premaps; the gather is strictly more general. Adds
regression tests for the source == target + LOCAL case (incl. with an extra kept dimension) and
a multi-connectivity gather test.

Requirements

All fixes and/or new features come with corresponding tests.
Important design decisions have been documented in the appropriate ADR inside the docs/development/ADRs/ folder.

Copilot

Copilot reviewed 9 out of 9 changed files in this pull request and generated no comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

egparedes

A couple of questions about the comments and design. I haven't looked at the actual embedded implementation yet

egparedes · 2026-05-27T12:51:10Z

-        A table showing the relation between the connectivity kind and the supported cases
-        is shown in :class:`common.ConnectivityKind`.
+        domain translation (i.e. only transform the domain without altering the data). Such
+        connectivities are not :class:`common.GatherConnectivity` and dispatch to `_domain_premap`;


Question: is this the most correct way to express it ("they are NOT GatherConnectivities") instead of saying "they are Cartesian Connectivites" or something along these lines?

egparedes · 2026-05-27T13:06:41Z

+    #  Reordering the bases would fix that but then `NdArrayField`'s arithmetic operators would
+    #  shadow `Connectivity`'s raising stubs (connectivities would start accepting `+` etc.). A bare
+    #  annotation documents the contract for type checkers with no runtime effect.
+    ndarray: core_defs.NDArrayObject


Isn't it needed a TODO comment somewhere saying the root Field or Connectivity classes shouldn't have the .ndarray property?

It's already on the Field.ndarray property, I'll link from here.

Copilot

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated no new comments.

Copilot

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

src/gt4py/next/embedded/nd_array_field.py:275

This Args section still documents the pre-PR restrictions and is now inconsistent with the new behavior. The bullets "be of the same kind or encode only compact domain transformations" and "for reshuffling operations, all connectivities must have the same domain", as well as the parenthetical "remapping operations only support a single connectivity argument", no longer apply: the unified _gather_premap supports multiple dimension-introducing connectivities and does not require them to share a domain (as demonstrated by new tests like test_gather_premap_multiple_connectivities and test_gather_premap_shared_domain_dim). Please update this Args section to describe the actual current contract (e.g. the gather vs. affine mixing rule and the chained-remap rejection implemented at lines 309–328).

        Args:
            *connectivities: connectivities to be used for the `premap` operation. If only one
                connectivity is passed, it will be expanded to fully defined connectivities for
                each dimension in the domain of the field according to the rules described above.
                If more than one connectivity is passed, they all must satisfy:
                - be of the same kind or encode only compact domain transformations
                - the codomain of each connectivity must be different
                - for reshuffling operations, all connectivities must have the same domain
                (Note that remapping operations only support a single connectivity argument.)

Copilot

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

egparedes

In general looks good, just a couple of comments to make docstrings more readable

egparedes · 2026-05-28T11:12:57Z

+        domain translation (i.e. only transform the domain without altering the data). Such affine
+        connectivities only relabel the domain; data-rearranging cases are handled as
+        advanced-indexing gathers (:class:`common.GatherConnectivity`).


Very minor nitpick:

Suggested change

domain translation (i.e. only transform the domain without altering the data). Such affine

connectivities only relabel the domain; data-rearranging cases are handled as

advanced-indexing gathers (:class:`common.GatherConnectivity`).

domain translation (i.e. only transform the domain without altering the data).

Such affine connectivities only relabel the domain; data-rearranging cases are

handled as advanced-indexing gathers (:class:`common.GatherConnectivity`).

egparedes · 2026-05-28T11:23:01Z

+def _gather_output_domain(
+    field_domain: common.Domain, connectivities: Sequence[common.GatherConnectivity]
+) -> common.Domain:
+    """Output domain of a gather: each connectivity's codomain becomes its own domain dimensions."""


I don't understand this docstring and would rephrase it.

havogt · 2026-05-29T06:08:45Z

beverin green, santis offline

fix[next]: Fix premap between same dimensions

bba3df6

havogt requested a review from Copilot May 27, 2026 11:37

Copilot AI reviewed May 27, 2026

View reviewed changes

havogt added 2 commits May 27, 2026 14:16

fix cupy

e0b936c

Update ADR about connectivities

0d4bbfb

egparedes reviewed May 27, 2026

View reviewed changes

egparedes requested a review from Copilot May 27, 2026 13:08

Copilot started reviewing on behalf of egparedes May 27, 2026 13:08 View session

Copilot AI reviewed May 27, 2026

View reviewed changes

havogt added 3 commits May 27, 2026 15:26

address review comments

458c8d2

relax restriction

21d02ba

add tests

c5afa63

havogt requested a review from Copilot May 28, 2026 05:42

Copilot started reviewing on behalf of havogt May 28, 2026 05:42 View session

Copilot AI reviewed May 28, 2026

View reviewed changes

havogt requested a review from Copilot May 28, 2026 09:57

Copilot started reviewing on behalf of havogt May 28, 2026 09:57 View session

havogt requested a review from egparedes May 28, 2026 09:59

Copilot AI reviewed May 28, 2026

View reviewed changes

egparedes approved these changes May 28, 2026

View reviewed changes

havogt and others added 2 commits May 28, 2026 14:34

address review comments

8efc90f

Merge branch 'main' into premap_clean_1583

5240652

havogt merged commit 8b3f022 into GridTools:main May 29, 2026
23 of 24 checks passed

Conversation

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

Uh oh!

Description

Requirements

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

egparedes left a comment

Choose a reason for hiding this comment

Uh oh!

egparedes May 27, 2026

Choose a reason for hiding this comment

Uh oh!

egparedes May 27, 2026

Choose a reason for hiding this comment

Uh oh!

havogt May 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

egparedes left a comment

Choose a reason for hiding this comment

Uh oh!

egparedes May 28, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

egparedes May 28, 2026

Choose a reason for hiding this comment

Uh oh!

havogt commented May 29, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

havogt commented May 27, 2026 •

edited

Loading