Skip to content

refactor[next]: MultiDimensionIndex and typed embedded positions - #2912

Draft
egparedes wants to merge 3 commits into
connectivities-as-types-6-class-keyed-providersfrom
connectivities-as-types-8-typed-positions
Draft

egparedes wants to merge 3 commits into
connectivities-as-types-6-class-keyed-providersfrom
connectivities-as-types-8-typed-positions

Conversation

@egparedes

@egparedes egparedes commented Sep 22, 2026 •

Copy link
Copy Markdown
Contributor

Stacked on #2910. Part 4 of 4 of the connectivities as types stack, and the last one (ADR 0030, updated here).

What

  • MultiDimensionIndex[D, *Ls] (gt4py.next.MultiDimensionIndex): a position in the product of a primary dimension and local dimensions, e.g. MultiDimensionIndex(Vertex(3), V2E.Local(1)), the second neighbor of vertex 3. It is a tuple subclass, so it indexes a field or a neighbor table directly, compares and hashes like the plain tuple, and survives pickle and copy. A TypeVarTuple cannot carry a bound, so the constructor checks the shape at runtime: one primary index followed by at least one local index, and an owned local dimension must index the neighbors of the primary index's dimension (MultiDimensionIndex(Edge(1), V2E.Local(1)) is a TypeError, since V2E.domain is Vertex).
  • Typed positions in the iterator-level embedded execution. Positions, named indices and domains are keyed by dimension classes instead of tag strings. SparseTag is replaced by an explicit SparseAxis(dim) offset part, and the tag round trip (_dimension_to_tag) and the resolve calls on sparse lookups are gone.

Tests

  • MultiDimensionIndex: indexing a table, tuple equality and hash, pickling, rejected shapes. Plus a typing test of its revealed type and unpacking.
  • Embedded domains reject axes that are not dimensions; a sparse axis steps along its local dimension.
  • The iterator-embedded changes are otherwise covered by the existing iterator unit and integration tests (embedded and roundtrip).

Verification

Run locally (CPU) on this head; 0 failures in every row that ran:

Check 3.12 3.13 3.14
test_next (internal, cpu, nomesh): embedded, roundtrip, gtfn 5078 passed not run passed (CI)
test_next (dace, cpu, nomesh) 1511 passed not run passed (CI)
src/gt4py/next doctests (in each session) 120 passed
test_typing_exports (mypy cases, pyright probes) 21 passed, pyright 0 errors 21 passed, pyright 0 errors 21 passed, pyright 0 errors
pre-commit run --all-files clean

GitHub CI on the pushed head passes every check: test_next (internal and dace; nomesh and atlas) on 3.12 and 3.14, test_cartesian, test_eve, test_storage, notebooks, package, typing-exports and code quality. The local 3.13/3.14 runs were stopped in favour of CI.

@egparedes
egparedes added this pull request to stack #2900 September 22, 2026 07:00
@egparedes
egparedes force-pushed the connectivities-as-types-8-typed-positions branch from 65fa529 to 71d978f Compare September 23, 2026 16:10
@egparedes
egparedes force-pushed the connectivities-as-types-8-typed-positions branch 2 times, most recently from 934c646 to 0238583 Compare September 24, 2026 10:12
@egparedes
egparedes removed this pull request from stack #2900 September 24, 2026 10:14
@egparedes
egparedes changed the base branch from connectivities-as-types-7-constlist-axisliteral to connectivities-as-types-6-class-keyed-providers September 24, 2026 10:14
@egparedes
egparedes added this pull request to stack #2917 September 24, 2026 10:14
@egparedes
egparedes force-pushed the connectivities-as-types-8-typed-positions branch from 0238583 to 22aab76 Compare September 24, 2026 10:17
havogt added a commit to C2SM/icon4py that referenced this pull request Sep 25, 2026
Temporary [tool.uv.sources] override to GridTools/gt4py#2912 at 22aab76b9, the
last PR of stack #2917 (#2899 -> #2907 -> #2910 -> #2912). Pinned by rev, not
branch, so a rebase of the stack cannot silently move what this is built
against. Revert once the stack is released.

The manifest `gt4py==` pins are left alone; a source override does not enforce
them. gt4py pins `dace==2.0.0a9`, which moves dace from 2.0.0a7; nothing else
in the resolution changes.
- MultiDimensionIndex survives pickle and copy (__getnewargs__); docstring says
  what it is and how it behaves under tuple operations
- embedded domains reject axes that are not dimensions, pointing at gtx.resolve;
  the annotations say dimensions only
- remove the unused get_ordered_indices; ADR 0029 wording on what stays string-keyed
- drop the working plan document committed by mistake again
A position in a product needs at least one local index, and an owned local
dimension must index the neighbors of the primary index's dimension. Tests for
both, for an owner-less local axis, and for a direct 'SparseAxis' shift.
@egparedes
egparedes force-pushed the connectivities-as-types-8-typed-positions branch from 22aab76 to 73339fd Compare October 2, 2026 02:14

This branch has not been deployed

No deployments
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