Array API pilot: backend dispatch for clean numerical kernels (rescoped from #19)

[Hugo-W]

Array API pilot: backend dispatch for clean numerical kernels (rescoped from #19)

Issue #19 ("Implement Python Array API Standard 2024.12 for array backend abstraction") was closed as not realistically achievable as written — full JAX/PyTorch support for all functionality collides with structural blockers (stride-based lag construction, SciPy/sklearn/compiled/RNG dependencies, ecosystem immaturity). See the comment on #19 for the full blocker analysis.

This issue tracks a bounded pilot to gain the real, achievable benefits of Array API adoption without overpromising whole-library portability.

Goal

Introduce array-api-compat namespace dispatch (xp = array_namespace(*arrays)) into a small set of explicitly identified pure numerical kernels, enabling NumPy/JAX/PyTorch arrays to flow through those kernels without implicit host conversion. Everything else stays NumPy-only with a documented portability matrix.

Scope — pilot kernels

1. pyeeg/models/var.py (cleanest candidate, ~80–90% mechanical)

   • Replace np with dispatched xp for array creation, atleast_2d/atleast_3d, and linalg.lstsq.
   • No SciPy/sklearn/stride dependencies.

2. Materialized lag-matrix prototype (alongside the zero-copy fast path)

   • Current as_strided + .strides implementation in pyeeg/utils.py is out of scope of the Array API and must remain the NumPy fast path.
   • Add a portable materialized (copy-based) lag implementation using standard indexing/stacking.
   • Gate on backend: NumPy uses as_strided; other backends use the materialized path.
   • Validate shape, lag ordering, edge padding, dtype/device preservation, and memory cost.

3. Dense TRF solver paths (pyeeg/solvers.py)

   • Port only the closed-form dense linear algebra (svd, eigh, solve, lstsq, inv, pinv — all covered by Array API linalg).
   • Rewrite the 6 np.einsum calls to matmul/tensordot/vecdot (einsum is not in the 2024.12 standard).
   • Keep NumPy-only: scipy.sparse (spilu, csc_matrix), scipy.optimize.least_squares (robust path). These become explicit boundary functions that reject or convert non-NumPy input.

Out of scope (NumPy-only boundary layers)

• scipy.* paths: sparse, optimize, signal, stats CDFs, generalized eigh
• scikit-learn estimators: CCA, PCA, KFold, covariance
• np.random.default_rng (stats.py, simulate.py) — RNG not standardized by Array API
• I/O / viz / dashboard / compiled gammatone & ratemap C-extensions
• llm_features.py (already PyTorch-native)

Approach

• Use array-api-compat (array_namespace()) for dispatch — no ad-hoc xp parameters threading through the public API.
• Use array-api-strict as a test gate to expose accidental NumPy dependencies.
• Test pilot kernels against NumPy, JAX, and PyTorch (optional extras).
• Cover dtype promotion, devices, immutable arrays, empty/degenerate inputs, output namespace, and absence of implicit host conversion.

Acceptance criteria

• var.py runs backend-native on NumPy, JAX, and PyTorch without implicit host conversion.
• Materialized lag-matrix produces identical output to the as_strided fast path for NumPy; works on JAX/PyTorch.
• Dense TRF solver paths run backend-native; robust/sparse paths explicitly reject or convert non-NumPy input with a clear error.
• A portability matrix is documented (backend-native / NumPy-fallback / NumPy-only / unsupported) and becomes the real acceptance contract.
• Existing NumPy-only test suite remains green.

Non-goals

• Whole-library JAX/PyTorch support.
• Autodiff / JIT / GPU acceleration claims (validate separately if pursued).
• Abstracting I/O, visualization, dashboard, simulation RNG, or compiled extensions.

Priority: low (carried over from #19).