# Package-wide simplification spec

Status: in progress
Branch: `worktree-simplify-package-cleanup`

## Progress

Landed (each as its own commit):

- **Tier A** — A1–A8, A11, A12, A13, plus one bonus finding (a `TestCompactOnlyVector`
  class defined twice, the first shadowed and never executed).
- **Tier B** — B1 (`_dtype_sfx` ×31), B2/B3/B4 (the `_jit_*` chunking, mode and
  stride helpers).

Deliberately skipped, with reasons:

- **A9 (`_registered_platforms`)** — reported as always equal to `set(self._kernels)`.
  True on the happy path only: the field records *successful* lowering registration,
  so if `_register_fallback_lowering` raises, the two diverge and a retry would
  re-attempt registration. Removing it changes error-path behaviour for the sake of
  one set. Not worth it.
- **A10 (write-only FFI handler attributes)** — entangled with C2, which rekeys the
  same cache. Should be done together, not separately.
- **Pallas compat shim** — see "Adjudicated conflicts" below; needs JAX-version
  verification first.
- **`dtype_suffix` raising on unmapped dtypes** — recommended by the reviewer, but it
  changes dispatch behaviour rather than preserving it. Documented and pinned by test
  instead; raising is a follow-up decision, not a cleanup.

Not started: B5–B17, all of Tier C, all of Tier D.

Verification at this point: full fast suite 1574 passed / 0 failed / 201 skipped;
mypy 204 errors, down from the 208-error baseline (no regression).

## Motivation

A four-angle quality review (reuse, simplification, efficiency, altitude) of the
whole `brainevent` package (~53k lines of non-test source across ~60 modules)
produced 45 findings. This spec records them, tiers them by risk, and defines the
verification bar for landing them.

The review targeted **quality only** — reuse, simplification, efficiency, and
altitude. It did not hunt for correctness bugs. A handful of probable bugs
surfaced incidentally and are recorded in "Out of scope" below rather than fixed
here.

## Baseline (measured before any change)

- `mypy` (CI gate config): **208 errors in 23 files** (75 source files checked).
  The gate is *not* currently clean, so the bar for this work is "no new errors";
  reductions are a bonus.
- `import brainevent`: ~1.94 s cold; brainevent's own module bodies 0.95–1.28 s.
- `detect_cpp_toolchain()`: 8.4 ms/call, uncached.
- Header glob + SHA-256 per compile-cache lookup: ~75 ms, paid even on a cache hit.
- Re-lowering an identical CSR kernel: ~100–140 ms each time (full numba recompile).

## Verification bar

Per `AGENTS.md`:

- Full test suite must pass (CI runs `pytest -m ""`, including `slow` variants).
- `mypy` must not regress past the 208-error baseline.
- Meaningful tests for edge cases and critical paths, not trivial line coverage.
- Tests co-located as `foo_test.py` beside `foo.py` (suffix style, never a
  `tests/` directory, never a `test_*.py` prefix).

## Adjudicated conflicts

Three findings were contested between reviewers and were resolved by direct
inspection:

1. **`_dtype_sfx` (31 copies).** One reviewer declined to flag it, arguing each
   table maps to a different `.cu` symbol family. Verified false: all 31 bodies
   are byte-identical dtype→suffix maps (`np.dtype` vs `jnp.dtype` keys resolve to
   the same objects). Only the *base kernel name* differs, and that stays at the
   call site. `_jit_uniform/binary.py:32` already imports the table cross-module,
   proving it is shareable. **Verdict: consolidate.**
2. **`cdiv`.** One reviewer called it dead code to delete; another wanted it reused.
   Verified: genuinely unreferenced outside `_misc.py`, while three `_jit_*/csr.py`
   modules hand-roll its body as `_n_chunks`. **Verdict: keep and adopt it.**
3. **Pallas compat shim.** `pallas_triton_params` / `pallas_mosaic_tpu_params` are
   defined, exported in `__all__`, and used nowhere; the two real Pallas call sites
   pass `backend=impl_backend`. Whether wiring the shim up is correct depends on
   JAX-version behavior not yet verified. **Verdict: defer pending verification;
   do not delete and do not blind-wire.**

## Tier A — mechanical, no behavior change

