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:

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:

# 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:

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:

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 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 — _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 — 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, 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	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	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	drop matrix_mode='mv' at L94; simplify the module docstring
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
_normalize_matrix_mode	_misc.py:41
_mode_infix	_misc.py:172
_MM_STRIDE	_misc.py:38
_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
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)

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:

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 does not accept matrix_mode at all and hardcodes stride = _MM_STRIDE (L1328), while _jitsmm_cuda_kernel 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 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 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, and _jit_normal/_test_util.py:3 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