Skip to content

Multi-backend implicit solve: Array-API dense fallback + SciPy/CuPy sparse registry #69

Description

@jc-macdonald

Summary

Give CoreSolver's implicit/IMEX linear-solve step a two-tier, backend-general dispatch instead of being SciPy-sparse-only, so "any array library meeting the Array API standard" is a real, enforced guarantee — not just true for the elementwise/dense paths #66 already covers.

Design (agreed 2026-08-18, supersedes the "two honestly-separate engines" framing originally in #26/#67)

Tier 1 — universal dense fallback (correctness guarantee, zero per-backend code). Represent the implicit-step operator as a dense array and solve via the Array API standard's linalg extension (xp.linalg.solve, available uniformly across NumPy, CuPy, JAX, and PyTorch via array-api-compat). Any backend implementing __array_namespace__ + the linalg.solve extension gets correct implicit/IMEX behavior automatically — matching op_system's "no compile-time backend selection" promise for the solve step itself, not just the elementwise ops #66 covers.

Tier 2 — sparse acceleration registry (opt-in performance, not required for correctness). A small, explicit registry — {ecosystem_id: SparseAdapter(issparse, factorize, solve, cache_key)} — data, not a class-hierarchy Protocol. Ship with two adapters day one:

  • scipy: scipy.sparse/scipy.sparse.linalg — today's existing _as_scipy_operator/implicit_solve code becomes this adapter, functionally unchanged.
  • cupy: cupyx.scipy.sparse/cupyx.scipy.sparse.linalg — near-mirror of the SciPy adapter; CuPy ships matching issparse/spsolve/splu/cg by design, since CuPy is built as a drop-in GPU mirror of NumPy/SciPy.

Detection tries each registered ecosystem's own issparse()scipy.sparse.issparse() alone does not recognize CuPy or JAX sparse types (confirmed), so a single hardcoded SciPy check would silently misroute other backends. Falls through to Tier 1 dense when nothing matches. Future ecosystems get sparse acceleration by adding a registry entry, not by subclassing an interface.

Why this isn't the premature abstraction the earlier LinearSolverBackend Protocol was (#19/#21/#22/#23/#25, closed): those were closed for having exactly one implementer and no second one in sight. This registry ships with two concrete, near-identical, real implementations from day one (SciPy CPU, CuPy GPU), and is a dict of functions rather than a class hierarchy requiring subclassing — cheap to extend, not ceremony-heavy.

JAX is intentionally not in this registry. Dense JAX already gets full Tier-1 support (GPU/TPU, jit/vmap/grad-compatible). JAX's sparse-at-scale story is structurally different (no drop-in splu-shaped call convention) and is lineax/diffrax, tracked separately as new CoreSolver methods in #67 — not a Tier-2 registry entry.

Scope

  • Refactor matrix_ops.py's _as_scipy_operator/implicit_solve boundary into the scipy registry adapter (behavior-preserving).
  • Add the cupy registry adapter (cupyx.scipy.sparse/.linalg), gated so it's only imported/used when CuPy is installed (optional extra, mirroring how jax is optional elsewhere in this ecosystem).
  • Add the Tier-1 dense fallback path using the Array API linalg.solve extension (via array-api-compat or the namespace's own .linalg.solve), used whenever no registry adapter recognizes the operator.
  • Registry lookup keyed on structural detection (each adapter's own issparse), not backend-name string matching.
  • LU/factorization caching (existing performance feature) lives inside each sparse adapter, same cache-key semantics as today's SciPy path.

Acceptance criteria

  • A cross-backend contract test: solve one small implicit system on every backend available in CI (NumPy always; CuPy/JAX if installed) and assert matching results within tolerance — a runnable guarantee of the "meets the Array API standard" promise, not just an aspiration.
  • Existing SciPy-sparse-backed behavior (accuracy, caching) is unchanged.
  • A dense-only backend with no registered sparse adapter still produces correct (if unaccelerated) implicit-method results.
  • just quality and just test pass.

Relationship to other issues

Sibling to #66 (dense/elementwise namespace inference) — together they make the entire non-diffrax CoreSolver surface backend-general by construction, not by convention. Blocks the "dense IMEX already works on JAX via Tier 1" claim in the revised #67. Part of the overall plan tracked in #26.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    architectureStructural or design-level changesenhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions