fix[next]: Fix premap between same dimensions#2607
Conversation
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
A couple of questions about the comments and design. I haven't looked at the actual embedded implementation yet
| 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`; |
There was a problem hiding this comment.
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?
| # 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 |
There was a problem hiding this comment.
Isn't it needed a TODO comment somewhere saying the root Field or Connectivity classes shouldn't have the .ndarray property?
There was a problem hiding this comment.
It's already on the Field.ndarray property, I'll link from here.
There was a problem hiding this comment.
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_premapsupports multiple dimension-introducing connectivities and does not require them to share a domain (as demonstrated by new tests liketest_gather_premap_multiple_connectivitiesandtest_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.)
egparedes
left a comment
There was a problem hiding this comment.
In general looks good, just a couple of comments to make docstrings more readable
| 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`). |
There was a problem hiding this comment.
Very minor nitpick:
| 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`). |
| 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.""" |
There was a problem hiding this comment.
I don't understand this docstring and would rephrase it.
|
beverin green, santis offline |
Description
Fixes #1583: embedded
premapproduced wrong results / errors when a connectivity's sourcedimension equals its target dimension while introducing a
LOCALneighbor dimension (e.g.C2E2CO: domain(Cell, C2E2CO), codomainCell). Such a connectivity was misclassified asreshuffling (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
premapchanges dimensions is actually aproperty of the (connectivity, field) pair, so any connectivity-only heuristic is wrong for some
field.
Changes:
_reshuffling_premapand_remapping_premapinto a single_gather_premap— the generaln-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.
ConnectivityKindand allkindproperties. Dispatch is now on connectivity type:affine connectivities (
CartesianConnectivity; cartesian shift / relocation) keep the no-copydomain-only path (
_domain_premap), everything else is aGatherConnectivity(theconnectivity-intrinsic "premap rearranges data" property) and goes through
_gather_premap.NeighborConnectivityintoNeighborTable(there is no non-ndarray-backed neighborconnectivity);
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+LOCALcase (incl. with an extra kept dimension) anda multi-connectivity gather test.
Requirements