diff --git a/content/index.md b/content/index.md index 13d9474..7687787 100644 --- a/content/index.md +++ b/content/index.md @@ -34,6 +34,7 @@ Work-in-progress proposals, organized by contributor. - [[personal/havogt/mesh-and-first-class-halos|A mesh concept with first-class halos]] — keywords: mesh, halos, unstructured, connectivities, offset-provider, domain-inference, distributed, halo-exchange, prior-art - [[personal/havogt/field-data-protocol|A FieldData protocol for gt4py.next embedded fields]] — keywords: fields, domain, data, protocol, embedded, function-fields, boundary-conditions, materialization, lazy, concat_where, origin, prior-art - [[personal/havogt/boundary-condition-syntax|Frontend syntax sugar for boundary conditions over concat_where]] — keywords: frontend, foast, boundary-conditions, concat_where, regions, syntax-sugar, if-elif-else, match, embedded, cartesian-parity, dsl-design, prior-art, python-versions, peps, piece-algebra, metaclass, replay +- [[personal/havogt/jax-connectivities/jax-connectivities|JAX support for connectivities and premap]] — keywords: jax, embedded, connectivities, premap, unstructured, domain-inference, pytree, tracing, distributed, autodiff, prior-art ### edopao diff --git a/content/personal/havogt/jax-connectivities/jax-connectivities.md b/content/personal/havogt/jax-connectivities/jax-connectivities.md new file mode 100644 index 0000000..d927e28 --- /dev/null +++ b/content/personal/havogt/jax-connectivities/jax-connectivities.md @@ -0,0 +1,214 @@ +--- +title: JAX support for connectivities and premap +author: havogt +tags: + [ + jax, + embedded, + connectivities, + premap, + unstructured, + domain-inference, + pytree, + tracing, + distributed, + autodiff, + prior-art, + ] +created: 2026-08-03 +status: draft +--- + +> **TL;DR** Make `premap` work under `jax.jit` by registering a connectivity as a pytree +> whose **child** is the neighbour table and whose **aux data** carries a handle to the +> same buffer, keyed on `id(ndarray)`. Domain inference reads the handle — concrete at +> trace time — so it works for **narrowing** as well as covering ranges, while the table +> travels as a runtime argument instead of being compiled into the module. + +## Problem + +A `Field` already crosses a `jax.jit` boundary cleanly: `Domain` is static pytree aux +data, the array is the traced child. Connectivities do not, and `premap` is where it +breaks, because it uses the table for two things that pull opposite ways: + +- **domain inference** — `inverse_image` derives the *output domain* from the table's + **contents** (`_hyperslice`: `nonzero`/`min`/`max`/`any`/`all` → Python `slice`s). Needs + concrete values. +- **the gather** — needs the table on device as a runtime argument, at ICON scale. + +A connectivity captured as a **closure** keeps contents concrete but JAX inlines the +whole table into the compiled module. Registered as a **plain pytree node** it becomes a +runtime argument but inference receives a tracer. + +Compounding this, `_hyperslice` is untraceable *by construction*: `jnp.nonzero` has a +data-dependent output shape and cannot be staged without an explicit `size=`. Inference +must run eagerly no matter what else changes. + +**The domain generally does not cover the connectivity's image.** Narrowing is the normal +case, not the exception. Any design whose fast path handles only the covering case is +partial support and is rejected. + +## Constraints + +1. **Inference depends on `image_range`**, which comes from the field being premapped and + differs between calls — it cannot be precomputed once per connectivity. +2. **Narrowing is the common case.** See above. This is the constraint that decides the + design. +3. **The table is large** — millions of rows, several per mesh. Nothing may scale with + table size per compiled program. +4. **Gathers must be differentiable** — 4D-var is the motivating use case. +5. `Field.__eq__` is elementwise by design, so fields and connectivities can never be + `static_argnums` (which require `__eq__` to be a real predicate). + +## Design + +Register `JaxArrayConnectivityField` as a pytree node: + +``` +children = (table,) +aux_data = (domain, codomain, skip_value, Handle(table)) +``` + +`Handle` is a thin wrapper whose `__eq__`/`__hash__` key on **`id(table)` — the buffer, +not the connectivity object**. `inverse_image` reads the table through the handle, which +is an ordinary Python reference and therefore concrete during tracing, and runs today's +`_hyperslice` under `jax.ensure_compile_time_eval()`. + +The same buffer is referenced twice in different capacities: as aux data it is read at +trace time only and never staged; as the child it is passed at call time and gathered +from on device. No duplication. + +**Verified for genuine narrowing.** Field on `Vertex[0:30000)`, table 120 000×2 with image +spanning `[0,60000)`: + +``` +jaxpr consts : [] +main params : %arg0: tensor<30000xf64>, %arg1: tensor<120000x2xi32> +StableHLO : 1.4 kB +out domain : Domain(Edge=(0:60000), E2V[local]=(0:2)) +``` + +Narrowing occurred (`Edge` 120 000 → 60 000), the table stayed a jit **argument**, nothing +was inlined. The `restrict()` that follows slices the *traced* child with Python slices, +which is fully traceable. + +### Why keying on the buffer matters + +Retrace granularity is per **buffer**, not per wrapper object: reconstructing the +connectivity around the same array does not retrace; re-uploading the array does. +gt4py's normal flow already reuses buffers — `FieldOffset.as_connectivity_field()` +memoizes on `id(offset_definition)` and returns the same object, fed from the user's +stable offset-provider dict. So in practice this is **one trace per mesh**. + +A content digest would additionally let two separately-uploaded copies of the same table +share a compiled program. It costs a full hash per new buffer and should not be added +speculatively. + +### Optional, free optimisation + +An eagerly-computed `(value_min, value_max, support_box)` descriptor lets `inverse_image` +skip the table read *when the range happens to cover* — pure integer arithmetic, exact +(brute-forced: 0 mismatches in 2 531 covering cases, given a `skip_value ∉ image_range` +guard). This is a shortcut inside the general path, never a fallback boundary, so it +changes no contract. Worth having; not load-bearing. + +## Rejected alternatives + +**Connectivity as a jit closure** (today). The table is inlined as `stablehlo.constant` — +unconditional in jax 0.6.2, no size threshold. Fatal independently of size: closure +constants are replicated per device, invisible to `in_shardings`, not donatable, and JAX +*refuses outright* to close over an array spanning non-addressable devices, so multi-node +is impossible. An upstream fix exists (`JAX_USE_SIMPLIFIED_JAXPR_CONSTANTS`, jax ≥ 0.7.1) +but is off by default and postdates our pin. + +**Static bounds descriptor as the design.** Exact only in the covering case, falling back +to eager inference otherwise — i.e. precisely the partial support that is out of scope. +Demoted to the optimisation above. + +**Caller-supplied output domain.** The `segment_sum(num_segments=)` / jraph / e3nn-jax +analogy does *not* transfer: those declare a trivial quantity the caller already knows, +whereas a narrowed gt4py domain is a non-trivial function of table contents that nobody +upstream of `inverse_image` knows. Declaring it means the caller re-runs the same O(n) +scan by hand, and field view gives up automatic domain inference — a core promise +(ADRs 0010, 0020). Keep it as an explicit **escape hatch for `as_offset`**, whose table +is built from a runtime field so no eager metadata can ever exist. + +**Host-mirrored per-row/block bounds.** A summary is not a distinct design — any O(n) +summary is not O(1)-comparable and needs the same handle treatment. Its one real benefit +is that host-side inference needs no `ensure_compile_time_eval`, dodging the `lax.scan` +and `shard_map` caveats below; if that is wanted, **mirror the whole table on host rather +than summarising**, since per-row min/max is not exact (a row straddling `image_range` +without intersecting it is indistinguishable from one that does, so it fails +conservatively — partial support again). A block decomposition +`(row_boundary, vmin, vmax)` exploiting ICON's renumbering would be O(#blocks) and +value-typed, restoring cross-mesh program sharing; it needs a measurement of #blocks on a +real grid and a fix for the straddling gap. Future work. + +## Two independent defects, both blocking + +Neither is caused by this design; both must be fixed regardless, and the second becomes +unavoidable once narrowing is the norm. + +**NaN gradients from skip values.** `_gather_premap` relies on `-1` wrapping around and +masks only later in `_make_reduction`, so under `grad` the sentinel gathers a real element +that flows through the operator's nonlinearities: `grad = [0.25, 0.167, 0.125, nan]`. Fix +is one `where` sanitising the *index* before the gather — which also removes an existing +oddity where the domain-start offset shifts `-1` onto a different wrapped element. + +**`neighbor_sum` after a narrowing `premap` is already broken — in NumPy too.** +`_make_reduction` broadcasts the full context table against the narrowed field: +`operands could not be broadcast together with shapes (4,2) (2,2)`. It must restrict the +offset definition to the field's domain before masking — a static slice given the Domain, +so it stays traceable. Related: it masks against the module-level +`common._DEFAULT_SKIP_VALUE` rather than the connectivity's own `skip_value`. + +## Risks and open questions + +- **Mesh replacement retains old tables.** The table is a jit argument, so the caller + keeps it alive anyway and aux data adds *zero* incremental retention (a per-rank R2B7 + decomposition is ~25 MB across ~10 tables, shared by all programs). But old compiled + programs pin old tables, and `jax.clear_caches()` does not release them — only dropping + the jitted callables does. Matters for ensemble and nested-grid workflows. +- **`ensure_compile_time_eval` is broken under `lax.scan`** in 0.6.2 + (`NotImplementedError: Evaluation rule for 'empty' not implemented`). If gt4py programs + get wrapped in `lax.scan` for time-stepping this is live, and it is the one concrete + argument for the host-mirror variant. +- **`shard_map` is deferred, not solved.** icon4py runs MPI with per-rank local index + spaces and halos (GHEX), not JAX's global-array model; there the handle *is* the + rank-local table and the inferred domain is rank-local, which is correct. The aux/child + mismatch only arises under `shard_map` over a JAX global array, and that world raises + the prior question of what a `Domain` means globally — a larger decision that should not + drive connectivity representation. Note `ensure_compile_time_eval` under `shard_map` + produces `Manual` HloShardings that reportedly do not work. +- `restrict()` on a traced table cannot recompute metadata; the constructor must tolerate + its absence. +- The traced connectivity must be **installed in the embedded context**, not merely passed + to `premap` — `_make_reduction` reads the offset provider from the contextvar, so the + design silently degrades to closure constants if a caller forgets. + +## Sketch of the change + +`common.py` — document that a `Connectivity` may never itself be aux data (`__eq__` raises +by design, producing a confusing `ValueError` from `jaxlib`). Optionally hang the bounds +descriptor off `NeighborConnectivityType`, next to `max_neighbors`. + +`nd_array_field.py` — add the buffer-keyed handle; register `JaxArrayConnectivityField` +as a pytree node; route `inverse_image` through the handle; sanitise skip indices in +`_gather_premap`; restrict the offset definition in `_make_reduction` and use the +connectivity's own `skip_value`. + +Tests worth pinning: a **narrowing** `premap` under `jit` with the table as an argument +and nothing in `jaxpr.consts`; `neighbor_sum` after a narrowing `premap` (currently +failing on NumPy); and `grad` through a gather with skip values producing no NaNs. + +## Related + +- [[personal/havogt/mesh-and-first-class-halos|A mesh concept with first-class halos]] — + overlaps directly; a mesh owning its connectivities is the natural home for any + precomputed metadata, and owns the local-index-space question deferred above. +- [[personal/havogt/dependent-local-dimensions|Dependent local dimensions and connectivity chains]] — + local dimensions and reductions consume `premap`'s output. +- [[personal/havogt/scan-redesign|Redesign of the vertical scan]] — same static-metadata + versus traced-data tension, same `lax.scan` caveat. +- [[personal/havogt/jax-connectivities/jax-connectivities_research|Research appendix]] — + measurements, prior-art survey, JAX-internals citations. diff --git a/content/personal/havogt/jax-connectivities/jax-connectivities_research.md b/content/personal/havogt/jax-connectivities/jax-connectivities_research.md new file mode 100644 index 0000000..a627407 --- /dev/null +++ b/content/personal/havogt/jax-connectivities/jax-connectivities_research.md @@ -0,0 +1,282 @@ +--- +title: JAX connectivities — research appendix +author: havogt +tags: [jax, connectivities, premap, pytree, tracing, distributed, autodiff, prior-art] +created: 2026-08-03 +status: draft +--- + +> Appendix to [[personal/havogt/jax-connectivities/jax-connectivities|JAX support for +> connectivities and premap]]. Measurements against **jax/jaxlib 0.6.2**, CPU-only +> jaxlib. Claims are labelled **[M]** measured, **[D]** documented, **[S]** read in +> source, **[I]** inferred. + +## Correction: narrowing is the common case + +The first pass of this research assumed *"input fields are allocated over the full +horizontal range; program `domain=` restrictions apply to the output"*, and on that basis +recommended a static bounds descriptor. **That premise is wrong** — realistic field +domains do not cover the connectivity's image, so narrowing is the normal case. The +bounds descriptor then degrades to eager inference on every real call, which is the +partial premap support that was explicitly out of scope. + +Two objections raised against the handle design were **overstated to wrong**, and the +corrections are recorded in place below: + +- **[M]** Retrace granularity is per **buffer**, not per wrapper object: + `1st=1, new-wrapper-same-buffer=1, new-buffer-same-contents=2`. The earlier + "reconstructed connectivity retraces" figure came from `jnp.asarray()` minting a fresh + buffer, not from wrapper reconstruction. `FieldOffset.as_connectivity_field()` memoizes + on `id(offset_definition)` (`fbuiltins.py:496-514`), so gt4py's normal flow reuses + buffers — one trace per mesh. +- **[M]** Cache pinning adds **zero** incremental retention: the table is a jit argument, + so the caller keeps it alive regardless. A per-rank ICON R2B7 decomposition is ~25 MB + across ~10 tables, shared by all programs. The residual footgun is mesh *replacement* — + old programs pin old tables and `jax.clear_caches()` does not release them. + +**[M]** Narrowing verified end-to-end on the handle design: field on `Vertex[0:30000)`, +table 120 000×2 with image over `[0,60000)` → `jaxpr consts: []`, table present as +`%arg1: tensor<120000x2xi32>`, StableHLO 1.4 kB, output domain narrowed to +`Edge=(0:60000)`. + +**[M]** A further pre-existing defect surfaces once narrowing is normal: `neighbor_sum` +after a narrowing `premap` fails **on NumPy as well as JAX** — +`operands could not be broadcast together with shapes (4,2) (2,2)` — because +`_make_reduction` (`nd_array_field.py:1005-1013`) broadcasts the full context table +against the narrowed field instead of restricting the offset definition to the field's +domain first. + +The covering-case measurements below remain valid; they simply describe an optimisation +inside the general path rather than a design. + +## Measured comparison + +Field 60 000×f64, table 120 000×2 int32 (0.96 MB): + +| | closure (today) | identity handle | bounds descriptor | +| --- | --- | --- | --- | +| jaxpr consts | `[(120000,2) int32]` | `[]` | `[]` | +| StableHLO text | 1.92 MB | small | **2.4 kB** | +| table in `func @main` | no | yes | **yes** | +| retraces, same conn object | 1 | 1 | 1 | +| retraces, reconstructed conn | 2 | 2 | **1** | +| retraces, different mesh, same shape+bounds | 3 | 3 | **1** | +| `grad` | — | works | works (`sum(grad) == table.size`) | + +**[M]** The covering-path equivalence was brute-forced over 18 458 `(table, +image_range)` pairs (random 2-D tables, with and without `skip_value=-1`) against the +real `_hyperslice`: **0 mismatches** in the 2 531 covering cases. Without the +`skip_value ∉ image_range` guard, 176 mismatch — all from ranges containing `-1`. + +**[M]** Of non-covering pairs, only **8.9%** yield a valid narrowing; the rest raise +`_hyperslice`'s "non-contiguous or empty". + +## Corrections to earlier measurements + +**[M]** The apparent 2× blow-up (6.4 MB table → 12.8 MB StableHLO) is an artefact of +`.as_text()` rendering decimal ASCII. `module.operation.write_bytecode()` is ~1×. The +cost is nonetheless real: **[S]** that bytecode is `sha256`'d for the persistent +compilation-cache key on every lowering (`jax/_src/cache_key.py:214-219`). Peak RSS +growth measured ≈110 MB for a 16 MB constant. + +**[S]** Constant inlining is unconditional in 0.6.2: `core.is_literalable` requires +rank 0 (`jax/_src/core.py:510`), so every shaped closure array becomes a constvar, and +`mlir.lower_jaxpr_to_fun` emits `ir_constant` for each +(`jax/_src/interpreters/mlir.py:1765`). No size threshold, no flag. + +**[D]** Upstream fix `JAX_USE_SIMPLIFIED_JAXPR_CONSTANTS=True` hoists large constants +into arguments, jax ≥ 0.7.1 — +[PR #30180](https://github.com/jax-ml/jax/pull/30180), +[docs](https://docs.jax.dev/en/latest/internals/constants.html). Default `False`. +Tracking issue [#3220](https://github.com/google/jax/issues/3220) open since 2020. + +**[S]** Multi-node is refused outright for closures: `jax/_src/array.py:1100-1114` — +*"Closing over jax.Array that spans non-addressable (non process local) devices is not +allowed. Please pass such arrays as arguments to the function."* + +Diagnostic available today: `JAX_CAPTURED_CONSTANTS_WARN_BYTES` (default 2 GB) plus +`JAX_CAPTURED_CONSTANTS_REPORT_FRAMES=-1` gives a per-constant capture-site report +(`jax/_src/interpreters/mlir.py:1111-1153`, jax 0.6.1, absent from the CHANGELOG). + +## JAX mechanisms + +**`ensure_compile_time_eval`** — **[S]** since 0.4.36 its body is +`with config.eager_constant_folding(True): yield` (`jax/_src/core.py:1383-1444`); the +mechanism is one branch in `partial_eval.JaxprTrace.process_primitive` +(`jax/_src/interpreters/partial_eval.py:1989-1992`): with no `Tracer` inputs, bind on +`core.eval_trace` instead of staging. + +- **[M]** The docstring is **wrong** about raising `ConcretizationTypeError` when eager + evaluation is impossible — tracer-valued ops silently fall through to staging. +- **[M]** Works under `grad`, `vmap`, `while_loop`, simple `shard_map`. **Broken under + `lax.scan`**: `NotImplementedError: Evaluation rule for 'empty' not implemented` + (cf. [#29996](https://github.com/jax-ml/jax/issues/29996)). Unverified on ≥0.7. +- **[D]** `shard_map` caveat from a maintainer + ([discussion #31461](https://github.com/jax-ml/jax/discussions/31461)): it *"creates + `Manual` HloShardings that don't really work in the execution part"*; workaround + `with jax.sharding.use_abstract_mesh(AbstractMesh((), ())):`. + +**Pytree aux data** — **[D]** the +[custom-pytree docs](https://docs.jax.dev/en/latest/custom_pytrees.html) require +"meaningful hashing and equality" and say the hash enters the jit cache key. +**[M] The hash part is false in jaxlib 0.6.2**: `PyTreeDef.__hash__` ignores custom-node +aux data entirely — identical hashes for aux `'a'`/`'b'`/`1`/`2`/`None`. Unhashable aux +(a `list`) works and never raises. Correctness rests entirely on `__eq__`; do not rely +on this, the docs assert the opposite. + +**[M]** A *bare* array in aux data works on the first call and raises on the second: +`ValueError: Exception raised while checking equality of metadata fields of pytree … +(Note: arrays cannot be passed as metadata fields!)`. Identity short-circuits the first +comparison, so single-call tests hide it. **[M]** `Connectivity.__eq__` raises +`TypeError` by design (`common.py:1039`), so a `Connectivity` can never be aux data. + +**[D]** `register_dataclass`'s own docstring (`jax/_src/tree_util.py:934-947`): +*"Metadata fields must be static, hashable, immutable objects … metadata fields cannot +contain `jax.Array` or `numpy.ndarray` objects."* + +**Cache retention** — **[S]** `jax/_src/util.py:317-337`: `weakref_lru_cache` holds a +weak ref to the **first argument only**, strong refs to the rest, `maxsize=2048`. +`_infer_params_cached(fun, jit_info, signature, …)` makes the treedef (hence aux data) +strongly held. **[M]** aux data survives `jax.clear_caches()` and `f._clear_cache()`; +only `del f` frees it, and `gc.get_referrers` cannot see the holder. Related: +[#16278](https://github.com/jax-ml/jax/issues/16278), +[#11448](https://github.com/jax-ml/jax/issues/11448). + +**`static_argnums`** — **[S]** `jax/_src/api_util.py:116-129` boxes static args +requiring `hash(val)` and `type(a) is type(b) and a == b`. **[M]** With a gt4py field: +*"static arguments should be comparable using `__eq__`"* then `ValueError: The truth +value of an array with more than one element is ambiguous`. Canonical issue +[#24204](https://github.com/jax-ml/jax/issues/24204) — errors on the *second* call. + +**Donation** — **[M]** irrelevant for index arrays: `UserWarning: Some donated buffers +were not usable`. **[I]** Donation needs an output to alias into; a read-only table has +none. + +## Prior art + +The rule, consistent across every library surveyed: **arrays are children; only small, +hashable, value-typed descriptors are aux data; and no library derives a shape from +index-array contents inside `jit`.** + +| library | aux data | children | static shapes | recompiles avoided | +| --- | --- | --- | --- | --- | +| **jraph** | *nothing* — plain `NamedTuple` | all 7 fields incl. `senders`/`receivers` | user declares ints to `pad_with_graphs` (*"do not support jax.jit"*); `segment_sum(num_segments=)` | topology is pure data → zero recompiles | +| **jax-md** | `max_occupancy` (**content-derived int**), `format`, `cell_size` | `idx`, `reference_position`, `error.code` | eager `allocate()` computes `int(occupancy * 1.25)` | fixed `[N, max_occupancy]`; recompile only on re-`allocate` | +| **e3nn-jax** | `irreps` (representation type) | `(array,)` | `jnp.where(…, size=)`; `jnp.unique(size=x.shape[0]) # Pigeonhole` | one treedef per irreps type | +| **Equinox** | `field(static=True)` values | everything `is_array` | n/a | `filter_jit`; **warns** if an array reaches a static field | +| **flax** | `pytree_node=False` fields | `variables` collections | Module is the static half | `nnx` **hard-errors** since 0.12.0 on data in a static attribute | +| **jax-cfd** | `(offset, grid)`; `Grid` is tuples of ints/floats | `(data,)` | structured — coordinates regenerated via `jnp.arange` | tiny descriptor is the whole key | +| **jax-fem** | *none* | element values, not indices | gather/scatter **hoisted out** of the kernel | inner kernel is index-free | +| **PyG** (contrast) | `EdgeIndex._sparse_size`, `Index._dim_size` | the index tensor | eager `int(edge_index.max())+1`, memoized | eager tolerates the sync; `torch.compile` **graph-breaks**, remedy is "pass the size explicitly" | + +Smoking guns, read in source: + +```python +# jraph/_src/models.py:166-168 +# Equivalent to jnp.sum(n_node), but jittable +sum_n_node = tree.tree_leaves(nodes)[0].shape[0] +``` +```python +# jax_md/partition.py:1090-1096 -- reachable only from allocate(), NOT jit +max_occupancy = int(occupancy * capacity_multiplier + _extra_capacity) +``` +```python +# jax_md/partition.py:1253 +"Converts a sparse neighbor list to dense ids. Cannot be JIT." +``` + +The two closest precedents are **jax-md**'s allocate/update split — `allocate` *"cannot +be compiled, since it uses the values of positions to infer the shapes"*, with +`max_occupancy` as a content-derived static int documented as *"Changing this will +invoke a recompilation"* — and **PyG**'s `EdgeIndex`, which exists precisely to attach +`_sparse_size` next to the index tensor. + +**jax-cfd** shows the structured-grid escape gt4py does not have: coordinates are +*regenerated* under jit. An unstructured mesh has no such compression. + +## Skip values and masking + +| library | sentinel | made safe by | +| --- | --- | --- | +| jraph | padding edges → node 0 | padding nodes carry zero features; `arange(static) < traced` masks | +| jax-md | index `N` (one past end) | gather clamping, then `out *= mask`; `safe_mask` double-`where` | +| e3nn-jax | `-1`, optionally remapped via `fill_src` | `mode="promise_in_bounds"` — refuses to tolerate it silently | +| **gt4py** | `-1`, relying on wraparound | masked *after* the gather, in `_make_reduction` | + +**[S/M]** gt4py's `-1` wraps in JAX exactly as in NumPy — +`jax/_src/numpy/indexing.py:188-200` emits `select(index < 0, index + axis_size, index)` +for traced indices too. Note `jnp.take`'s `mode='fill'` does **not** catch it: after +normalisation the index is in-bounds. + +**[M]** `_gather_premap` subtracts the domain start from the index array +(`nd_array_field.py:679`), so with a field domain starting at 2 the sentinel `-1` +becomes `-3` — a *different* wrapped element. Harmless for the primal, but it selects +which element receives spurious cotangent. + +### The NaN-gradient bug + +**[M]** With `x = [4, 9, 16, -1]`, `tbl = [[0,1],[2,-1]]`: + +``` +primal : 9.0 <- correct +grad : [0.25 0.1667 0.125 nan] +``` + +Both remedies verified: the double-`where` sanitising the **operand** (jax-md's +`safe_mask`; [JAX FAQ](https://docs.jax.dev/en/latest/faq.html#gradients-contain-nan-where-using-where)), +or sanitising the **index** before the gather. The latter is the better fit — one +`where` in `_gather_premap`, and it also removes the domain-start shift above. + +## Distributed + +**[M]** Auto-sharding a gather with sharded operand *and* sharded indices (8 CPU +devices): the optimised HLO contains `all-gather(%param), replica_groups=[1,8]` — every +device materialises the entire global field. Confirmed at N=64 and N=2²⁰. +**[S]** XLA's `gather_scatter_handler.cc` has five partitioning methods; the +communication-free ones need iota-like monotone indices, so an unstructured table falls +back to all-gather or a whole-output all-reduce. **`jit` auto-sharding will not scale; +`shard_map` is mandatory.** + +**[S/M]** In explicit-sharding mode gather has *no* sharding rule +(`jax/_src/lax/slicing.py:1950-1963`): *"Use `.at[...].get(out_sharding=)` to provide +output PartitionSpec for the gather indexing."* + +**[D]** There is no halo-exchange example in the 0.6.2 `shard_map` tutorial — only a +passing mention under `ppermute`. The working example is pmap-based +([Wave_Equation.ipynb](https://github.com/jax-ml/jax/blob/main/cloud_tpu_colabs/Wave_Equation.ipynb)). +**[I]** Per-shard index renumbering has no documented JAX pattern; ICON already does it, +so this is a port rather than research. + +**[M]** The decisive constraint: inside `shard_map` the child is a per-shard tracer of +local shape while an aux-data table reference keeps the **global** shape. Any path +reading the aux table instead of the child silently sees global data — which is exactly +what the identity-handle design does in `inverse_image`. The bounds descriptor has no +such path. + +## Differentiability + +**[S]** `_gather_transpose_rule` is `scatter_add` (`jax/_src/lax/slicing.py:2007-2027`); +`ad.defjvp(gather_p, _gather_jvp_rule, None)` marks indices non-differentiable, so index +provenance is irrelevant to `grad`. + +- **`unique_indices` must never be set true.** **[S]** Comment at `slicing.py:1989`: + *"We don't consume `unique_indices` directly in gather(), only in its transpose."* A + connectivity is non-unique by construction; setting it leaves the forward pass correct + and **silently drops gradient contributions**. +- **[D]** Scatter-add on GPU is non-deterministic for conflicting updates + (`--xla_gpu_deterministic_ops=true` fixes it at a cost, + [#17844](https://github.com/jax-ml/jax/issues/17844)); performance issue + [#8637](https://github.com/jax-ml/jax/issues/8637) open since 2021. **[I]** The VJP of + an edge→vertex gather is the high-contention case; the standard remedy is a + `custom_vjp` whose backward uses the inverse connectivity. Future work, not a blocker. +- **[D]** OOB gather clamps but OOB scatter is *dropped*, and reverse-mode AD + [does not preserve OOB semantics](https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html#out-of-bounds-indexing). + gt4py's `-1` is in-bounds after normalisation, so it does not benefit. + +## Non-issue + +**[M]** `_gather_premap` builds one full-output-shape index array per field dimension, +which looked like a scalability problem. It is not — XLA fuses it: for a 3-D gather +(8000×2×64 output) compiled `bytes accessed` is **10.3 MB for both** `_gather_premap` +and an idiomatic `jnp.take(arr, tbl, axis=0)`. CPU only; worth one GPU re-check, but no +rewrite indicated.