idea[havogt]: JAX support for connectivities and premap

[havogt]

Design proposal for making premap work under jax.jit, plus a research appendix.

The problem

premap uses the neighbour table for two things that pull in opposite directions under jit: domain inference reads its contents (inverse_image → _hyperslice), while the gather needs it on device as a runtime argument. A connectivity captured as a closure keeps it concrete but JAX inlines the whole table into the compiled module; registered as a plain pytree node it becomes a runtime argument but inference gets a tracer.

The proposal

Give a connectivity a small eagerly-computed bounds descriptor — (value_min, value_max, support_box), all Python ints — as pytree aux data, and put the table in children. Whenever the field's range covers the table's image (and excludes skip_value), inverse_image answers with pure integer arithmetic and never touches the buffer. The content-inspecting _hyperslice remains as a fallback for the narrowing case.

Brute-forced over 18 458 (table, image_range) pairs: 0 mismatches in the 2 531 covering cases. The skip_value ∉ image_range guard is load-bearing — without it, 176 disagree.

Measured, field 60 000×f64 / table 120 000×2:

	closure (today)	identity handle	bounds descriptor
jaxpr consts	table inlined	[]	[]
StableHLO text	1.92 MB	small	2.4 kB
retraces, different mesh, same shape+bounds	3	3	1

The last row is the point: one compiled program serves every mesh of the same shape and bounds. Neither alternative can do that. The closure form is additionally fatal for icon4py — JAX refuses outright to close over an array spanning non-addressable devices, so multi-node is impossible.

Independent bug, worth fixing on its own

_gather_premap relies on -1 wrapping around and masks only later in _make_reduction. Under grad that yields NaN gradients, since the sentinel gathers a real element that then flows through the operator's nonlinearities — a correctness blocker for 4D-var, unrelated to the rest of this. Fix is one where sanitising the index before the gather.

Notes

• The research appendix carries the measurements, a prior-art table (jraph, jax-md, e3nn-jax, Equinox, flax, jax-cfd, jax-fem, PyG) and JAX-internals citations. The recurring pattern across all of them: arrays are children, only small value-typed descriptors are aux data, and nobody derives a shape from index-array contents inside jit. The closest precedents are jax-md's max_occupancy and PyG's EdgeIndex._sparse_size.
• Biggest open question is distributed: bounds describe the global index space, so under shard_map bounds and table must be constructed together per rank. Flagged rather than solved.
• Everything is against jax 0.6.2 and some of it is version-sensitive — notably that PyTreeDef.__hash__ ignores custom-node aux data (contradicting the docs), and that the upstream constant-hoisting fix needs jax ≥ 0.7.1.
• Overlaps [[personal/havogt/mesh-and-first-class-halos]]; a mesh object owning its connectivities would be the natural home for the descriptor.

[havogt]

Revised in cd8e838 — the recommendation has changed.

The first draft rested on an assumption I should have checked: "input fields are allocated over the full horizontal range; program domain= restrictions apply to the output." That is wrong — realistic field domains do not cover the connectivity's image, so narrowing is the normal case. The static bounds descriptor then falls back to eager inference on every real call, which is exactly the partial premap support that was out of scope to begin with. Withdrawn as the recommendation; kept only as a free shortcut inside the general path.

New recommendation: the buffer-keyed handle design. children = (table,), aux data = (domain, codomain, skip_value, Handle(table)) with Handle keyed on id(table). Inference reads the handle — an ordinary Python reference, concrete at trace time — so it works for narrowing by construction. Verified: field on Vertex[0:30000), table 120 000×2 with image over [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 happened, table stayed a jit argument, nothing inlined.

Two objections I previously raised against this design were overstated to wrong:

• Retrace granularity is per buffer, not per wrapper object (1st=1, new-wrapper-same-buffer=1, new-buffer-same-contents=2). My earlier "reconstruction retraces" figure came from jnp.asarray() minting a fresh buffer. FieldOffset.as_connectivity_field() memoizes on id(offset_definition), so gt4py's normal flow is one trace per mesh.
• Cache pinning adds zero incremental retention — the table is a jit argument, so the caller keeps it alive anyway. ~25 MB per rank for R2B7 across ~10 tables, shared by all programs. Residual footgun is mesh replacement.

shard_map is deferred rather than answered: icon4py runs MPI with per-rank local index spaces (GHEX), where the handle is the rank-local table and the inferred domain is rank-local. The global-array case raises the prior question of what a Domain means globally, which shouldn't drive connectivity representation.

A second pre-existing defect, which only becomes unavoidable under narrowing: neighbor_sum after a narrowing premap is already broken, on NumPy as well as JAX — operands could not be broadcast together with shapes (4,2) (2,2), because _make_reduction broadcasts the full context table against the narrowed field instead of restricting the offset definition first. That plus the NaN-gradient bug are both independent of the JAX design and block it.