Affine/CAR: one channel-grouping spec, and kernels chosen from the weights

[cboulay]

Closes #198. Closes #210.

Both issues share a root: channel_clusters was doing two unrelated jobs, and the one it should never have had could change the answer.

#198 — a cluster hint could silently corrupt output

channel_clusters both built the weight matrix (for RereferenceKind/callable weights) and hinted its block structure for the matmul. A hint finer than W's true blocks made two input groups derive the same output indices, and the scatter assigned rather than accumulated, so the last group won. Wrong output of plausible shape, dtype and magnitude, with no warning.

Rather than validate the hint, this removes its ability to matter. Block structure is now always read off W itself (util/blockdiag.py), so no grouping a caller supplies can change the result. channel_groups keeps only the matrix-building job — which the issue correctly flagged as load-bearing on the kind/callable path.

The issue's channel_clusters=None workaround in ezmsg-learn#20 can go away.

Cost of always deriving structure: none in the common case. Contiguity detection is pure numpy in one pass over the nonzero mask; the connected_components search (and its scipy import) runs only for matrices above 1M weights whose blocks aren't already contiguous.

#210 — the block path was slower than dense below ~2048 channels

The block path gathered each cluster with fancy indexing and scattered the result back, per cluster, per message. Now:

• contiguous blocks slice into views on both sides and fill a preallocated buffer with matmul(out=) — no gather, no scatter, and since the blocks tile the output exactly, empty instead of zeros
• non-contiguous blocks take one whole-array permutation rather than a gather per cluster
• plan_block_matmul chooses between those and a dense matmul, and picks the merge granularity, from a cost model over weights-touched, chunk length and per-call overhead

min_cluster_size is gone — the min_cluster_size > n_ch trick for forcing dense was a workaround for a decision the library is better placed to make. kernel="auto"|"dense"|"blocks" is the explicit override.

case (float32)	before	after
256ch, 4×64 blocks, n=30	38.5 µs	7.4 µs
256ch, 4×64 blocks, n=3000	3391	190
1024ch, 16×64 blocks, n=30	155	20.7
1024ch, 16×64 blocks, n=3000	17022	1016

CommonRereference

mean-mode became x - (x @ project) @ spread — two skinny matmuls giving every channel its group's mean without a gather. Cost no longer depends on group contiguity or group count, and there's no Python loop per message. Leave-one-out collapses to N/(N-1) · (x - mean). median keeps the gather/scatter loop; a median has no matmul form.

case	before	after
global CAR, 256ch, n=30	27.0 µs	6.1 µs
4 groups, 256ch, n=30	40.0	6.0
32 groups, 256ch, n=30	161	6.4

One grouping spec

Three spellings collapse into channel_groups, a ChannelGroupSpec resolved and validated in util/channels.py:

spec	meaning
None	no grouping
"bank"	group by that metadata field
("array", "bank")	group by the tuple of those fields
[[0, 1], [2, 3]]	explicit index groups
fn(message, axis)	anything else

Groups are validated as in-range and pairwise disjoint — an overlap meant a channel referenced twice, or two weight blocks written to the same output.

Breaking changes

• AffineTransformSettings: channel_clusters → channel_groups (now only builds kind/callable weights), min_cluster_size removed, kernel added
• CommonRereferenceSettings: channel_clusters + cluster_by_field → channel_groups
• set_weights(..., recalc_clusters=) → recalc_structure= — ezmsg-learn's LRRTransformer calls this
• car_matrix / rereference_matrix: clusters= → groups=
• util.channels: channel_clusters_from_field / validate_channel_clusters → channel_groups_from_field / validate_channel_groups
• _state.clusters → _state.blocks; _find_block_diagonal_clusters, _max_cross_cluster_weight, _merge_small_clusters removed

Two behaviour changes worth a look on review:

1. Channels in no group pass through unchanged instead of being zeroed, matching car_matrix, which leaves them identity.
2. float32 in → float32 out. The old unconditional promotion to float64 doubled the bandwidth of every downstream stage.

Not mutating inputs

Message data may be a view shared with other branches of the graph. matmul(out=) only ever targets a buffer just allocated. The median path — the only one that writes elementwise — seeded its output with an asarray guarded by output is data, an identity test that silently fails for an ndarray subclass, where asarray returns a distinct object sharing the buffer. It now branches on dtype, where each branch is guaranteed to allocate. test_does_not_mutate_input marks inputs read-only across all 11 kernel paths, for both ndarray and a subclass; reverting the fix produces exactly the 4 expected failures.

Known limitation

The cost model counts weights touched, not access patterns, so it doesn't see that a many-block loop re-walks a strided view of a chunk too large to cache. Above ~2000 channels with several-thousand-sample chunks it picks a finer blocking than optimal — up to ~40% off the best merge, still ~3× faster than dense. Documented in blockdiag.py; kernel="dense" is the escape hatch.

Also not done: a batched kernel (reshape + one broadcast matmul) beat the slice loop ~2× at ≥512 channels with short chunks, but it needs equal-size blocks and is a wash at 256 channels. Left as a possible follow-up.

Testing

4026 passed, 6 skipped. New: tests/unit/test_blockdiag.py, the #198 regression test with genuinely non-contiguous blocks, and the input-mutation suite. benchmarks/benchmark_affine_kernels.py re-derives the cost-model constants and reports what the planner picks, so a stale constant shows up as a row where it chose the slower kernel.