# JITC mv/mm unification spec

Status: implemented
Branch: `remove-jit-mv/mm-mode`
Version: stays **0.2.0** — this lands inside the unreleased 0.2.0 work, so the
mv/mm split is edited out of the 0.2.0 changelog rather than deprecated in a
follow-up release.

## Goal

Make every JITC kernel — scalar, normal, uniform — draw **one** connectivity
matrix, the mv (32-lane) one. Once `mm` no longer exists as a distinct matrix,
delete `matrix_mode`, the `mat.mv` / `mat.mm` views and the `todense()` /
`tocsr()` ambiguity errors **completely** — no deprecation shim, no
compatibility keyword, no changelog trace. After this work, reading the tree
must give no indication that a mode selector ever existed.

Final state, stated as an invariant:

> For any `(weight, prob, seed, shape, corder)`, the matrix generated by
> `jits`, `jitsmv`, `jitsmm`, `binary_jitsmv`, `binary_jitsmm`, `jits_to_csr`
> and `jitsmv_dt2t` is bit-identical, on every backend (`numba`, `cuda_raw`),
> and equals today's `matrix_mode='mv'` matrix.

Anything that today distinguishes mv from mm is either deleted or converted to
32-lane. Kernel *organization* (AW-T4 grouping, bit-packing, warp reduction)
stays wherever it wins on throughput — the constraint is on the **drawn
matrix**, not on the launch shape.

**Interface — align to v0.1.2.** Signatures, keywords and the set of public
names return to the v0.1.2 tag. The materialization path returns there too: see
"The chunk width: keyed to the walked dimension" below, which removes the reason
0.2.0 had for deviating.

## Background: what mv/mm actually is

The light-RNG walk assigns each thread a residue class:

```
local_j = lane + STRIDE * q
```

- mv: one warp = one `(row, chunk_id)` task, `lane ∈ [0,32)`, `STRIDE = 32`.
- mm: one warp = 8 groups of 4 threads, one group = one `(row, chunk_id, col)`
  task, `sub_lane ∈ [0,4)`, `STRIDE = 4` (AW-T4).

The RNG is keyed by `(seed, row, chunk_id, lane)`. With `STRIDE = 4` and
`lane ∈ [0,4)`, both the stream keys and the residue classes differ, so mm
samples a *different* matrix. That is the entire difference — everything else
(`stationary_initial_q`, `fast_bounded_u32`, `light_rng_init_wpr`, chunking on
`shape[1]` with `target_chunks=4`) is already shared.

Consequently the fix is: **key the RNG by a 32-valued lane and step by 32**,
while keeping whatever thread grouping the kernel wants for parallelism.

## Design note: the 32-lane walk

**Decision: shape (A), warp-per-`(row, chunk, col)`, unconditionally.** Matching
the mv matrix is the only constraint that matters; small-`n` throughput is not
traded against it. No performance variant is kept in reserve, and no benchmark
gate stands between this work and landing.

**(A) Warp-per-(row, chunk, col) — mirrors `binary_jitsmv.cu` and the existing
`jitsmm_mv_*` kernels in `float_jitsmm.cu`.**

```
lane   = threadIdx.x & 31          // 0..31, one residue class each
row    = blockIdx.x * warps_per_block + warp_id
chunk  = blockIdx.y
col    = blockIdx.z
rng    = light_rng_init_wpr(seed, row, chunk, lane)
local_j = lane + 32 * q
… warp_reduce_sum → lane 0 atomics into output[row * n + col]
```

Correctness is immediate: identical keying, identical stride, identical
iteration order to `binary_jitsmv.cu`. `float_jitsmm.cu` already ships exactly
this kernel (`jitsmm_mv_notrans/trans_*`), which is the existence proof that a
32-lane matmat works.

AW-T4 exists because for small `n` it gives 8 independent column tasks per warp
instead of 1, so (A) may be slower there. That is accepted: the AW-T4 grouping
is deleted outright rather than preserved behind a replay scheme, and `jitsmm` /
`binary_jitsmm` become exactly "`n` independent mvs on one matrix" in the launch
shape as well as in the drawn matrix.

