Epic: make MLX a first-class, end-to-end backend #324

Goal

Make the MLX backend first-class and trustworthy end to end, not just feature-complete in isolation. Epic #278 brought AMICAMLXNG to fitting and non-fitting parity with AMICATorchNG, but every MLX user is still a raw-backend user: the AMICA and AMICAICA wrappers build torch unconditionally (#313), MLX has no explicit pcakeep/pcadb (#323) and no params-file path (#304). We are the first users (per-session AMICA on 64-channel EEG on Apple Silicon, pcakeep = nbchan - 1 after average reference), so the bar is a workflow we trust.

Phases

Phase 1: shared pcakeep/pcadb validation and reduction-request predicate in pamica/rank.py, used by torch, NumPy and MLX; MLX constructor parameters, upfront mir_step gate and persistence; cross-backend test on real EEG.

Phase 2: AMICAICA exports the full PCA basis so MNE's own apply restores the PCA residual by default; ADR, differences-page section and docstrings stating plainly that this goes beyond the Fortran reference (whose output carries no representation of the discarded subspace) and how to get the reference behavior (n_pca_components=ica.n_components_).

Phase 3: one shared params-file reader returning canonical keys for JSON and Fortran text; AMICA_NumPy accepts input.param; NumPy raises NotImplementedError for pdftype != 0.

Phase 4: AMICA(backend="mlx") and AMICAICA(backend="mlx"), including save/load, from_params_file and pcakeep through the wrapper, with end-to-end tests on both backends.

Phase 5: degenerate-fit and input guards on the raw torch and MLX accessors (#306).

Phase 6: validate_implementations.py --backend {torch,numpy,mlx}, an end-to-end test mirroring our own pipeline, and an MLX getting-started path (#315).

Phases 7-9 (added 2026-09-22 after a read-only investigation found a component-orientation defect; user decision to fold them into this epic): pamica stores each model's mixing block transposed relative to the reference, so components are rows of the stored block, but doscaling (default on), the share metric and the share merge act on stored columns. Phase 7 fixes doscaling (default trajectories change toward the reference; native-oracle test; ADR amending ADR 0001), Phase 8 moves to a component-row layout that makes sharing correct (byte-identical without sharing; save-format conversion), Phase 9 confirms and fixes a suspected one-iteration Newton start offset. After Phase 9 the per-backend parity rows are re-measured.

Phases 10-16 (added 2026-09-22/23 as reviews and native-oracle tests surfaced further reference mismatches): Phase 10 writes the sphere column-major as the reference reads it; Phase 11 follows the reference's iteration order (E-step, checks and the LL-decrease response, exit, then update) and its unconditional A-freeze windows; Phase 12 normalizes the drawn initial mixing matrix as the reference does; Phase 13 uses the reference's single-precision density constants; Phase 14 makes AMICA_NumPy reject unknown options. Phase 15 re-measures every parity figure (validation rows, the paper's Table 1, the multi-model ensemble, ADR 0003) under the finished code; Phase 16 audits the documentation against it; Phase 17 unifies the lrate default across entry points on the compiled binary's 0.1 (user decision), makes the raw PyTorch backend fall back to CPU for float64 on Apple Silicon, and documents pamica's defaults against the compiled binary and EEGLAB's runamica15.m.

Evidence that motivated the scope (real sample EEG, 32 channels)

• AMICAICA.apply after pcakeep=20 with nothing excluded loses 8.0% of the signal (relative error 0.080); exporting the full PCA basis gives 1.2e-15. Fixed in Restore the PCA residual in AMICAICA.apply #325.
• AMICATorchNG does not validate pcakeep/pcadb: pcakeep=-3 silently fits 29 sources, pcakeep=2.7 truncates to 2, pcakeep=0 and pcadb<=0 end in nan_ll. Fixed in Phase 1: shared pcakeep/pcadb policy and MLX support #326 (shared validator on every backend).
• The torch mir_step gate rejects pcakeep=32 on 32-channel data, the value both bundled param files carry. Fixed in Phase 1: shared pcakeep/pcadb policy and MLX support #326.
• sample_params.json and input.param agree on all 47 shared keys (except the data path) once JSON-schema spellings are canonicalized, but the wrapper does not apply max_decs/min_grad_norm/share_int from the JSON today. Fixed in Phase 3: shared params-file reader for every backend #327 (shared reader with canonical keys).

Each box is checked when the phase that fixes it merges.

Definition of Done

• Every phase PR merged with its pre-registered gate tests green, reviewed with the pr-review-toolkit, all findings addressed.
• AMICA(backend="mlx") and AMICAICA(backend="mlx") run fit, transform, save/load, params-file, pcakeep, EEGLAB export and MNE apply end to end, covered by tests.
• validate_implementations.py reports torch, NumPy and MLX against the Fortran binary.
• Every deliberate divergence from the reference introduced here is recorded in docs/guides/amica-differences.md (with an ADR where a default changes).
• doscaling and share_comps act on components exactly as the reference does, pinned by native-oracle tests (doscaling normalizes stored columns, not components (all backends) #333, share_comps compares and merges the wrong vectors (all backends) #334); the Newton start matches the reference (Newton appears to start one iteration late (0- vs 1-based newt_start) #335); per-backend parity rows re-measured afterwards.
• Every parity figure in the docs and paper.md re-measured under the finished code (Re-measure parity and sweep docs for epic #324 #351), and the documentation audited against it (Audit the documentation against the finished epic #324 #352).
• Epic branch merged into dev with a regular merge commit after user approval.

The plan of record, with decided policies, deliverables and pre-registered test gates per phase, is .context/issue-324/plan.md on the epic branch.