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

docs/development/ADRs/next/0019-Connectivities.md

@@ -31,7 +31,7 @@ We update and introduce the following concepts
 
 **NeighborTable** is a _GatherConnectivity_ that is a 2D mapping of the N neighbors of a Location A to a Location B, backed by a buffer.
 
-**ConnectivityType**, **NeighborConnectivityType** contains all information that is needed for compilation.
+**ConnectivityType**, **NeighborTableType** contain all information that is needed for compilation. A `NeighborTableType` is the type of a table bound to a `NeighborConnectivity` declaration (ADR 0030).
 
 ### Full definitions
 
@@ -48,7 +48,7 @@ Embedded execution of iterator (local) view supports only `NeighborTable`s.
 
 ### IR transformations and compiled backends
 
-All transformations and code-generation should use `ConnectivityType`, not the `Connectivity` which contains the runtime mapping.
+All transformations and code-generation should use `ConnectivityType` / `NeighborTableType`, not the `Connectivity` which contains the runtime mapping.
 
 Note, currently the `global_tmps` pass uses runtime information, therefore this is not strictly enforced.
 
@@ -60,3 +60,7 @@ The only supported `Connectivity`s in compiled backends (currently) are `Neighbo
 
 - Removed the abstract `NeighborConnectivity` concept; `NeighborTable` is now the single neighbor-connectivity concept (there is no non-buffer-backed neighbor connectivity in use).
 - Added `GatherConnectivity` (a `Connectivity` whose `premap` rearranges data via a gather), which the embedded field-view `premap` dispatches on. It replaces the former `ConnectivityKind` flag and unifies the previous reshuffling/remapping `premap` implementations into a single advanced-index gather.
+
+### 2026-09-24
+
+- `NeighborConnectivityType` is renamed `NeighborTableType` and typed by the `NeighborConnectivity` declaration its table is bound to; a `NeighborTable`'s own `__gt_type__()` is the structural `ConnectivityType` (ADR 0030).

docs/development/ADRs/next/0030-Connectivities_As_Types.md

@@ -0,0 +1,225 @@
+---
+tags: []
+---
+
+# Connectivities as Types
+
+- **Status**: proposed
+- **Authors**: Enrique González Paredes (@egparedes)
+- **Created**: 2026-09-21
+- **Updated**: 2026-10-02
+
+A neighbor connectivity is declared as a **class**, and its local dimension as a
+class **nested** in it:
+
+```python
+class V2E(gtx.NeighborConnectivity[Vertex, Edge], max_neighbors=6, min_neighbors=5):
+    class Local(gtx.LocalDimensionIndex): ...
+
+
+@gtx.field_operator
+def f(a: Field[Dims[Edge], float]) -> Field[Dims[Vertex], float]:
+    return neighbor_sum(a(V2E), axis=V2E.Local) + a(V2E[0])
+```
+
+The declaration is written in DSL code, owns its local dimension, and states the
+constraints a neighbor table bound to it has to satisfy. It holds no data. It builds on [ADR 0029](0029-Dimensions_As_Nominal_Types.md): the
+connectivity, like a dimension, is identified by its type, and `V2E.Local` is an
+ordinary dimension class with the tag `<module>.V2E.Local`.
+
+## Context
+
+An unstructured connectivity used to be spelled by four independently authored
+names that had to agree, none of them checked against the others: the
+`FieldOffset` tag, the Python variable it was bound to, the local dimension's
+name and the offset-provider key. The `V2EDim = Dimension("V2E")` convention made
+all four equal, which hid which one each execution path actually used; the
+regression tests in `test_offset_dimensions_names.py` break the convention one
+name at a time. Nothing tied a local dimension to the table it indexes, so the
+backends recovered that link by string equality, and the table's shape, codomain
+and skip values were never checked against the `FieldOffset` declaration.
+
+## Decision
+
+### The declaration
+
+- `NeighborConnectivity[Domain, Codomain]` is a PEP 695 generic whose subclasses
+  are declarations: for each `Domain` element, a list of `Codomain` neighbors. Its
+  metaclass, `ConnectivityMeta`, forbids instantiation. The two dimensions are
+  the class attributes `V2E.domain` and `V2E.codomain`.
+- The local dimension is the nested class `Local`, a subclass of
+  `LocalDimensionIndex`. Declaring it is required, and `NeighborConnectivity`
+  sets `Local.owner` to the connectivity when the class is created. A local
+  dimension can have at most one owner; a declaration redefined under the same
+  name (a re-run notebook cell) takes ownership over again, and for a local
+  dimension adopted rather than nested, the first declaration wins.
+- A local dimension with no table, such as the coefficient axis of a fixed-size
+  stencil, is declared on its own: `class LsqCoeff(LocalDimensionIndex, size=3)`.
+  Its `owner` is `None`. A declaration can also *adopt* such a module-level local
+  dimension, written `Local: TypeAlias = LsqCoeff`, which then keeps its own tag.
+- A connectivity can *share* another one's local dimension,
+  `Local: TypeAlias = C2E.Local`.
+  This is the flattened sparse pattern, e.g. cell-to-cell-edge (`C2CE: Cell -> CellEdge`) indexing the same neighbor axis as `C2E`, so that its results
+  combine with `C2E`-shaped sparse fields. The owner stays `C2E`, and the
+  neighbor counts and skip-value structure are the owner's. A sharer must have
+  the owner's domain; its codomain is free. The local dimension records its
+  sharers (`Local.sharers`) as it records its owner.
+- `max_neighbors` and `min_neighbors` are optional class keywords, not type
+  parameters: Python has no integer type parameters, and nothing static needs
+  the count. A declared count is a constraint on the bound table; an undeclared
+  one is taken from the table. `min_neighbors < max_neighbors` means that the
+  table must use skip values.
+- `common.check_neighbor_table(V2E, table)` checks a table, or just its type
+  (which is all an ahead-of-time compilation has), against the declaration, and
+  returns the table's `NeighborTableType` (below): the domain is
+  `(Domain, V2E.Local)`, the codomain is `Codomain`, the dtype is integral, and
+  the neighbor counts and skip values agree. Skip values are checked on the
+  table's type: a table with a `skip_value` counts as having skip values whether
+  or not an entry uses it.
+
+`Domain` and `Codomain` name the two index spaces the declaration maps between.
+A bound table is a field over `(Domain, Local)` with values in `Codomain`: the
+table's domain is the declaration's domain extended by the local axis, which is
+the same use of the word as `Connectivity.domain` and
+`CartesianConnectivity.domain_dim`. "Origin" would have been the other natural
+name for the first dimension, but gt4py already uses it for the start of a
+buffer (`__gt_origin__`).
+
+### The type of a bound table
+
+Transformations and code generation see types, never tables (ADR 0019). The type
+of a table bound to a declaration is a `common.NeighborTableType`:
+`connectivity` (the declaration), `dtype`, `skip_value` and `max_neighbors`. Its
+`domain` and `codomain` are derived from the declaration,
+`(connectivity.domain, local_dimension_of(connectivity))` and
+`connectivity.codomain`, so they cannot disagree with it. The mapping from
+offset-provider keys to these records is `common.TableTypes`, and it can be
+given instead of the tables for ahead-of-time compilation.
+
+A table cannot tell which declaration it is bound to: the table of a sharer
+(`C2CE`) has the same domain as its owner's (`C2E`), with another codomain. So a
+`NeighborTableType` is built where a table is bound, from its offset-provider key:
+`check_neighbor_table(C2CE, table)`, or `offset_provider_to_type`, which finds
+the declaration whose `offset_tag` is the key among the owner and the sharers of
+the table's local dimension. `NeighborTable.__gt_type__()` returns only what the
+table knows, the structural `common.ConnectivityType` (domain, codomain, dtype,
+skip value).
+
+A table bound under a key that no declaration answers to -- hand-written IR
+names its offsets by plain strings -- has no declaration. Its
+`NeighborTableType` then has the table's structural `ConnectivityType` as its
+`connectivity`, and `domain` and `codomain` are read from that. This keeps the
+IR level, which does not know declarations, working unchanged.
+
+A `NeighborTableType` is fingerprinted through its fields, so the declaration
+takes part in the fingerprint of everything compiled for it: the owner's and a
+sharer's tables, identical as tables, produce different artifact keys.
+
+### `NeighborConnectivity` is not a `Connectivity`
+
+`common.Connectivity` is a *data* protocol (`ndarray`, `domain`, `asnumpy`); a
+declaration holds no data. The neighbor table stays a `Connectivity`
+implementation, and the declaration is only the type the table is checked
+against. `Field.premap` and `Field.__call__` accept either, as they already
+accepted a `FieldOffset`, which is not a `Connectivity` either.
+
+### `LocalDimensionIndex` subclasses `DimensionIndex`
+
+A separate root would force every `type[DimensionIndex]` annotation in the tree
+(`ts.FieldType.dims`, `Domain`, `ConnectivityType.domain`, ...) to widen, and
+would then accept local dimensions wherever a primary one is meant anyway. The
+tree tells local dimensions apart at runtime, and constructors whose parameter
+must be a primary dimension (`NeighborConnectivity[Domain, Codomain]`) check it;
+`Staggered[D]` rejects a local dimension statically too, through its bound on a
+declared Cartesian axis (ADR 0029).
+
+### Localness is the class, and `DimensionKind.LOCAL` is removed
+
+A dimension is local if and only if it subclasses `LocalDimensionIndex`
+(`common.is_local_dimension(dim)`), so `DimensionKind.LOCAL` is removed and
+`DimensionKind` is `HORIZONTAL | VERTICAL`. A local dimension's `kind` is `None`,
+and declaring one with `kind=` is a `TypeError`. `None` rather than `HORIZONTAL`
+keeps every `kind == HORIZONTAL` / `kind != VERTICAL` comparison in the backends
+meaning what it meant: with `HORIZONTAL`, a sparse field would silently count its
+local axis as horizontal. Two consequences:
+
+- `None` does not order against the enum, so `order_dimensions` sorts by an
+  explicit rank — horizontal, then local, then vertical — the order `kind` used to
+  encode. It must not move, or the memory layout of sparse fields changes with it.
+- Displays derive the label from the class: `str(V2E.Local)` is still
+  `Local[local]`, the IR pretty printer still marks a local axis with `ₗ`, and DaCe
+  map variables of a local dimension keep their `_gtx_localdim` suffix.
+
+What remains of `kind` is the layout sort key and the scan axis.
+
+### `Local` is not annotated anywhere
+
+Neither `NeighborConnectivity` nor `ConnectivityMeta` annotates `Local`, and
+that is load-bearing: an annotation makes a declaration's `Local` a *variable*
+for the checkers, so `Field[Dims[Vertex, V2E.Local], float]` is rejected by
+pyright ("Variable not allowed in type expression") for a nested `Local`, and by
+mypy ("not valid as a type") for an adopted or shared one. A real nested `Local`
+on the base is not an option either: pyright reports an incompatible override in
+every declaration. With no annotation, all three spellings are types for both
+checkers, which `typing_tests/pyright_probes.py` pins for pyright and
+`typing_tests/test_next.yaml` for mypy.
+
+The cost is that `conn.Local` is not an attribute the checkers know for a
+*generic* `conn`. Library code reads it through `common.local_dimension_of(conn)`
+instead, and code that has to name a local dimension generically uses a
+`TypeVar` bound to `LocalDimensionIndex`. Writing an adopted or shared local as
+`Local: TypeAlias = ...` (rather than a plain assignment) is what keeps mypy
+treating it as a type.
+
+### Frontend integration
+
+A declaration is typed like the `FieldOffset` it replaces: `V2E.__gt_type__()`
+is a `ts.ShiftType`, which takes a field over the codomain to one over the
+domain, `Shift[<tag>: Edge -> (Vertex, V2E.Local)]`. `V2E[i]` has the domain
+`(Vertex,)`, and so does a Cartesian shift `KDim + 1`, over `KDim` and without a
+tag. The tag is the connectivity's `offset_tag`:
+
+- **the local dimension's tag**, `V2E.Local.tag`, for the connectivity that
+  declares it. This is the single string that shifts, neighbor reductions and
+  sparse arguments already use to find the table in the offset provider, so
+  existing backends need no change.
+- **its own tag**, `C2CE.tag`, for a connectivity that shares another one's local
+  dimension, since the local dimension's tag already names the owner's table.
+  Shifts find the table by that tag. Reductions and sparse arguments know only
+  the local dimension, and take its neighbor count and skip values from a table
+  over it (`common.connectivity_key_over`): the owner's if bound, else the
+  sharer with the smallest tag. Connectivities sharing a local dimension must
+  therefore have the same neighbor *structure* — the same count, and a skip value
+  at the same positions — which is what sharing a neighbor axis means.
+
+`V2E.Local` inside DSL code types as that local dimension, and
+`FieldOffset.Local` names the same thing on a legacy offset, so the spelling
+works for both. The other frontend touch points treat the class like the
+`FieldOffset` it derives: grid-type deduction (`transform_utils`, `past_to_itir`)
+counts it as unstructured, and embedded `premap` accepts it. `V2E[i]` subscripts the
+metaclass, which forwards type-parameter subscription (`NeighborConnectivity[V, E]`) to `__class_getitem__`, since a metaclass `__getitem__` shadows it.
+
+## Consequences
+
+- An unstructured connectivity is spelled once. The provider key, the offset tag
+  and the local dimension are all derived from the declaration.
+- A table bound to a connectivity can be checked against its declaration.
+- `V2E.Local` in DSL code is resolved from the shift type, because the type of
+  `V2E` is not the class.
+- Code generation sees which declaration a table is bound to, not only its
+  shape; a table without a declaration is typed by its structure.
+- A declaration is fingerprinted by its name *and* its declared dimensions and
+  counts, so redefining it under the same name (e.g. re-running a notebook
+  cell) does not reuse artifacts compiled for the old declaration.
+- `FieldOffset` remains during migration; a `FieldOffset` and a
+  `NeighborConnectivity` sharing a local dimension are interchangeable.
+
+## Alternatives considered
+
+- **The local dimension generated by the metaclass**, e.g. `V2E.Local` created
+  from `V2E`'s name. Type checkers cannot see a generated class, so it could not
+  be used in `Field[Dims[Vertex, V2E.Local], ...]`.
+- **Neighbor counts as type parameters.** Python has no integer type parameters,
+  and a `Literal[6]` argument would add a type parameter nothing statically uses.
+- **`NeighborConnectivity` as a `Connectivity` subclass.** Mixes the
+  declaration with the data protocol; see above.

docs/development/ADRs/next/README.md

@@ -23,6 +23,7 @@ Writing a new ADR is simple:
 - [0023 - Fingerprinting](0023-Fingerprinting.md)
 - [0026 - Staggered Dimensions](0026-Staggered_Dimensions.md)
 - [0029 - Dimensions as Nominal Types](0029-Dimensions_As_Nominal_Types.md)
+- [0030 - Connectivities as Types](0030-Connectivities_As_Types.md)
 
 ### Frontend and Parsing #frontend
 

docs/user/next/QuickstartGuide.md

@@ -227,7 +227,7 @@ Another way to look at it is that transform uses the edge-to-cell connectivity t
 You can use the field offset `E2C` below to transform a field over cells to a field over edges using the edge-to-cell connectivities:
 
 ```{code-cell} ipython3