## The chunk width: keyed to the walked dimension

**Decision: derive `chunk_size` from the dimension the walk actually runs over,
not from the caller's `shape[1]`. `_materialize_params` is then unnecessary and
is deleted.**

An earlier draft of this spec kept 0.2.0's `_materialize_params` helper and
claimed v0.1.2 was buggy here. Measured against the released 0.1.2 that claim is
**false**, and the fix belongs one level down. Both points are recorded here
because the wrong version was acted on first.

### What `_materialize_params` was

It returned the `(shape, transpose)` pair `jits` had to be called with so that
the dense form matched what `mat @ v` computes:

```python
class JITCScalarR:  def _materialize_params(self): return self.shape,       False
class JITCScalarC:  def _materialize_params(self): return self.shape[::-1], True
```

The column-oriented classes need the swapped pair because that is what their
`__matmul__` passes — in v0.1.2 exactly as in 0.2.0:

```python
# JITCScalarC.__matmul__, unchanged since v0.1.2
jitsmv(..., shape=self.shape[::-1], transpose=True, corder=self.corder)
```

### v0.1.2 was not broken; 0.2.0 broke it

The two calls produce the same generation dimensions:

```python
out_shape      = shape if not transpose else shape[::-1]
n_rows, n_cols = out_shape if corder else out_shape[::-1]
```

For `shape=(M, N), transpose=False` and `shape=(N, M), transpose=True` the
output shape is `(M, N)` either way, so `(n_rows, n_cols)` is identical. The
*only* thing that differed was the chunk width, which 0.2.0 derived from the
caller's `shape[1]` — `N` in one call, `M` in the other — and `chunk_id` keys
the RNG stream.

v0.1.2's kernel had no chunking at all (cuRAND Philox, `curand_init(seed, row, 0)`,
one thread per row), so its matrix did not depend on the generation shape.
Measured on the released 0.1.2:

```
0.1.2:  jits((M,N),transpose=False) vs jits((N,M),transpose=True)   identical = True
        mat @ v == mat.todense() @ v   for JITCScalarC/R, JITCUniformC/R, both corder
                                                                    all True
0.2.0:  same comparison                                             identical = False
        nnz 66 vs 61 (corder=True), 71 vs 73 (corder=False)
```

So `_materialize_params` was not a fix for a v0.1.2 defect — it was a
*compensation* for a shape dependence 0.2.0 had just introduced.

### The fix

`chunk_size` is derived from the **walked** dimension:

```python
def _walk_length(shape, transpose, corder):
    n_rows, n_cols = (shape[1], shape[0]) if transpose else (shape[0], shape[1])
    return n_cols if corder else n_rows

def _chunk_size(walk_length):
    return max(1, cdiv(walk_length, _TARGET_CHUNKS))   # _TARGET_CHUNKS = 4
```

`_walk_length` is the same for both `(shape, transpose)` pairs, so the drawn
matrix depends only on the generated matrix's own geometry. Consequences:

- `_materialize_params` collapses to `(self.shape, False)` for the row- *and*
  column-oriented classes and is deleted; `todense()` / `tocsr()` return to
  v0.1.2's bodies verbatim.
