diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 94069d67..3eeb2dd0 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -110,7 +110,7 @@ repos: args: [ "--select=F821,F822", - "--builtins=n,t,ct,levels,internal,leaves,edges,pairs,farleaves,blocks,blocksize,orders,slots,nodes,degrees,targets,sources,w,sw,coarse,crossleaves,rows,cols,_", + "--builtins=n,t,ct,levels,internal,leaves,edges,pairs,farleaves,blocks,blocksize,orders,slots,nodes,degrees,targets,sources,w,sw,coarse,crossleaves,rows,cols,chunkflat,_", ] # NOT scoped: the whole repo, `examples/` and `bench/` and `tests/` included. # It was `files: ^jaccpot/` when introduced, because diff --git a/agent_guides/STYLE_GUIDE.md b/agent_guides/STYLE_GUIDE.md index 77baa6b9..d6ff159b 100644 --- a/agent_guides/STYLE_GUIDE.md +++ b/agent_guides/STYLE_GUIDE.md @@ -211,7 +211,7 @@ must also be added to the flake8 hook's `--builtins` list — see 4.4. | `srcslots` | padded neighbour count per target leaf in the materialised source-particle layout | | `edges` | entries of the flattened neighbour list | | `pairs` | entries of a precomputed leaf-pair schedule | -| `chunks`, `chunkflat` | the 2-D chunked scatter schedule | +| `chunks`, `chunkflat` | the 2-D chunked scatter schedule. `chunkflat` also stands alone, for the arrays of ONE chunk in `_m2l.py:_chunk_segment_scatter_add`, where it is the fixed `chunk_size` width shared by the contributions, their target indices and their validity mask | | `farleaves` | the **far-field** leaf view, which is not `leaves`: they differ on the octree backend | | `crossleaves` | the CROSS-domain near view in `distributed/_force_scale.py`, which is not `leaves` either: it degenerates to length 1 when a rank has no cross neighbours | | `coarse` | the remote coarse (LET) tree's nodes, which are a different tree from the local `nodes` | diff --git a/docs/annotation_pilot_phase2_2026-08-30.md b/docs/annotation_pilot_phase2_2026-08-30.md index 7da28321..e31636f6 100644 --- a/docs/annotation_pilot_phase2_2026-08-30.md +++ b/docs/annotation_pilot_phase2_2026-08-30.md @@ -335,3 +335,126 @@ Two things to settle before annotating any of them, neither of which is a measur `compute_leaf_p2p_accelerations_target_block_pairs_only` in `_large_n_blocks.py` remains unmeasurable for a different reason: it is public, exported, and has no direct test at all, so the pilot never records it. It needs a fixture first. + +--- + +## The three pallas/kernel modules, re-recorded — 2026-09-07 + +The table above put `pallas/m2l_complex_fused.py` at 30%, `pallas/nearfield_fused_leaf.py` +at 29% and `runtime/kernels/_m2l.py` at 8%. All three were measured **before** the PRs that +annotated them (`19f1539`, `575c13a`, `cbc99f1`), which is the staleness this document has +now been caught by three times — items 2 and 3 in August, and these three. Re-recorded +together on `main` at `4f3abe2`, 36 module-level targets, `PILOT_MAX_PER_FN=3`, same scope, +1870 passed / 94 skipped, replayed across all four xdist shards: + +| module | tested | accepted | rate | was | ok/inc/unrep | +|---|---|---|---|---|---| +| `pallas/nearfield_fused_leaf.py` | 108 | 9 | **8%** | 29% | 5/0/0 | +| `pallas/m2l_complex_fused.py` | 122 | 9 | **7%** | 30% | 8/0/2 | +| `runtime/kernels/_m2l.py` | 244 | 2 | **1%** | 8% | 15/0/0 | + +**474 perturbations, 20 accepted — 4%**, against the 48% these nine modules opened at. The +three are not 152 parameters of backlog; they are essentially closed. + +### Where the 20 sit + +``` + 9 Pallas ENTRY points nearfield_fused_leaf_pallas 4, leafpair_pallas_decoupled 3, + leafpair_pallas 1, m2l_complex_fused_pallas 1 + 8 broadcast helpers _matvec 3, _matvec_T 3, _block_matmul 2 + 2 the chunked scatter _chunk_segment_scatter_add + 1 a pure-JAX twin nearfield_leafpair_jax +``` + +Two of those groups have a mechanism, and they are different mechanisms. + +**THE PALLAS CLAMP — the 9 entry-point acceptances.** A `pallas_call` derives its grid +from ONE array's shape and indexes every other operand with the same block index, so a +short operand is read out of bounds and **JAX clamps** rather than raising. Verified in +`nearfield_fused_leaf_pallas`, where `num_leaves` and `tile_t` both come from +`target_positions`: + +* perturbing **positions** shrinks the grid, so the output shape shrinks and the last leaf + is silently dropped — visible to a caller who checks; +* perturbing the **mask** leaves the grid at 4, so the BlockSpec indexes mask block 3 out of + bounds and leaf 3 reuses leaf 2's mask. Real particles masked out or phantom ones + included, with **no shape change to notice**. + +The pure-JAX reference twins reject the identical perturbation, because ordinary +broadcasting fails. So the Pallas lane is strictly weaker at shape validation than the lane +`_m2l.py`'s docstring requires it to equal — `nearfield_fused_leaf_jax` accepted **zero** of +its 21 perturbations while its `_pallas` twin accepted 4. + +**THE BROADCAST — the 8 helper acceptances.** Covered in #336: a reduction written as +`jnp.sum(mat * vec[None, :], axis=1)` accepts a length-1 operand and spreads it, and a +rank change (an extra leading axis, or a flattened operand) changes which axis is reduced +rather than raising. Arithmetic objects to a *mismatched* length but not to a broadcastable +one, which is why these eight sit in the three helpers that reduce by hand and not in the +nine `src_mult`/`deltas` functions that reduce by matmul -- all of which accepted zero. + +A corollary that spans both groups: **an annotated axis only bites when a SECOND parameter +binds the same name.** `target_mask: Bool[Array, "leaves w"]` beside a +BARE `target_positions` binds `leaves`/`w` alone, so there is nothing for it to disagree +with. Same for `multipoles`' `sh` in `m2l_complex_fused_pallas`. Half-annotating a family is +worse than it looks: it reads as covered and validates nothing. + +### A category this document did not have: the free axis + +Not every acceptance is a defect. Where an axis is bound by exactly one parameter and +nothing derives from it, perturbing it produces a **well-formed call** and the pilot counts +it as silently accepted — correctly, by its own definition, but it is not work. + +Proven rather than argued, on `_chunk_segment_scatter_add`: 12 recorded calls show +`local_accum`'s leading axis at 7, 15, 31, 127, 255, 511 and 1023 against an **unchanged** +`contribs`. Same for `rows` in `_matvec` and `cols` in `_matvec_T` — a (31, 16) operator +with a (16,) vector is a perfectly good matvec. + +Of the 20, **3 are free-axis artefacts**. Both PRs from this pass assert the free axes stay +accepted, so the pilot's own report cannot later talk someone into cross-binding an axis the +evidence does not support. + +### Nine predictions, written down first: four right, four wrong, one partial + +Recorded because the error has a direction. P1 called `nearfield_fused_leaf_jax` the biggest +gap in the module (it accepted 0); P2 called its `_pallas` twin low (it has the most); P4 +called the dispatcher high (0); P8 called the nine `src_mult`/`deltas` functions open (all +0). Every miss is the same mistake: **predicting acceptance from "nothing annotates it", +forgetting that ordinary array arithmetic rejects most shape errors by failing to +broadcast.** Which is exactly why §4.1 needs this tool and not a grep — and why the four +correct predictions were all about *broadcast* reductions, where arithmetic does not object. + +### What was closed, and the one decision left + +Closed: the three `m2l_complex_fused` broadcast helpers (#336) and +`_chunk_segment_scatter_add` (this PR). `_m2l.py` is done at 1% — the August verdict +("last, and possibly not worth a PR") was right and is now more so. + +`m2l_complex_fused_pallas`'s `multipoles` is **not closable by annotation**: `sh` is bound by +that parameter alone and the output's `sh` derives from it, so the two move together and +stay consistent. `_m2l_one`/`_m2l_one_vjp` stay UNREPLAYABLE (opaque `t`) and so unmeasured, +but every reduction in them goes through the helpers #336 annotates, so they are +transitively protected without being decorated. + +**The remaining 9 are not an annotation job at all — decided 2026-09-07.** Eight are the +Pallas entry points in `pallas/nearfield_fused_leaf.py` and one is its `leafpair_jax` twin. +Annotating the position arrays would give the masks something to agree with, but each of +those parameters raises a documented `ValueError` in its own body, and checklist item 13 +plus `DELIBERATELY_BARE`'s first entry set the policy: changing which exception a caller +sees is a behaviour change. + +Unlike #310's case this was not a pure exception swap, which is what made it a real choice. +The guard checks `ndim != 3 or shape[-1] != 3` and **never the leading extent**, while the +docstring promises `ValueError` "if the input shapes are mutually inconsistent" — a +consistency the body does not actually check. So the docstring over-promises and the +annotation would have delivered it, under a different exception type. + +**Resolved in favour of the guard.** The body check is strengthened to verify what its own +docstring already promises, the `ValueError` stays reachable and stays the exception callers +see, and no annotation is added. That keeps the Raises contract intact and fixes the defect +in the same place the contract is documented. It is a `fix:` PR with its own test, not part +of the annotation burn-down, and the census does not move for it. + +The general rule this settles, for the next module that hits it: where a bare parameter's +own body already documents a `ValueError` over its shape, an annotation is the WRONG +instrument even when it would close something real. Strengthen the guard; the census is not +the objective. diff --git a/jaccpot/runtime/kernels/_m2l.py b/jaccpot/runtime/kernels/_m2l.py index 462a4163..0949086a 100644 --- a/jaccpot/runtime/kernels/_m2l.py +++ b/jaccpot/runtime/kernels/_m2l.py @@ -39,7 +39,7 @@ import jax.numpy as jnp import numpy as np from beartype import beartype -from jaxtyping import Array, Float, Inexact, Int, jaxtyped +from jaxtyping import Array, Bool, Float, Inexact, Int, jaxtyped from yggdrax.grouped_interactions import ( GroupedInteractionBuffers, ) @@ -391,11 +391,12 @@ class count. return blocks_to, blocks_from +@jaxtyped(typechecker=beartype) def _chunk_segment_scatter_add( - local_accum: Array, - contribs: Array, - tgt_chunk: Array, - valid: Array, + local_accum: Inexact[Array, "nodes sh"], + contribs: Inexact[Array, "chunkflat sh"], + tgt_chunk: Int[Array, "chunkflat"], + valid: Bool[Array, "chunkflat"], *, chunk_size: int, ) -> Array: @@ -412,13 +413,17 @@ def _chunk_segment_scatter_add( Parameters ---------- - local_accum : Array - Local coefficient accumulator to add into. - contribs : Array - Per-pair M2L contributions for this chunk. - tgt_chunk : Array + local_accum : Inexact[Array, 'nodes sh'] + Local coefficient accumulator to add into. `nodes` is bound by this parameter + alone -- nothing else in the signature carries it, and the recording shows it + varying over 7, 15, 31, 127, 255, 511 and 1023 against an unchanged + `contribs` -- so it is deliberately NOT cross-checked against anything. + contribs : Inexact[Array, 'chunkflat sh'] + Per-pair M2L contributions for this chunk. Shares `sh` with `local_accum`, + which is the coefficient count the two are added along. + tgt_chunk : Int[Array, 'chunkflat'] Target node index per pair in the chunk. - valid : Array + valid : Bool[Array, 'chunkflat'] Validity mask; the tail chunk is padded. chunk_size : int Fixed chunk width. Static -- it is what makes every chunk the same shape. diff --git a/tests/unit/runtime/test_m2l_shape_contracts.py b/tests/unit/runtime/test_m2l_shape_contracts.py index 7118ab7e..b2e72b64 100644 --- a/tests/unit/runtime/test_m2l_shape_contracts.py +++ b/tests/unit/runtime/test_m2l_shape_contracts.py @@ -20,6 +20,7 @@ from jaxtyping import TypeCheckError from jaccpot.runtime.kernels._m2l import ( + _chunk_segment_scatter_add, _m2l_chunk_contributions, _rotation_blocks_for_grouped_classes, ) @@ -161,3 +162,99 @@ def test_a_two_component_centre_is_rejected(): args["centers"] = args["centers"][:, :-1] with pytest.raises(TypeCheckError): _m2l_chunk_contributions(**args) + + +# --------------------------------------------------------------------------- +# The chunked scatter, from the 2026-09-07 re-recording. +# +# That run put this module at 2 silent acceptances of 244 perturbations -- 1%, down +# from the 8% on 286 that the section above was written against, because the class +# and accumulator families are now closed. BOTH remaining acceptances are in +# `_chunk_segment_scatter_add`, and only ONE of them is a defect. +# +# The defect: `contribs[sort_idx]` gathers with a `sort_idx` whose length comes from +# `tgt_chunk`, so a `contribs` one row short is an out-of-bounds gather, and JAX +# CLAMPS it -- the last row is silently used twice and one pair's contribution is +# scattered into the wrong target. +# +# The non-defect: `local_accum`'s leading axis. 12 recorded calls show it at 7, 15, +# 31, 127, 255, 511 and 1023 against an unchanged `contribs`, so it is genuinely free +# and the pilot was perturbing a free axis. It is asserted below to stay accepted. +# --------------------------------------------------------------------------- + +CHUNK, CHUNK_SH, CHUNK_NODES = 512, 25, 255 + + +def _scatter_args(dtype=jnp.complex128): + """Build one valid chunked-scatter argument set. + + Parameters + ---------- + dtype : Any + Coefficient dtype. The recording shows complex128, complex64 AND float64 here, + which is why the annotation is `Inexact` and not `Float`. + + Returns + ------- + dict + Keyword arguments for :func:`_chunk_segment_scatter_add`. + """ + return { + "local_accum": jnp.zeros((CHUNK_NODES, CHUNK_SH), dtype=dtype), + "contribs": jnp.ones((CHUNK, CHUNK_SH), dtype=dtype), + "tgt_chunk": jnp.zeros((CHUNK,), dtype=jnp.int64), + "valid": jnp.ones((CHUNK,), dtype=bool), + } + + +@pytest.mark.parametrize("dtype", [jnp.complex128, jnp.float64]) +def test_the_chunked_scatter_accepts_both_bases(dtype): + """The control, and the dtype half of it. + + Parameters + ---------- + dtype : Any + Complex for the complex basis, float for the real one. `Float` here would + reject the complex basis outright -- the mistake #293 shipped. + """ + args = _scatter_args(dtype) + out = _chunk_segment_scatter_add(**args, chunk_size=CHUNK) + assert out.shape == (CHUNK_NODES, CHUNK_SH) + + +def test_contributions_shorter_than_their_target_list_are_rejected(): + """The one measured defect: an out-of-bounds gather that JAX clamps. + + On `main` this returned a full (255, 25) accumulator, having silently gathered + row 510 twice. + """ + args = _scatter_args() + args["contribs"] = jnp.ones((CHUNK - 1, CHUNK_SH), dtype=jnp.complex128) + with pytest.raises(TypeCheckError): + _chunk_segment_scatter_add(**args, chunk_size=CHUNK) + + +def test_the_accumulator_and_the_contributions_must_agree_on_sh(): + """Adding coefficients of two different expansion orders. + + Already rejected on `main`, by broadcasting rather than by annotation, so this + accepts either exception: `sh` is pinned here for the name, not for a new check. + The recording agrees on it in all 9 distinct combinations, at 4, 9, 25 and 81. + """ + args = _scatter_args() + args["local_accum"] = jnp.zeros((CHUNK_NODES, CHUNK_SH - 1), dtype=jnp.complex128) + with pytest.raises((TypeCheckError, ValueError)): + _chunk_segment_scatter_add(**args, chunk_size=CHUNK) + + +def test_the_accumulators_node_axis_stays_free(): + """`nodes` is bound by `local_accum` alone and must NOT be cross-checked. + + The second of the pilot's two acceptances is this, and it is not a defect: the + accumulator's length is the target-node count, which has nothing to do with the + chunk. Asserted so nobody "closes" it later on the strength of the pilot's report. + """ + args = _scatter_args() + args["local_accum"] = jnp.zeros((CHUNK_NODES - 1, CHUNK_SH), dtype=jnp.complex128) + out = _chunk_segment_scatter_add(**args, chunk_size=CHUNK) + assert out.shape == (CHUNK_NODES - 1, CHUNK_SH)