-class E2CDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL): ...
+class E2CDim(gtx.LocalDimensionIndex): ...
 E2C = gtx.FieldOffset(E2CDim.tag, source=CellDim, target=(EdgeDim, E2CDim))
 ```
 
@@ -379,7 +379,7 @@ print("where nested tuple return: {}".format(((result_1.asnumpy(), result_2.asnu
 As explained in the section outline, the pseudo-laplacian needs the cell-to-edge connectivities as well in addition to the edge-to-cell connectivities. Though the connectivity table has been filled in above, you still need to define the local dimension, the field offset, and the offset provider that describe how to use the connectivity table. The procedure is identical to the edge-to-cell connectivity from before:
 
 ```{code-cell} ipython3
-class C2EDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL): ...
+class C2EDim(gtx.LocalDimensionIndex): ...
 C2E = gtx.FieldOffset(C2EDim.tag, source=EdgeDim, target=(CellDim, C2EDim))
 
 C2E_offset_provider = gtx.as_connectivity([CellDim, C2EDim], codomain=EdgeDim, data=cell_to_edge_table, skip_value=-1)

docs/user/next/workshop/exercises/helpers.py

@@ -11,7 +11,14 @@
 import gt4py.next as gtx
 from gt4py.next.iterator.embedded import MutableLocatedField
 from gt4py.next import neighbor_sum, where, Dims