- `n_chunks` stays exactly 4, so the CSR count pass's `chunk_counts[n_rows,
  n_chunks]` buffer, the `grid.y <= 65535` headroom and the per-lane walk length
  are all unchanged. A fixed chunk *constant* would have achieved the same
  shape-independence but at a cost: at 1M x 1M the CSR buffer grows from 15 MB
  to 3.6 GB (C=1024) or 116 GB (C=32), and small constants leave each lane
  emitting at most one element per stream.
- The chunk width is no longer a caller-visible knob: `chunk_size` and
  `target_chunks` are removed from `jits_to_csr`, `jits_csr_count_p_call`,
  `jits_csr_fill_p_call` and `jitsmv_dt2t_p_call` (and the `jitn*` / `jitu*`
  siblings). `jits_to_csr`'s parameter list is then identical to v0.1.2's.

Pinned by `test_generation_is_shape_pair_independent` (the two pairs draw the
same matrix) and `test_todense_matches_matmul` (`mat.todense() @ v == mat @ v`
for every class and `corder`).

## Phase 1 — `binary_jitsmm.cu` → 32-lane

Convert [binary_jitsmm.cu](../../brainevent/_jit_scalar/binary_jitsmm.cu) in
place. Keep the file, the `pack` kernel (`packed[n][ceil(k/32)]`, unchanged),
and the six `@BE` entry points (`notrans_{f32,f64,f16,bf16}`, `trans_{…}`) so
no Python dispatch name moves.

Edits, per kernel (`_notrans_f32_kern`, `_trans_f32_kern`, and both macros
`DEFINE_BINARY_JITSMM_NOTRANS` / `_TRANS`):

| Current | New |
| --- | --- |
| `int sub_lane = lane & 3; int group = lane >> 2;` | delete; use `lane` directly |
| `row = warp_task * AW_T4_GROUPS_PER_WARP + group` | `row = row_block * warps_per_block + warp_id` |
| `light_rng_init_wpr(seed0, row, chunk_id, sub_lane)` | `…, lane)` |
| `local_j = sub_lane + AW_T4_GROUP_SIZE * q` (×2 sites) | `local_j = lane + 32U * q` |
| `group4_reduce_sum_f32(acc, group)` | `warp_reduce_sum_f32(acc)` |
| `if (sub_lane == 0)` | `if (lane == 0)` |

And in the four launch wrappers (`notrans_f32`, `trans_f32`,
`FFI_BINARY_JITSMM`):

| Current | New |
| --- | --- |
| `rows_per_block = warps_per_block * AW_T4_GROUPS_PER_WARP` | delete |
| `row_group_blocks = (m + rows_per_block - 1) / rows_per_block` | `row_warp_blocks = (m + warps_per_block - 1) / warps_per_block` |

`dim3 blocks(row_warp_blocks, n_chunks, n)` — the z-dimension over columns is
what makes this mm rather than mv, and it stays. `group4_reduce_sum_f32/f64`
become unused and are deleted along with `AW_T4_GROUP_SIZE` /
`AW_T4_GROUPS_PER_WARP`. The stale header comment (`n <= 32`, "AW backend with
4-thread WPR groups") is rewritten.

Python side, same phase:

- [binary.py:866](../../brainevent/_jit_scalar/binary.py#L866) —
  `_jitsmm_numba_kernel`: `stride = _MM_STRIDE` → 32, and update the comment at
  lines 867–869 that explains the 4-vs-32 divergence.
- [binary.py:970–973](../../brainevent/_jit_scalar/binary.py#L970) — the comment
  claiming gradients must delegate to `float.jitsmm` in mm mode is now wrong;
  rewrite.

### Phase 1 tests (the batch=1 check)

New tests in [binary_test.py](../../brainevent/_jit_scalar/binary_test.py),
parametrized over every available backend × `transpose` × `corder` × dtype:

1. `test_binary_jitsmm_batch1_matches_binary_jitsmv` — for `B` of shape
   `(k, 1)`, `binary_jitsmm(w, p, B, seed, …)[:, 0]` equals
   `binary_jitsmv(w, p, B[:, 0], seed, …)`. Exact `==` for f32/f64;
   `allclose` for f16/bf16, where atomic ordering perturbs the low bits.
2. `test_binary_jitsmm_columns_are_independent_mvs` — for `B` of shape
   `(k, n)`, column `c` of the result equals `binary_jitsmv(…, B[:, c])` for
   every `c`. This is the strong form: mm is literally "n independent mvs on
   one matrix".
3. `test_binary_jitsmm_structure_matches_jits` — `binary_jitsmm` with an
   all-ones `B` reproduces the row/column sums of the `jits` matrix. Written
   against the post-unification single `jits`, so it survives Phase 3 unchanged.

The existing `test_binary_jitsmm_*` reference tests compare against
`float.jitsmm`, which is still in mm mode after Phase 1 alone — so **Phases 1
and 2 land as a single commit** for `_jit_scalar`. They are the same semantic
change and splitting them would leave the tree red.

## Phase 2 — propagate to `float`, `csr`, `dt2t` (`_jit_scalar`)

Nothing here needs new kernels; every mv variant already exists.

### `float_jitsmm.cu`

`jitsmm_mv_notrans/trans_*` (8 `@BE` symbols) already implement shape (A).
Promote them:

- rename `jitsmm_mv_*` → `jitsmm_*`
- delete the AW-T4 `jitsmm_notrans/trans_*` kernels, their macros and their 8
  `@BE` symbols
- net: 16 → 8 `@BE` symbols

The launch parameters differ between the two families (`FFI_JITSMM(…, 1, m)`
vs the AW-T4 form); take the `mv` family's verbatim.

### `float_jits.cu`

Delete `jits_mm_aw_t4_notrans/trans_*` (8 `@BE` symbols) and the
`_jits_mm_aw_t4*_kern` bodies. Keep `jits_mv_*`, renaming to drop the `_mv`
infix so `_jits_cuda_kernel` builds one unconditional name. 16 → 8.

### `csr.cu`

Delete `count_chunks_{notrans,trans}_mm_aw_t4_*` and
`fill_{notrans,trans}_mm_aw_t4_*` (16 `@BE` symbols), the four
`_*_mm_aw_t4_f32_kern` bodies, and the mm arms of the four `DEFINE_*_WRAPPER`
expansions. 32 → 16.

### `dt2t.cu`

Already mv-only. No kernel change; the comments that say "dt2t always
materializes the mv matrix" become vacuous and are simplified.

### Python

| File | Change |
| --- | --- |
| [float.py](../../brainevent/_jit_scalar/float.py) | drop `matrix_mode` from `jits`, `jitsmm`, `jitsmm_p_call`, `_jits_numba_kernel`, `_jits_cuda_kernel`, `_jitsmm_cuda_kernel`, every JVP/transpose/batching forwarder, and both `_benchmark_data` dicts; stride becomes an unconditional 32 (L426, L1328); kernel names lose the mode infix (L511, L1415) |
| [csr.py](../../brainevent/_jit_scalar/csr.py) | drop `matrix_mode` from `jits_to_csr`, `jits_csr_count_p_call`, `jits_csr_fill_p_call` and the four kernel builders; drop the `_mode_infix` calls (L85, L285); stride unconditional (L133, L333) |
| [dt2t.py](../../brainevent/_jit_scalar/dt2t.py) | drop `matrix_mode='mv'` at L94; simplify the module docstring |
| [binary.py](../../brainevent/_jit_scalar/binary.py) | handled in Phase 1 |

## Phase 3 — delete the mode selection, restore the 0.1.2 interface

### Deletions — nothing left behind

| Symbol | Location |
| --- | --- |
| `MatrixMode` | [_typing.py:82](../../brainevent/_typing.py#L82) |
| `_normalize_matrix_mode` | [_misc.py:41](../../brainevent/_misc.py#L41) |
| `_mode_infix` | [_misc.py:172](../../brainevent/_misc.py#L172) |
| `_MM_STRIDE` | [_misc.py:38](../../brainevent/_misc.py#L38) |
| `_MV_STRIDE` | keep the constant, rename to `_LANE_STRIDE`, drop the mv/mm rationale from its comment |
| `_JITCScalarModeView` + the `.mv` / `.mm` properties | [main.py:40,404,410](../../brainevent/_jit_scalar/main.py#L40) |
| 3× `NotImplementedError` in `todense` / `tocsr` | main.py L399, L713, L1215 |

`matrix_mode` becomes an unknown keyword — passing it raises the ordinary
`TypeError` from Python, with no special-casing, message or warning. No
deprecation path is provided.

A grep gate closes the phase:

```
grep -ri "matrix_mode\|_mode_infix\|MatrixMode\|_MM_STRIDE\|mm_aw_t4\|aw.t4\|\.mv\.\|\.mm\." \
     brainevent/ changelog.md
