diff --git a/agent_guides/STYLE_GUIDE.md b/agent_guides/STYLE_GUIDE.md index 77baa6b9..e7fb0566 100644 --- a/agent_guides/STYLE_GUIDE.md +++ b/agent_guides/STYLE_GUIDE.md @@ -676,3 +676,19 @@ Genuine configurable features and the documented environment gates are **not** c and the `ValueError` below it unreachable for any caller who honours the type. When the two disagree, decide which one should do the rejecting and delete the other -- do not leave a check that cannot fire. + + **And when the annotation WOULD close something real, the guard is still the + answer.** Settled 2026-09-07 on `pallas/nearfield_fused_leaf.py`, which is the + harder version of this: those position parameters' bodies checked + `ndim != 3 or shape[-1] != 3` and never the leading extent, so an annotation + would have caught nine real silent acceptances *and* made the documented + `ValueError` unreachable. It was not the #310 case of closing nothing. The + resolution was to strengthen the body -- the check now verifies the mutual + consistency its own docstring already promised -- because the Raises contract is + public and the census is not the objective. An annotation is the wrong instrument + wherever the parameter's own body documents a `ValueError` over its shape, even + when it would work. + + That PR then made the same mistake one level down, which is worth knowing about: + a generic shape check placed *before* an existing specific one made THAT message + unreachable and turned five tests red. Ordering is part of the check. diff --git a/jaccpot/pallas/nearfield_fused_leaf.py b/jaccpot/pallas/nearfield_fused_leaf.py index 0b409102..6ed8bb97 100644 --- a/jaccpot/pallas/nearfield_fused_leaf.py +++ b/jaccpot/pallas/nearfield_fused_leaf.py @@ -95,6 +95,50 @@ def _resolve_subtile(target_subtile: int | None, leaf_width: int) -> int: _POS_WIDTH = 4 +# Why a shape disagreement between these operands is a WRONG NUMBER and not an error, +# stated once and quoted into every message below. Measured 2026-09-07 by +# `bench/annotation_pilot.py`: 9 of the 20 silent acceptances in that pass were this, +# all of them in this file. +_WHY_MUTUAL = ( + "The grid is derived from ONE array's shape and every other operand is indexed " + "with the same block index, so an operand that disagrees is read OUT OF BOUNDS -- " + "and JAX clamps an out-of-bounds index rather than raising, which turns a shape " + "mistake into a plausible wrong answer instead of an error." +) + + +def _require_shape( + name: str, + array: Array, + expected: tuple[int, ...], + because: str = _WHY_MUTUAL, +) -> None: + """Refuse an operand whose shape disagrees with the one the grid was derived from. + + Parameters + ---------- + name : str + Parameter name, so the message names the argument the caller passed. + array : Array + The operand to check. + expected : tuple[int, ...] + The shape implied by the arrays the grid dimensions were read from. + because : str + Why a disagreement matters, appended to the message. Defaults to + ``_WHY_MUTUAL``, which is the reason for every caller in this module. + + Raises + ------ + ValueError + If ``array``'s shape is not ``expected``. + """ + got = tuple(int(d) for d in array.shape) + if got != tuple(expected): + raise ValueError( + f"{name} must have shape {tuple(expected)}, got {got}. {because}" + ) + + def pallas_nearfield_fused_supported() -> bool: """Return whether the active accelerator can run the fused leaf kernel. @@ -408,6 +452,16 @@ def nearfield_fused_leaf_pallas( tile_t = int(target_positions.shape[1]) num_sources = int(source_positions.shape[1]) + # The ndim/trailing checks above are not enough: they never look at the LEADING + # extent, so `target_mask` and the source tables could disagree with the grid and + # be clamped into it. The annotations do not close this either -- `target_mask` + # carries `leaves w` with a BARE `target_positions` beside it, so those axis names + # are bound by one parameter and have nothing to disagree with. + _require_shape("target_mask", target_mask, (num_leaves, tile_t)) + _require_shape("source_positions", source_positions, (num_leaves, num_sources, 3)) + _require_shape("source_masses", source_masses, (num_leaves, num_sources)) + _require_shape("source_mask", source_mask, (num_leaves, num_sources)) + if num_leaves == 0 or tile_t == 0 or num_sources == 0: return jnp.zeros((num_leaves, tile_t, _OUT_WIDTH), dtype=dtype) @@ -616,6 +670,12 @@ def nearfield_leafpair_jax( G : Array Scalar gravitational constant, applied as a plain multiplier. + Raises + ------ + ValueError + If ``leaf_positions`` is not ``[num_leaves, W, 3]``, or if any other operand + disagrees with it on ``num_leaves`` or ``W``. + Returns ------- Array @@ -640,6 +700,26 @@ def nearfield_leafpair_jax( (``tests/unit/operators/test_pallas_nearfield_fused.py`` draws sources from ``x != i``); it is a precondition, not a check. """ + # The reference lane had no shape check at all, and it needs one for a reason the + # Pallas twins do not share: a `leaf_positions` of trailing width 2 makes `accels` + # two-wide, and the final `concatenate` with the potential then returns a 3-wide + # result where 4 is the contract -- an acceleration missing a component, at the + # right rank. Measured on `main`: (3, 2, 2, 3) in gave (3, 2, 3) out, silently. + if leaf_positions.ndim != 3 or leaf_positions.shape[-1] != 3: + raise ValueError( + "leaf_positions must have shape (num_leaves, W, 3), got " + f"{tuple(int(d) for d in leaf_positions.shape)}. A trailing width other " + "than 3 propagates into the accelerations and is then concatenated with " + "the potential, so the result keeps its rank and loses a component." + ) + num_leaves = int(leaf_positions.shape[0]) + leaf_width = int(leaf_positions.shape[1]) + num_source_slots = int(source_leaf_ids.shape[1]) + _require_shape("leaf_masses", leaf_masses, (num_leaves, leaf_width)) + _require_shape("leaf_mask", leaf_mask, (num_leaves, leaf_width)) + _require_shape("source_leaf_ids", source_leaf_ids, (num_leaves, num_source_slots)) + _require_shape("source_valid", source_valid, (num_leaves, num_source_slots)) + safe_sids = jnp.where(source_valid, source_leaf_ids, 0) src_pos = leaf_positions[safe_sids] # (L, S, W, 3) src_mass = leaf_masses[safe_sids] # (L, S, W) @@ -967,6 +1047,13 @@ def nearfield_leafpair_pallas( leaf_width = int(leaf_positions.shape[1]) num_source_slots = int(source_leaf_ids.shape[1]) + # `leaf_positions` is the table the grid comes from AND the gather target, so a + # disagreement here is read twice over. See `_WHY_MUTUAL`. + _require_shape("leaf_masses", leaf_masses, (num_leaves, leaf_width)) + _require_shape("leaf_mask", leaf_mask, (num_leaves, leaf_width)) + _require_shape("source_leaf_ids", source_leaf_ids, (num_leaves, num_source_slots)) + _require_shape("source_valid", source_valid, (num_leaves, num_source_slots)) + if num_leaves == 0 or leaf_width == 0 or num_source_slots == 0: return jnp.zeros((num_leaves, leaf_width, _OUT_WIDTH), dtype=dtype) @@ -1221,6 +1308,17 @@ def nearfield_leafpair_pallas_decoupled( num_sources = int(source_positions.shape[0]) num_source_slots = int(source_leaf_ids.shape[1]) + # The TARGET side, checked here. `num_targets` and `num_sources` are deliberately + # independent in this variant -- that separation is the decoupled form's whole + # point -- so each operand is checked against the one it belongs to rather than + # against a single leaf count. The source side is checked BELOW, after the + # width guard, so that #297's specific message keeps priority over this generic + # one; `test_the_decoupled_source_pool_is_its_own_leading_axis` asserts that + # ordering and caught it when these four sat here. + _require_shape("target_mask", target_mask, (num_targets, leaf_width)) + _require_shape("source_leaf_ids", source_leaf_ids, (num_targets, num_source_slots)) + _require_shape("source_valid", source_valid, (num_targets, num_source_slots)) + # THE SOURCE POOL MUST BE EXACTLY AS WIDE AS THE TARGET BLOCK, and until this check # existed neither violation said anything. The source gather tables' `BlockSpec` # below is built from `leaf_width` -- the TARGET width -- so the kernel reads exactly @@ -1253,6 +1351,11 @@ def nearfield_leafpair_pallas_decoupled( "silently drops the surplus columns." ) + # Now the source tables, against the POOL and not the target block: the width is + # settled by the check above, the leaf count is the pool's own. + _require_shape("source_masses", source_masses, (num_sources, leaf_width)) + _require_shape("source_mask", source_mask, (num_sources, leaf_width)) + if num_targets == 0 or leaf_width == 0 or num_source_slots == 0 or num_sources == 0: return jnp.zeros((num_targets, leaf_width, _OUT_WIDTH), dtype=dtype) diff --git a/tests/unit/pallas/test_fused_leaf_axis_contracts.py b/tests/unit/pallas/test_fused_leaf_axis_contracts.py index 162cce5f..459f27d4 100644 --- a/tests/unit/pallas/test_fused_leaf_axis_contracts.py +++ b/tests/unit/pallas/test_fused_leaf_axis_contracts.py @@ -5,10 +5,17 @@ (num_leaves, W_t, 3)` -- and a decorator would run first and replace that with a generic `TypeCheckError`. `test_the_position_checks_still_fire` is the guard on that decision. -What the bodies do not check is the masks, masses, ids and validity arrays that have to +The bodies did not use to check the masks, masses, ids and validity arrays that have to agree with those positions slot for slot, and that is exactly where this module's real defect lived: #297, where the decoupled lane's source pool had a different width from the target block and the kernel read the surplus out of bounds, silently. + +They check it NOW. The 2026-09-07 pilot re-recording found nine more of the same kind in +this file, and they were closed by strengthening those bodies rather than by annotating -- +see `test_fused_leaf_shape_guards.py`, which owns that half. The decision, and why an +annotation was the wrong instrument even though it would have worked, is in +`docs/annotation_pilot_phase2_2026-08-30.md`. The two files divide as: annotations here, +body guards there. """ from __future__ import annotations diff --git a/tests/unit/pallas/test_fused_leaf_shape_guards.py b/tests/unit/pallas/test_fused_leaf_shape_guards.py new file mode 100644 index 00000000..17269732 --- /dev/null +++ b/tests/unit/pallas/test_fused_leaf_shape_guards.py @@ -0,0 +1,246 @@ +"""Mutual shape consistency in the fused near-field leaf entry points. + +These are the 9 silent acceptances the 2026-09-07 pilot re-recording left in this file, +closed in the BODY rather than with annotations. The reasoning is written out in +`docs/annotation_pilot_phase2_2026-08-30.md`; the short form is that each of these +positions parameters already raises a documented `ValueError` over its shape, so a +decorator would run first and replace it -- and `DELIBERATELY_BARE`'s first entry makes +that a behaviour change rather than a docs change. Strengthening the guard delivers the +same check under the exception the docstring already promises. + +WHY THE OLD CHECKS WERE NOT ENOUGH. They tested `ndim != 3 or shape[-1] != 3` and never +the LEADING extent, while the docstrings promised `ValueError` "if the input shapes are +mutually inconsistent" -- a consistency the bodies did not verify. So the docstring +over-promised, and the gap it left is the dangerous kind: + + perturbing POSITIONS shrinks the Pallas grid, so the output shape shrinks too and a + caller who checks can see it; + + perturbing a MASK leaves the grid alone, so the BlockSpec indexes the mask out of + bounds, JAX CLAMPS, and leaf 3 silently reuses leaf 2's mask -- real particles + masked out or phantom ones included, at the right output shape. + +Measured on `main` before this change: `target_mask` one row short returned a full +(3, 2, 4); so did `target_mask` one column short and `source_positions` one leaf short. + +The reference lane needed a check too, for a different reason: `nearfield_leafpair_jax` +had none at all, and a `leaf_positions` of trailing width 2 produced a 2-wide +acceleration which the final `concatenate` turned into a 3-wide result where 4 is the +contract -- (3, 2, 2, 3) in, (3, 2, 3) out, silently. +""" + +from __future__ import annotations + +import jax.numpy as jnp +import pytest + +from jaccpot.pallas.nearfield_fused_leaf import ( + nearfield_fused_leaf_pallas, + nearfield_leafpair_jax, + nearfield_leafpair_pallas, + nearfield_leafpair_pallas_decoupled, +) + +LEAVES, WT, SRCSLOTS, SRCLEAVES = 3, 2, 4, 2 +SOFT = jnp.asarray(1e-4) +GRAV = jnp.asarray(1.0) + + +def _fused(**over): + """Build one valid `nearfield_fused_leaf_pallas` call. + + Parameters + ---------- + **over : Any + Arguments to replace, one per perturbation. + + Returns + ------- + dict + Keyword arguments with `leaves`, `w` and the source slot count all distinct. + """ + args = { + "target_positions": jnp.zeros((LEAVES, WT, 3), dtype=jnp.float64), + "target_mask": jnp.ones((LEAVES, WT), dtype=bool), + "source_positions": jnp.zeros((LEAVES, SRCSLOTS, 3), dtype=jnp.float64), + "source_masses": jnp.ones((LEAVES, SRCSLOTS), dtype=jnp.float64), + "source_mask": jnp.ones((LEAVES, SRCSLOTS), dtype=bool), + "softening_sq": SOFT, + "G": GRAV, + "interpret": True, + } + args.update(over) + return args + + +def _pair(**over): + """Build one valid `nearfield_leafpair_pallas` call. + + Parameters + ---------- + **over : Any + Arguments to replace. + + Returns + ------- + dict + Keyword arguments for the leaf-pair entry point. + """ + args = { + "leaf_positions": jnp.zeros((LEAVES, WT, 3), dtype=jnp.float64), + "leaf_masses": jnp.ones((LEAVES, WT), dtype=jnp.float64), + "leaf_mask": jnp.ones((LEAVES, WT), dtype=bool), + "source_leaf_ids": jnp.zeros((LEAVES, SRCSLOTS), dtype=jnp.int32), + "source_valid": jnp.ones((LEAVES, SRCSLOTS), dtype=bool), + "softening_sq": SOFT, + "G": GRAV, + "interpret": True, + } + args.update(over) + return args + + +def _decoupled(**over): + """Build one valid `nearfield_leafpair_pallas_decoupled` call. + + Parameters + ---------- + **over : Any + Arguments to replace. + + Returns + ------- + dict + Keyword arguments with the source pool a DIFFERENT leaf count from the targets, + which is the variant's purpose and must keep working. + """ + args = { + "target_positions": jnp.zeros((LEAVES, WT, 3), dtype=jnp.float64), + "target_mask": jnp.ones((LEAVES, WT), dtype=bool), + "source_positions": jnp.zeros((SRCLEAVES, WT, 3), dtype=jnp.float64), + "source_masses": jnp.ones((SRCLEAVES, WT), dtype=jnp.float64), + "source_mask": jnp.ones((SRCLEAVES, WT), dtype=bool), + "source_leaf_ids": jnp.zeros((LEAVES, SRCSLOTS), dtype=jnp.int32), + "source_valid": jnp.ones((LEAVES, SRCSLOTS), dtype=bool), + "softening_sq": SOFT, + "G": GRAV, + "interpret": True, + } + args.update(over) + return args + + +def test_the_valid_calls_all_still_go_through(): + """The control, on all four entry points, including the decoupled independence. + + The decoupled call has 2 source leaves against 3 target leaves on purpose: that + separation is what the variant exists for, and a guard that "fixed" it by tying the + two together would break the lane while looking like an improvement. + """ + assert nearfield_fused_leaf_pallas(**_fused()).shape == (LEAVES, WT, 4) + assert nearfield_leafpair_pallas(**_pair()).shape == (LEAVES, WT, 4) + assert nearfield_leafpair_pallas_decoupled(**_decoupled()).shape == (LEAVES, WT, 4) + plain = {k: v for k, v in _pair().items() if k != "interpret"} + assert nearfield_leafpair_jax(**plain).shape == (LEAVES, WT, 4) + + +@pytest.mark.parametrize( + "over, culprit", + [ + ( + {"target_positions": jnp.zeros((LEAVES - 1, WT, 3), dtype=jnp.float64)}, + "target_mask", + ), + ({"target_mask": jnp.ones((LEAVES - 1, WT), dtype=bool)}, "target_mask"), + ({"target_mask": jnp.ones((LEAVES, WT - 1), dtype=bool)}, "target_mask"), + ( + { + "source_positions": jnp.zeros( + (LEAVES - 1, SRCSLOTS, 3), dtype=jnp.float64 + ) + }, + "source_positions", + ), + ], +) +def test_the_fused_entry_refuses_an_operand_the_grid_would_clamp(over, culprit): + """Each of these returned a full-size, plausible result on `main`. + + `source_mask` and `source_masses` are deliberately absent: they carry + `srcleaves srcslots` and so already cross-check each other, which the pilot + confirmed by REJECTING both -- they were never among the 9. Adding them here would + claim credit for a check that predates this PR. + + Parameters + ---------- + over : dict + The single argument to perturb. + culprit : str + Substring the message must name, so the error points at the argument the caller + actually got wrong rather than at whichever one the kernel noticed. + """ + with pytest.raises(ValueError, match=culprit): + nearfield_fused_leaf_pallas(**_fused(**over)) + + +def test_the_leafpair_entry_refuses_a_short_leaf_table(): + """`leaf_positions` is both the grid source AND the gather target here. + + One leaf short, `main` returned (2, 2, 4) -- the last leaf's contributions dropped + entirely, and every array that still had 3 rows read into the wrong block. + """ + with pytest.raises(ValueError): + nearfield_leafpair_pallas( + **_pair(leaf_positions=jnp.zeros((LEAVES - 1, WT, 3))) + ) + + +@pytest.mark.parametrize( + "over", + [ + {"target_positions": jnp.zeros((LEAVES - 1, WT, 3), dtype=jnp.float64)}, + {"target_mask": jnp.ones((LEAVES, WT - 1), dtype=bool)}, + {"source_positions": jnp.zeros((SRCLEAVES - 1, WT, 3), dtype=jnp.float64)}, + ], +) +def test_the_decoupled_entry_checks_each_side_against_its_own_count(over): + """Targets against `num_targets`, sources against `num_sources`, not one leaf count. + + Parameters + ---------- + over : dict + The single argument to perturb. + """ + with pytest.raises(ValueError): + nearfield_leafpair_pallas_decoupled(**_decoupled(**over)) + + +def test_the_reference_lane_refuses_a_two_component_position(): + """`nearfield_leafpair_jax` had no shape check at all, and lost a component quietly. + + Trailing width 2 gives a 2-wide acceleration, and `concatenate` with the potential + returns a 3-wide result -- the RANK is right, the contract is 4, and nothing raised. + """ + plain = {k: v for k, v in _pair().items() if k != "interpret"} + plain["leaf_positions"] = jnp.zeros((LEAVES, WT, 2), dtype=jnp.float64) + with pytest.raises(ValueError, match="leaf_positions"): + nearfield_leafpair_jax(**plain) + + +def test_the_new_guard_does_not_outrank_the_source_width_message(): + """#297's specific message must keep priority over the generic one. + + The source-width check explains a mechanism the generic message cannot -- a narrower + pool reads out of bounds, a wider one silently drops real particles -- so the source + table checks are placed AFTER it. This caught a real ordering mistake in this PR: + with them placed before, `test_decoupled_source_pool_is_its_own_leading_axis` and + four parametrisations of `test_decoupled_rejects_a_source_pool_of_a_different_width` + all went red. + """ + narrow = _decoupled( + source_positions=jnp.zeros((SRCLEAVES, WT - 1, 3), dtype=jnp.float64), + source_masses=jnp.ones((SRCLEAVES, WT - 1), dtype=jnp.float64), + source_mask=jnp.ones((SRCLEAVES, WT - 1), dtype=bool), + ) + with pytest.raises(ValueError, match="same leaf width"): + nearfield_leafpair_pallas_decoupled(**narrow)