-from gt4py.next import CartesianAxisIndex, Dimension, DimensionIndex, DimensionKind, FieldOffset
+from gt4py.next import (
+    CartesianAxisIndex,
+    Dimension,
+    DimensionIndex,
+    LocalDimensionIndex,
+    DimensionKind,
+    FieldOffset,
+)
 from gt4py.next.program_processors.runners import roundtrip
 from gt4py.next.program_processors.runners.gtfn import (
     run_gtfn as gtfn_cpu,
@@ -389,31 +396,31 @@ class E(DimensionIndex): ...
 class K(CartesianAxisIndex, kind=gtx.DimensionKind.VERTICAL): ...
 
 
-class C2EDim(DimensionIndex, kind=DimensionKind.LOCAL): ...
+class C2EDim(LocalDimensionIndex): ...
 
 
 C2E = FieldOffset(C2EDim.tag, source=E, target=(C, C2EDim))
 
 
-class V2EDim(DimensionIndex, kind=DimensionKind.LOCAL): ...
+class V2EDim(LocalDimensionIndex): ...
 
 
 V2E = FieldOffset(V2EDim.tag, source=E, target=(V, V2EDim))
 
 
-class E2VDim(DimensionIndex, kind=DimensionKind.LOCAL): ...
+class E2VDim(LocalDimensionIndex): ...
 
 
 E2V = FieldOffset(E2VDim.tag, source=V, target=(E, E2VDim))
 
 
-class E2CDim(DimensionIndex, kind=DimensionKind.LOCAL): ...
+class E2CDim(LocalDimensionIndex): ...
 
 
 E2C = FieldOffset(E2CDim.tag, source=C, target=(E, E2CDim))
 
 
-class E2C2VDim(DimensionIndex, kind=DimensionKind.LOCAL): ...
+class E2C2VDim(LocalDimensionIndex): ...
 
 
 E2C2V = FieldOffset(E2C2VDim.tag, source=V, target=(E, E2C2VDim))

docs/user/next/workshop/slides/slides_2.ipynb

@@ -273,7 +273,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "class E2CDim(gtx.DimensionIndex, kind=gtx.DimensionKind.LOCAL): ...\n",
+    "class E2CDim(gtx.LocalDimensionIndex): ...\n",
     "\n",
     "\n",
     "E2C = gtx.FieldOffset(E2CDim.tag, source=Cell, target=(Edge, E2CDim))"

noxfile.py

@@ -349,6 +349,9 @@ def test_typing_exports(session: nox.Session) -> None:
         "typing_tests",
         *session.posargs,
     )