| # | Site | Change |
|---|------|--------|
| A1 | `_misc.py:244–887` | Delete 9 never-called helpers (~443 lines): `is_known_type` (dup of live `_event/base.py:51`), `_block_csr_tocsr`, `_block_csr_tocoo`, `estimate_block_size`, `count_blocks`, `_count_blocks`, `_nonzero_blocks`, `_coordinate_index_dtype`. Keep `cdiv` (see A-adjacent B3). |
| A2 | package-wide | Remove 18 `ruff F401` unused imports and `F841` unused locals (`_dense/plasticity_binary.py:300,632`; `n_rows` in three `_jit_*/csr.py`). |
| A3 | `_jit_{normal,scalar,uniform}/csr.py` | `row_counts` is `device_get`'d twice; do it once. Removes a redundant device→host sync per conversion. |
| A4 | `_data.py:41` | `_buffer_registry.update(...)` is immediately followed by a loop that registers each name anyway. Delete. |
| A5 | `_csr/dt2t.py:341` | `_csrmv_dt2t_transpose_rule` is never registered; delete. |
| A6 | `_jit_uniform/binary.py:416` | `_spike_sfx` dict, sole occurrence in repo; delete. |
| A7 | `_csr/float.py:775–777` | `_hetero` arm is unreachable inside `if is_homo:`; inline the `_homo` literal. |
| A8 | `_op/main.py:437` | Invert `if ...: pass / else:` to drop a nesting level. |
| A9 | `_op/main.py:184,372` | `_registered_platforms` is always `set(self._kernels)`; delete the field. |
| A10 | `_op/numba_ffi.py:600`, `_op/numba_cuda_ffi.py:355,836` | 7 handler attributes are stored and never read; drop from the three `__init__` signatures (keep as locals for the memoisation key). |
| A11 | `_op/main.py:1217–1273` | `catch_errors` arms duplicate the timing call and the 13-line success record; unify under one `try`. |
| A12 | `_op/numba_ffi.py:97` | Header scan runs 1576 regex attempts with no early exit; add `break`, make lazy. |
| A13 | `_csr/float.py:868` | Stray `TODO` above a complete implementation; delete. |

## Tier B — consolidate duplicates

| # | Duplication | Target |
|---|-------------|--------|
| B1 | `_dtype_sfx` ×31 (+ `spk_suffix` inline ×11, `homo_suffix` ×82) | `dtype_suffix()` / `spike_suffix()` in `_op/util.py`, raising on unmapped dtype instead of silently falling back to `_f32` |
| B2 | `_normalize_chunk_size` ×4 | `_misc.py` (RNG-stream keying hazard is documented in only one copy) |
| B3 | `_is_static_zero`, `_n_chunks`, `_mode_infix` ×3 | `_misc.py`; `_n_chunks` reimplemented on existing `cdiv` |
| B4 | `_normalize_matrix_mode` ×3 | `_misc.py`; `MatrixMode` alias to `_typing.py` |
| B5 | Operator-forwarding block ×4 subclasses (~950 lines) | `DataRepresentation` in `_data.py`. Fixes drift: `%` currently works on `fcn`/`JITCMatrix` but not `csr`/`dense` |
| B6 | `_binary_op`/`_binary_rop` dispatcher ×20 | `DataRepresentation._elementwise_binary`. Fixes the `"mul with object of shape"` message copied verbatim into all 20 regardless of operator |
| B7 | CSR `*_p_call` validation preamble ×8 | `check_csr_matmul_shape()` in `_misc.py`, mirroring the existing `check_fixed_conn_num_shape` |
| B8 | `__matmul__`/`__rmatmul__` ×4 | Make `EventRepresentation.__matmul__`/`__rmatmul__` concrete |
| B9 | XLA FFI callback body ×3 | `_handle_metadata_extension()` + `_extract_buffers()` in `numba_ffi.py` |
| B10 | `_jax_x64_enabled` ×6 test modules (+3 variants) | One `enable_x64()` in `_test_util.py` |
| B11 | `_csr/main_test.py:43,59` | Import `gen_events`/`ones_like` from `_test_util.py` |
| B12 | `coo_to_csc_index` / `coo2csr`, each ×2 (np/jnp branches) | One `_group_by_axis` with the `mod = np if ... else jnp` selector already used at `_misc.py:1181` |
| B13 | `_op/benchmark.py:631–747` | `_format_hierarchical` delegates to `_format_vary_by` (verified byte-identical output) |
| B14 | `kernix_pipeline.py:269–319` vs `474–513` | One `_build_and_cache(...)` |
| B15 | `_csr/main.py:93–175` | Replace 7-way positional closure factory with module-level defs, mirroring `_fcn/main.py:102–129` |
| B16 | `_csr/main.py:58`, `_csr/binary.py:75` | Two interchangeable workspace types (frozen dataclass vs namedtuple) that convert into each other mid-autodiff; keep one |
| B17 | `_jit_*/main.py` R/C classes | Hoist 4 methods each onto the shared base using `type(self)` |

