From e1f77f73e8cbd42147c1aaf105e0ca05d5c307fd Mon Sep 17 00:00:00 2001 From: TobiBu Date: Mon, 7 Sep 2026 10:02:23 +0200 Subject: [PATCH] refactor(pallas): the fused complex-M2L helpers, where a broadcast hides a bug `m2l_complex_fused.py` carries real and imaginary parts as SEPARATE real arrays rather than as a complex dtype, so every reduction is written twice and the two halves must agree by hand. These three helpers are where that agreement lives, and all three reduce with an explicit broadcast rather than a matmul: _matvec(mat, vec) jnp.sum(mat * vec[None, :], axis=1) _matvec_T(mat, vec) jnp.sum(mat * vec[:, None], axis=0) _block_matmul(...) jnp.sum(block_r * vec_r[:, None, :], -1) - ...(block_i) A broadcast accepts a length-1 operand and SPREADS it, which is why a wrong shape here is quiet. Measured on main: a length-1 `vec` into a (32, 16) operator returned a full (32,) result; `_matvec_T` took a FLATTENED (512,) operator and returned (512,); and a `vec_i` of (4, 1) beside a `vec_r` of (4, 8) returned a plausible (4, 8) with the first column's imaginary part in every lane. That last one is the dangerous member of the set and the pilot never tried it -- its perturbations are leading/trailing/extra/flattened, never length-1 -- so it is a gap this closes beyond what was measured. AXES BY EXECUTION, READ PER CALL. Two sets of extents cannot establish a relation, so the 2026-09-07 recording was read call by call: _matvec mat (16, 32) vec (32,) | mat (32, 16) vec (16,) -> `vec` is mat's SECOND axis, at two extents with the roles SWAPPED _matvec_T mat (16, 32) vec (16,) | (32, 16) vec (32,) mat (32, 128) vec (32,) | (128, 32) vec (128,) -> `vec` is mat's FIRST axis, at four extents `rows`/`cols` therefore name a RELATION between two arguments, not a width -- one helper is applied to three different operators at four extents -- which is why the two new vocabulary entries say so. WHAT THIS DOES NOT CLAIM. Swapping the pair is ALREADY caught without annotations: a (32,) vector will not broadcast against a 16-long axis, so main raises TypeError there and the only change is which exception. Its test accepts either, because asserting TypeCheckError alone would dress an exception-type change up as a closed gap. And two of the pilot's six acceptances on this pair are `rows` and `cols` being perturbed where nothing else binds them -- (31, 16) with a (16,) vector is a well-formed matvec. Those are NOT defects and a test asserts they stay accepted. `_block_matmul_vjp` is deliberately left BARE: it rejected all six perturbations, so 4.1 says leave it alone. Its recording is still what supplies the second extent for squareness, (4, 8, 8) and (8, 16, 16), which `_block_matmul`'s own single recorded extent cannot. Verified the Pallas lane still traces: `_m2l_one` reaches these helpers from inside `pallas_call`, where the operands are kernel tracers, and interpret mode agrees with the JAX twin to 5.6e-16. Eight tests, five red against main. Co-Authored-By: Claude Opus 5 --- .pre-commit-config.yaml | 11 +- agent_guides/STYLE_GUIDE.md | 1 + jaccpot/pallas/m2l_complex_fused.py | 35 ++-- .../test_fused_m2l_helper_axis_contracts.py | 173 ++++++++++++++++++ 4 files changed, 206 insertions(+), 14 deletions(-) create mode 100644 tests/unit/pallas/test_fused_m2l_helper_axis_contracts.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 9a017d52..94069d67 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -85,6 +85,15 @@ repos: # blocksize the target block size, `JACCPOT_LARGE_N_TARGET_BLOCK_SIZE` # farleaves the FAR-field leaf view, which is not the near-field one: # they differ on the octree backend (5 vs 3) + # rows a dense operator's OUTPUT extent, in the two generic + # matvec helpers only. Not a tree or particle quantity: the + # same helper is applied to three different operators at + # four different extents, so the pair `rows`/`cols` names a + # RELATION between its two arguments rather than a width + # cols the same operator's INPUT extent, i.e. the length the + # vector must have. `_matvec` reduces against `cols` and its + # adjoint `_matvec_T` against `rows`; that one swap is the + # entire difference between them # _ an anonymous axis: rank asserted, extent deliberately not. # Needed as a builtin only for the SINGLE-axis spelling # `Int[Array, "_"]`; `"_ 2"` and `"leaves _ _"` are @@ -101,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,_", + "--builtins=n,t,ct,levels,internal,leaves,edges,pairs,farleaves,blocks,blocksize,orders,slots,nodes,degrees,targets,sources,w,sw,coarse,crossleaves,rows,cols,_", ] # 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 2caff49c..77baa6b9 100644 --- a/agent_guides/STYLE_GUIDE.md +++ b/agent_guides/STYLE_GUIDE.md @@ -224,6 +224,7 @@ must also be added to the flake8 hook's `--builtins` list — see 4.4. | `degrees` | spherical-harmonic degrees of a per-degree summary, `p+1` of them | | `orders` | the candidate expansion orders an adaptive policy scores | | `levels` | block-step levels, `k_max + 1` of them | +| `rows`, `cols` | a generic dense operator's output and input extents, in the two `m2l_complex_fused.py` matvec helpers ONLY. Not a tree or particle quantity: one helper is applied to three different operators at four extents, so the pair names a **relation between two arguments** rather than a width. `_matvec` reduces against `cols`, its adjoint `_matvec_T` against `rows`, and that single swap is the whole difference between them | | `2`, `3` | literals -- the `(start, end)` pair and the spatial dimension | | `_` | anonymous: deliberately unnamed, see below | diff --git a/jaccpot/pallas/m2l_complex_fused.py b/jaccpot/pallas/m2l_complex_fused.py index 6f34fb14..d06e8591 100644 --- a/jaccpot/pallas/m2l_complex_fused.py +++ b/jaccpot/pallas/m2l_complex_fused.py @@ -234,8 +234,12 @@ def m2l_complex_fused_tables(order: int) -> dict: # -------------------------------------------------------------------------- +@jaxtyped(typechecker=beartype) def _block_matmul( - block_r: Array, block_i: Array, vec_r: Array, vec_i: Array + block_r: Float[Array, "degrees blockdim blockdim"], + block_i: Float[Array, "degrees blockdim blockdim"], + vec_r: Float[Array, "degrees blockdim"], + vec_i: Float[Array, "degrees blockdim"], ) -> tuple[Array, Array]: """Complex block-diagonal matmul: out[b,i] = sum_j block[b,i,j] vec[b,j]. @@ -245,13 +249,13 @@ def _block_matmul( Parameters ---------- - block_r : Array + block_r : Float[Array, 'degrees blockdim blockdim'] Real part of the per-degree rotation blocks, shape ``(Bp, mdp, mdp)``. - block_i : Array + block_i : Float[Array, 'degrees blockdim blockdim'] Imaginary part, same shape. - vec_r : Array + vec_r : Float[Array, 'degrees blockdim'] Real part of the packed coefficients, shape ``(Bp, mdp)``. - vec_i : Array + vec_i : Float[Array, 'degrees blockdim'] Imaginary part, same shape. Returns @@ -269,15 +273,17 @@ def _block_matmul( return out_r, out_i -def _matvec(mat: Array, vec: Array) -> Array: +@jaxtyped(typechecker=beartype) +def _matvec(mat: Float[Array, "rows cols"], vec: Float[Array, "cols"]) -> Array: """out[i] = sum_j mat[i,j] * vec[j] (gather-free; Triton-GPU friendly). Parameters ---------- - mat : Array + mat : Float[Array, 'rows cols'] Dense operator, shape ``(rows, cols)``. - vec : Array - Vector, shape ``(cols,)``. + vec : Float[Array, 'cols'] + Vector, shape ``(cols,)``. Reduced against ``mat``'s SECOND axis, which is + what distinguishes this from :func:`_matvec_T`. Returns ------- @@ -361,16 +367,19 @@ def _m2l_one( ) -def _matvec_T(mat: Array, vec: Array) -> Array: +@jaxtyped(typechecker=beartype) +def _matvec_T(mat: Float[Array, "rows cols"], vec: Float[Array, "rows"]) -> Array: """out[j] = sum_i mat[i,j] * vec[i] == (mat^T @ vec); the adjoint of _matvec. Parameters ---------- - mat : Array + mat : Float[Array, 'rows cols'] The SAME operand :func:`_matvec` takes, shape ``(rows, cols)`` -- reducing the other axis avoids materialising a transpose, which keeps Triton happy. - vec : Array - Cotangent, shape ``(rows,)``. + vec : Float[Array, 'rows'] + Cotangent, shape ``(rows,)``. `rows` and not `cols`: that ONE difference from + :func:`_matvec` is the whole content of the adjoint, so the two axis names are + what stop the pair being swapped on a non-square operator. Returns ------- diff --git a/tests/unit/pallas/test_fused_m2l_helper_axis_contracts.py b/tests/unit/pallas/test_fused_m2l_helper_axis_contracts.py new file mode 100644 index 00000000..24dfef14 --- /dev/null +++ b/tests/unit/pallas/test_fused_m2l_helper_axis_contracts.py @@ -0,0 +1,173 @@ +"""Axis contracts for the fused complex-M2L inner helpers. + +`m2l_complex_fused.py` carries real and imaginary parts as SEPARATE real arrays rather +than as a complex dtype, so every reduction in it is written twice and the two halves +have to agree by hand. These three helpers are where that agreement lives, and all +three reduce with an explicit broadcast rather than a matmul:: + + _matvec(mat, vec) jnp.sum(mat * vec[None, :], axis=1) + _matvec_T(mat, vec) jnp.sum(mat * vec[:, None], axis=0) + _block_matmul(...) jnp.sum(block_r * vec_r[:, None, :], axis=-1) - ... + +A broadcast accepts a length-1 operand and spreads it, which is why a wrong shape here +is silent rather than loud. Measured on `origin/main` before these annotations: a +length-1 `vec` into a (32, 16) operator returned a full (32,) result, and a `vec_i` of +(4, 1) beside a `vec_r` of (4, 8) returned a full (4, 8) -- the imaginary part of every +column silently equal to the imaginary part of the first. + +The axes come from the 2026-09-07 recording (`tests/unit` + `tests/integration`), read +per CALL so a relation is taken from the pairing and not from two sets of extents: + + _matvec mat (16, 32) vec (32,) | mat (32, 16) vec (16,) + -> `vec` is `mat`'s SECOND axis, at two extents with the roles + swapped, so `rows` and `cols` are independent and not a + coincidence of one payload size + _matvec_T mat (16, 32) vec (16,) | mat (32, 16) vec (32,) + mat (32, 128) vec (32,) | mat (128, 32) vec (128,) + -> `vec` is `mat`'s FIRST axis, at four extents + _block_matmul block (4, 8, 8) vec (4, 8) + +`_block_matmul`'s own recording has ONE extent, so squareness is not established from +it alone; its reverse twin `_block_matmul_vjp` recorded (4, 8, 8) and (8, 16, 16) for +the same operator, and `_m2l_one` the same two, which is where the second extent comes +from. `_block_matmul_vjp` itself is deliberately left BARE: it rejected all six of the +pilot's perturbations, so section 4.1 says leave it alone. +""" + +from __future__ import annotations + +import jax.numpy as jnp +import pytest +from jaxtyping import TypeCheckError + +from jaccpot.pallas.m2l_complex_fused import _block_matmul, _matvec, _matvec_T + +ROWS, COLS = 32, 16 +DEGREES, BLOCKDIM = 4, 8 + + +def _mat(rows=ROWS, cols=COLS): + """A non-square dense operator, so `rows` and `cols` cannot be confused. + + Parameters + ---------- + rows : int + Output extent. + cols : int + Input extent -- the length `_matvec` reduces against. + + Returns + ------- + Array + Shape ``(rows, cols)``, all ones. + """ + return jnp.ones((rows, cols)) + + +def test_the_matched_calls_still_go_through(): + """The control: both helpers and the block matmul accept their recorded shapes.""" + assert _matvec(_mat(), jnp.ones((COLS,))).shape == (ROWS,) + assert _matvec_T(_mat(), jnp.ones((ROWS,))).shape == (COLS,) + out_r, out_i = _block_matmul( + jnp.ones((DEGREES, BLOCKDIM, BLOCKDIM)), + jnp.ones((DEGREES, BLOCKDIM, BLOCKDIM)), + jnp.ones((DEGREES, BLOCKDIM)), + jnp.ones((DEGREES, BLOCKDIM)), + ) + assert out_r.shape == out_i.shape == (DEGREES, BLOCKDIM) + + +@pytest.mark.parametrize("fn", [_matvec, _matvec_T]) +def test_a_length_one_vector_is_no_longer_broadcast(fn): + """One column spread across all of them is the silent failure this closes. + + `jnp.sum(mat * vec[None, :], axis=1)` is happy to broadcast a (1,) operand, and on + `main` both helpers returned a full-length result from it. + + Parameters + ---------- + fn : Callable + `_matvec` or its adjoint; the broadcast is silent in both directions. + """ + with pytest.raises(TypeCheckError): + fn(_mat(), jnp.ones((1,))) + + +def test_matvec_reduces_against_cols_and_its_adjoint_against_rows(): + """The one difference between the pair, asserted on a NON-square operator. + + GREEN ON MAIN, deliberately, and it is the only test here that is. Swapping the + pair does not need an annotation to be caught: a (32,) vector will not broadcast + against a 16-long axis, so `main` already raises -- as `TypeError` rather than + `TypeCheckError`, which is the only thing these annotations change about it. This + test exists to pin the two axis NAMES to the right arguments, so it accepts either + exception; asserting `TypeCheckError` alone would make an exception-type change + look like a closed gap. + + On a square operator the swap is shape-identical and no annotation can see it, + which is why this uses (32, 16). + """ + with pytest.raises((TypeCheckError, TypeError)): + _matvec(_mat(), jnp.ones((ROWS,))) # wants COLS + with pytest.raises((TypeCheckError, TypeError)): + _matvec_T(_mat(), jnp.ones((COLS,))) # wants ROWS + + +def test_an_extra_leading_axis_is_rejected_on_either_operand(): + """A stray batch axis changed the RANK of the result instead of raising. + + Measured on `main`: `_matvec(mat[1, 32, 16], vec[16])` returned (1, 16), and + `_matvec(mat[32, 16], vec[1, 16])` returned (1, 16) -- neither the (32,) the + caller's next reshape expects. + """ + with pytest.raises(TypeCheckError): + _matvec(jnp.ones((1, ROWS, COLS)), jnp.ones((COLS,))) + with pytest.raises(TypeCheckError): + _matvec(_mat(), jnp.ones((1, COLS))) + + +def test_a_flattened_operator_is_rejected(): + """`_matvec_T` took a flat (512,) operator and returned (512,). + + The rank check is the whole of it: with `mat` flat, `vec[:, None]` broadcasts + against it and the reduction runs over the wrong axis entirely. + """ + with pytest.raises(TypeCheckError): + _matvec_T(jnp.ones((ROWS * COLS,)), jnp.ones((ROWS,))) + + +def test_the_real_and_imaginary_halves_must_agree(): + """The dangerous one: a real/imag mismatch inside a complex kernel. + + `_block_matmul` computes `sum(block_r * vec_r) - sum(block_i * vec_i)`. Both terms + are reduced independently, so a `vec_i` of (4, 1) beside a `vec_r` of (4, 8) + produced two same-shaped arrays and a plausible (4, 8) result carrying the first + column's imaginary part in every lane. The pilot never tried this one -- its + perturbations are leading/trailing/extra/flattened, not length-1 -- so it is a gap + the annotation closes beyond what was measured. + """ + blocks = jnp.ones((DEGREES, BLOCKDIM, BLOCKDIM)) + with pytest.raises(TypeCheckError): + _block_matmul( + blocks, blocks, jnp.ones((DEGREES, BLOCKDIM)), jnp.ones((DEGREES, 1)) + ) + with pytest.raises(TypeCheckError): + _block_matmul( + jnp.ones((1, DEGREES, BLOCKDIM, BLOCKDIM)), + blocks, + jnp.ones((DEGREES, BLOCKDIM)), + jnp.ones((DEGREES, BLOCKDIM)), + ) + + +def test_the_free_axes_stay_free(): + """`rows` in `_matvec` and `cols` in `_matvec_T` are bound by ONE parameter. + + Nothing else in either signature carries them, so a (31, 16) operator with a (16,) + vector is a perfectly well-formed matvec and is accepted. Two of the pilot's six + acceptances on this pair were exactly that, and they are NOT defects -- the tool + perturbed a free axis. Asserted so that a later change does not "close" them by + cross-binding an axis the evidence does not support. + """ + assert _matvec(_mat(rows=ROWS - 1), jnp.ones((COLS,))).shape == (ROWS - 1,) + assert _matvec_T(_mat(cols=COLS - 1), jnp.ones((ROWS,))).shape == (COLS - 1,)