+    # A second checker, on code that must type-check for a downstream user: mypy and pyright
+    # disagree about what counts as a type, which the mypy-only cases above cannot catch.
+    session.run("pyright", "--project", "typing_tests", "typing_tests/pyright_probes.py")
 
 
 # -- DaCe codegen determinism check --

pyproject.toml

@@ -65,6 +65,7 @@ typing = [
 typing_exports = [
   # to test typing with gt4py in downstream code
   {include-group = "typing"},
+  'pyright>=1.1.400',  # the second checker: it disagrees with mypy about what counts as a type
   'types-six>=1.17.0.20251009',  # can not let mypy auto-install types as that leads to unexpected stderr output (which means test failure)
   'pytest-mypy-plugins>=4.0.0',  # pytest plugin for running mypy on code snippets
   "xarray>=2024.1.0"  # one of the regression tests requires xarray
@@ -252,8 +253,6 @@ markers = [
   'uses_ir_if_stmts',
   'uses_lift: tests that require backend support for lift builtin function',
   'uses_negative_modulo: tests that require backend support for modulo on negative numbers',
-  'uses_offset_tag_differing_from_local_dim: tests that shift with an offset tag that differs from the name of its local dimension',
-  'uses_offset_tag_differing_from_local_dim_in_reduction: tests that reduce over an offset whose tag differs from the name of its local dimension',
   'uses_origin: tests that require backend support for domain origin',
   'uses_reduce_with_lambda: tests that use lambdas as reduce functions',
   'uses_reduction_with_only_sparse_fields: tests that require backend support for with sparse fields',