## Tier C — efficiency

| # | Site | Change | Measured cost today |
|---|------|--------|---------------------|
| C1 | `_op/main.py:459` | Memoize `entry.kernel_generator(...)` on its static kwargs | ~100–140 ms per re-lowering |
| C2 | `numba_ffi.py:715`, `numba_cuda_ffi.py:533,1004` | Drop shapes from the FFI signature (they are never read; shape comes from `buf_ptr.dims`) | One permanent FFI target + numba dispatcher per (kernel, shape), never evicted |
| C3 | `kernix_toolchain.py:634,743` | Cache toolchain detection, invalidated by the existing `set_nvcc_discovery`/`set_compute_capabilities` | 8.4 ms/call |
| C4 | `kernix_pipeline.py:92`, `kernix_cache.py:132` | Cache header digests by `(path, mtime_ns, size)` | ~75 ms per lookup, paid on cache *hits* |
| C5 | `kernix_pipeline.py:346,382` | Memoize `.cu` source reads | per lowering |
| C6 | `kernix_pipeline.py:50` | Defer `CompilationCache` `mkdir` out of import | creates a dir on `import brainevent`, breaks read-only filesystems |
| C7 | `__init__.py:18–140` | Lazy submodule imports via the existing PEP 562 `__getattr__` | ~570 ms for a user who only touches `CSR` |

## Tier D — architectural

These are design changes, not cleanups. Each needs its own validation and is
sequenced last; any that prove too risky to land safely will be written up as
follow-up specs rather than forced in.

| # | Finding | Deeper fix |
|---|---------|-----------|
| D1 | `_csr/main.py:977–2809` — `CSR`/`CSC` are a 900-line mirror, 24 methods each | Implement once on `CompressedSparseData` against `_is_row_major` / `_transposed_cls` hooks |
| D2 | `_jit_{normal,scalar,uniform}` are 71–97% textually identical | A `Distribution` descriptor parameterizing one generic family; the CPU seam (`get_numba_light_rng_funcs`) already proves it viable |
| D3 | Unit split/reattach at ~70 entry points (`split_mantissa_unit` ×82, `maybe_decimal` ×67) | A decorator on the existing `@namescope` seam |
| D4 | 191 hand-threaded `*_info=ShapeDtypeStruct(...)` kwargs across 23 files | `_register_fallback_lowering` computes them from the avals it already has |
| D5 | `_fcn/binary.py:896` — caller predicts the lowering-time backend to compute abstract output shape | Backend resolved once at bind time as a static param, or the cuda_raw generator emits its own transpose |
| D6 | `_op/main.py:453` — per-backend preconditions hardcoded in the dispatcher | `KernelEntry.precondition` supplied at registration |
| D7 | `def_kernel('jax_raw', ...)` ×75 | `def_jax_kernel` / `platform='*'` sentinel |
| D8 | `config.py` setters mutate globals read at lowering time, no invalidation | Resolve at bind time as static params, or setters clear caches |
| D9 | `_dense/binary.py:39–52` — CUDA wheel discovery reimplemented outside the toolchain layer | `kernix_toolchain.find_cuda_runtime_library(name)` |
| D10 | `_misc.py:1683` — bare `except Exception` turns a CUDA build failure into a silent CPU fallback | Cheap `gpu_available()` predicate; real toolchain errors propagate |

## Out of scope — probable correctness bugs

Recorded, not fixed here. These need reproducing tests first per `AGENTS.md` rule 4.

- `clen` int32/`ceil` normalization applied in `_jit_uniform/binary.py:741,1367` but
  not in `_jit_normal/binary.py:563,977` or `_jit_scalar/binary.py:704,1267`.
- `_is_static_zero_prob` short-circuit present in `_jit_normal` (3/3 entry points),
  partial in `_jit_uniform` (1/3), absent in `_jit_scalar` (0/3).
- `JITCScalarR * JITCScalarR` works; `JITCNormalR * JITCNormalR` raises.
- Version guards below the supported floor (`jax < 0.4.38`, `< 0.4.34`) are
  unreachable given `jax>=0.8.0`; `check_pallas_jax_version` guards `(0,7,1)` and can
  never raise. One of these is also a live mypy error (`_compatible_import.py:55`).
- `_pallas_random.py` (1,244 lines) is imported only by `__init__.py`; no `_jit_*`
  primitive registers a pallas backend.