```

must return nothing outside two allowances:

| Allowed hit | Why |
| --- | --- |
| `brainevent/_event/compact.py` | local `def mv(...)` closures, unrelated to JITC |
| `changelog.md:675` | 0.0.x naming convention for the `*mv` / `*mm` function *names*, which stay |

**`docs/` is not in the gate.** Pre-existing documents there
(`docs/superpowers/plans/`, earlier `docs/specs/`) are dated records of past
decisions; they are neither rewritten nor whitelisted, they are simply outside
the scope of this work. The invariant is about the code tree. `changelog.md` is
the one document that *is* edited, because it describes the current unreleased
version rather than a past decision.

### Restorations (target = v0.1.2 signatures)

```python
def jits(weight, prob, seed, *, shape, transpose=False, corder=True, backend=None) -> Data
def jits_to_csr(weight, prob, seed, *, shape, corder, backend=None)
mat.todense()   # works, no mode
mat.tocsr()     # works, no mode
mat.tocsc()     # works, no mode
mat.tocoo()     # works, no mode
```

`jitsmv`, `jitsmm`, `binary_jitsmv`, `binary_jitsmm`, `jitsmv_dt2t` and all
`*_p` primitives already match v0.1.2 once `matrix_mode` is dropped.

### Interface conformance check

`test_public_api_matches_0_1_2` asserts via `inspect.signature` that the
parameter lists of the public JITC entry points equal a table transcribed from
the v0.1.2 tag. This discharges the "接口与 0.1.2 一致" requirement mechanically
rather than by review.

The check is **signature-level only**; the behaviour behind those signatures is
pinned separately by `test_generation_is_shape_pair_independent` and
`test_todense_matches_matmul`.

Implemented per-module rather than as one cross-package test, to respect the
co-location rule (AGENTS.md #11): `float_test.py` checks the `jit*` / `jit*mv` /
`jit*mm` parameter lists, `csr_test.py` the `*_to_csr` one, `dt2t_test.py` the
`*mv_dt2t` one. All match v0.1.2 exactly. Only the low-level `*_p_call`
primitives still differ, and only in parameter *names* (`w0`/`w1` →
`w_loc`/`w_scale`, `indptr` → `offsets`).

### Test updates

| File | Work |
| --- | --- |
| `csr_test.py` | 33 sites; drop the `MATRIX_MODES` parametrization, delete `test_to_csr_requires_matrix_mode`, and **invert** `test_mv_mm_differ` (L117) into `test_mv_mm_identical` in Phase 2, then delete it in Phase 3 once the concept is gone |
| `float_test.py` | 20 sites; delete `test_jits_requires_matrix_mode`, collapse the mv/mm duplicate blocks (L75–L136, L153–L241) into one |
| `main_test.py` | `.mv`/`.mm` view tests → direct `todense()` / `tocsr()`; add `test_column_oriented_todense_matches_matmul` pinning the `_materialize_params` behaviour |
| `dt2t_test.py` | 4 sites, mechanical |
| `_data_test.py` | L815–823 `m.mv.todense()` → `m.todense()` |
| `binary_test.py` | keep the Phase 1 additions |

### Changelog

0.2.0 is unreleased on this branch, so the mv/mm split is **edited out of the
0.2.0 notes**, not recorded as a reversal. Sites in
[changelog.md](../../changelog.md):

| Line | Action |
| --- | --- |
| 44 | delete the `mat.todense()` → `mat.mv.todense()` breaking row |
| 45 | delete the `matrix_mode` required-keyword breaking row |
| 63–67 | delete the "**`mv` and `mm` are genuinely different matrices**" paragraph, and fix the "Three of these deserve a note" lead-in that counts it |
| 113–119 | delete the `mat.mv` / `mat.mm` Added bullet — but **relocate its last sentence** (the column-oriented swapped generation shape) into a new *Fixed* entry, since that behaviour stays 0.2.0-aligned and must remain documented: "`todense()` on `JITCScalarC` / `JITCNormalC` / `JITCUniformC` now generates with the matvec's shape; the light kernels' `chunk_size` derives from `shape[1]`, so the old direct materialization drew a different matrix than `mat @ v`." |
| 128 | delete the `matrix_mode` on-CSR-entry-points Added bullet |
| 263 | drop `_normalize_matrix_mode`, `_MV_STRIDE` / `_MM_STRIDE`, `_mode_infix` and `MatrixMode` from the consolidated-internals list; keep `_normalize_chunk_size`, `_is_static_zero`, `_n_chunks` |
| 675 | leave alone — 0.0.x naming convention for `*mv` / `*mm` function names |

A new *Fixed* entry also covers the three latent bugs in the next section.

## Latent bugs fixed by this work

Each is real on the current tree, each disappears under unification, and each
gets a regression test that keeps it from coming back (AGENTS.md rule 4).

### B1 — `numba` `jitsmm` silently ignores `matrix_mode`

[`_jitsmm_numba_kernel`](../../brainevent/_jit_scalar/float.py#L1298) does not
accept `matrix_mode` at all and hardcodes `stride = _MM_STRIDE` (L1328), while
[`_jitsmm_cuda_kernel`](../../brainevent/_jit_scalar/float.py#L1415) honours it.
So `jitsmm(..., matrix_mode='mv', backend='numba')` returns the **mm** matrix
and disagrees with `backend='cuda_raw'` — a silent CPU/GPU divergence of exactly
the kind 0.2.0's parity work set out to remove.

*Regression test:* `numba` and `cuda_raw` `jitsmm` agree (CUDA-gated); satisfied
automatically once one stride remains.

### B2 — `vmap(jitsmv)` draws a different matrix than `jitsmv`

[`_jitsmv_batching`](../../brainevent/_jit_scalar/float.py#L1099) forwards to
`jitsmm_p_call` with a hardcoded `matrix_mode='mm'` (L1099 and L1113). Since
`jitsmv` is mv, `vmap(jitsmv)(V)` and a Python loop of `jitsmv` over the columns
of `V` produce different numbers. The `jitsmm_mv_*` CUDA kernels exist precisely
to fix this — their header comment says "used by `vmap(jitsmv)`" — but they were
never wired up: no caller anywhere in the repo, tests included, ever passes
`matrix_mode='mv'` to `jitsmm`.

*Regression test:* `vmap(jitsmv)(V)` equals `stack([jitsmv(v) for v in V.T])`.

### B3 — same defect on the binary path

[`_jitsmv_batching` in binary.py](../../brainevent/_jit_scalar/binary.py#L621)
forwards to `binary_jitsmm_p_call` (L621, L634). `binary_jitsmm` has no mode
parameter at all — its `.cu` is AW-T4 unconditionally — so
`vmap(binary_jitsmv)` has the same divergence with no way to opt out.

*Regression test:* `vmap(binary_jitsmv)(V)` equals the `binary_jitsmv` loop.

### B4 — the same three defects, replicated in `_jit_normal` and `_jit_uniform`

`_jit_normal/binary.py` pins `matrix_mode='mm'` at L871, L891, L901, L911 (four
sites — gradient rules as well as batching), and `_jit_uniform/binary.py` at
L1142 and L1213. Same regression tests, per package.

## Phase 4 — propagate to `_jit_normal` and `_jit_uniform`

The two packages are structural mirrors of `_jit_scalar` (`_JITCNormalModeView`
/ `_JITCUniformModeView`, the same `csr.py` / `float.py` / `dt2t.py` split, the
same `*mm.cu` AW-T4 kernels). Apply Phases 1–3 file-for-file:

| Scalar | Normal | Uniform |
| --- | --- | --- |
| `binary_jitsmm.cu` | `binary_jitnmm.cu` | `binary_jitumm.cu` |
| `float_jitsmm.cu` | `float_jitnmm.cu` | `float_jitumm.cu` |
| `float_jits.cu` | `float_jitn.cu` | `float_jitu.cu` |
| `csr.cu`, `dt2t.cu` | same names | same names |

Extra work not present in scalar:

- The six extra `matrix_mode='mm'` pin sites of B4.
- **`stride_for_mode()` and the cross-package `_test_util` dependency.**
  `stride_for_mode(mode)` (`'mv'→32`, `'mm'→4`) is the pure-numpy stride the
  parity tests use to recompute the connectivity by hand. It collapses to the
  constant 32 and is deleted, along with the `matrix_mode='mv'` parameter of
  `dense_normal_reference()` / `dense_uniform_reference()`.

  It is currently defined **once**, in
  [_jit_uniform/_test_util.py:63](../../brainevent/_jit_uniform/_test_util.py#L63),
  and [_jit_normal/_test_util.py:3](../../brainevent/_jit_normal/_test_util.py#L3)
  reaches across package boundaries for it together with `conn_length`,
  `default_chunk_size`, `hash_uniform01` and `iter_edges`. **Each package tests
  on its own**: that import is removed and `_jit_normal/_test_util.py` becomes
  self-contained, carrying its own copy of the RNG-walk helpers it needs
  (`mix32`, `fast_bounded_u32`, `light_rng_next/init`, `stationary_initial_q`,
  `conn_length`, `default_chunk_size`, `iter_edges`, plus `hash_uniform01` since
  `hash_normal01` is built on it). Duplication is the intended outcome — a
  reference implementation that two packages share cannot catch a divergence
  between them, which is exactly the class of bug (B1–B4) this work exists to
  remove.

  After the split, no `_jit_*` test module imports from another `_jit_*`
  package; the only shared test import left is the repo-wide
  `brainevent._test_util` (`allclose`, `requires_gpu`, `gen_events`), which is
  backend-agnostic plumbing, not a reference implementation.
- The weight streams differ (`hash_normal01`, `hash_uniform01`) but are keyed by
  `(seed, row, col)`, independent of lane — so the stride change moves the
  *structure* only, exactly as in scalar. No new weight-hash work.

Order: normal first (closer to scalar), then uniform.

## Performance: explicitly out of scope

`jitsmm` / `binary_jitsmm` lose the AW-T4 grouping in Phases 1–2 and may regress
at small `n`. No benchmark phase gates this work and no alternative kernel shape
is kept: correctness — one matrix, on every backend, for every entry point — is
the whole deliverable. If a regression later proves material in practice, it is
new work against the then-current tree, not a variant carried in this spec.

## Commit plan

Landed as **one** commit rather than the five below. The phases are still the
right way to *read* the change, but they cannot be split into green commits:
`_misc.py`, `_typing.py` and `_data_test.py` are shared by all three families,
so `_data_test.py`'s parametrized materialization test is red for any package
that has been converted while its siblings have not, and deleting `MatrixMode` /
`_mode_infix` requires all three packages to have stopped using them.

The phase order was still followed during implementation, with the invariant
re-verified on `numba` *and* `cuda_raw` after each package:

1. Phases 1+2 — `_jit_scalar` CUDA + Python
2. Phase 3 — `_jit_scalar` `main.py` (folded in: `_JITCScalarModeView` calls
   `jits(matrix_mode=...)`, so it cannot outlive Phase 2)
3. Phase 4a/4b — `_jit_normal`, `_jit_uniform`
4. Shared deletions (`_misc.py`, `_typing.py`), `_data_test.py`, changelog

CUDA-only assertions stay `skipif`-guarded, so the work is verifiable on CPU via
the `numba` backends; the CUDA/numba parity tests pin the two implementations
together.

## Deviations from this spec, as implemented

- **`float_test.py`'s "duplicate mv/mm blocks" are not duplicates.** The spec
  calls for collapsing L75–136 and L153–241 into one block each. They are the
  `jitsmv` and `jitsmm` forward/gradient tests — different entry points that
  happened to pass different `matrix_mode` values. Both are kept; only the
  keyword is dropped.
- **Each package gained a self-contained `_test_util.py`, including a new one
  for `_jit_scalar`** (`dense_scalar_reference`). The pure-numpy replay of the
  32-lane walk pins the *drawn matrix* rather than merely the agreement of the
  kernels with each other, and duplicating it per package means a divergence
  between families cannot hide behind a shared helper.
- **The interface conformance check is per-module, not one cross-package test**,
  to respect the co-location rule (AGENTS.md #11): `float_test.py` asserts the
  `jit*` / `jit*mv` / `jit*mm` parameter lists, `csr_test.py` the `*_to_csr`
  one, `dt2t_test.py` the `*mv_dt2t` one. All match v0.1.2 exactly. The
  `*_p_call` primitives keep 0.2.0's extra `chunk_size` / `target_chunks`
  keyword-only arguments, which have defaults and are unrelated to mv/mm.

## Edge cases to cover in tests

- `n = 1` (the batch=1 identity), `n = 32`, `n = 33` (crosses the packed-word
  boundary), `n > 65535` must still hit the existing grid-overflow abort.
- `k` not a multiple of 32 — the last packed word is partially valid.
- `prob = 0` (empty matrix) and `prob = 1` (`clen = 2`, the `cl < 2U` clamp).
- `chunk_size` boundaries: `shape[1] < target_chunks`, `shape[1] = 1`.
- `corder ∈ {True, False}` × `transpose ∈ {True, False}` — all four, since
  notrans/trans are separate kernels.
- dtypes `f32, f64, f16, bf16` — f16/bf16 accumulate in f32 and need `allclose`.
- Column-oriented matrices (`JITCScalarC`) for the `_materialize_params` path.
- `vmap` over both axis 0 and axis 1 (the two arms of each batching rule).

## Scope accounting

All counts measured on this tree (`grep -o`, occurrences not lines), not
estimated.

### `@BE` CUDA symbols

Per package, before → after. The three packages are exact structural mirrors, so
one table covers all three (`jits`↔`jitn`↔`jitu`).

| File | Before | After | Delta |
| --- | --- | --- | --- |
| `csr.cu` | 32 | 16 | −16 (the `*_mm_aw_t4_*` half) |
| `float_jits.cu` | 16 | 8 | −8 (`jits_mm_aw_t4_*`) |
| `float_jitsmm.cu` | 16 | 8 | −8 (AW-T4 family; `jitsmm_mv_*` promoted) |
| `binary_jitsmm.cu` | 9 | 9 | 0 (rewritten in place, `pack` kept) |
| `binary_jitsmv.cu` | 9 | 9 | 0 |
| `float_jitsmv.cu` | 8 | 8 | 0 |
| `dt2t.cu` | 8 | 8 | 0 |
| **per package** | **98** | **66** | **−32** |
| **×3 packages** | **294** | **198** | **−96** |

### `matrix_mode` occurrences

| Scope | Before | After |
| --- | --- | --- |
| `_jit_scalar`, product code | 98 | 0 |
| `_jit_normal`, product code | 107 | 0 |
| `_jit_uniform`, product code | 109 | 0 |
| `_misc.py` + `_typing.py` | 18 | 0 |
| **product total** | **332** | **0** |
| `_jit_scalar` tests | 69 | 0 |
| `_jit_normal` tests | 82 | 0 |
| `_jit_uniform` tests | 84 | 0 |
| **test total** | **235** | **0** |

### API surface and test scaffolding

| | Before | After |
| --- | --- | --- |
| Public functions with a `matrix_mode` parameter | 8 per package, 24 total | 0 |
| ↳ `float.py` | `jits`, `jitsmm`, `jits_p_call`, `jitsmm_p_call` | — |
| ↳ `csr.py` | `jits_to_csr`, `jits_csr_count_p_call`, `jits_csr_fill_p_call` | — |
| ↳ `dt2t.py` | `jitsmv_dt2t` | — |
| `NotImplementedError` ambiguity raises in `main.py` | 3 per package, 9 total | 0 |
| `.mv` / `.mm` attribute sites, `main_test.py` | 32 / 38 / 38 (scalar/normal/uniform) | 0 |
| `.mv` / `.mm` attribute sites, `_data_test.py` | 6 | 0 |
| `MATRIX_MODES` parametrization sites | 26 | 0 |
| `test_*_requires_matrix_mode` | 6 (3× `float_test`, 3× `csr_test`) | 0 (deleted) |
| `stride_for_mode` | 1 definition (`_jit_uniform/_test_util.py:63`) + 1 cross-package import (`_jit_normal/_test_util.py:8`) | 0 |
| Cross-package `_jit_*` test-helper imports | 1 (`_jit_normal/_test_util.py` → `_jit_uniform/_test_util.py`, 5 names) | 0 — each package self-contained |