Skip to content

Add CuPy contract test to prove the 'any Array-API backend' claim beyond NumPy/JAX #182

Description

@jc-macdonald

Motivation

op_system's README claims the compiled eval_fn/pytree_eval_fn "runs identically on NumPy, JAX (concrete and traced), or any other Array-API backend — without recompiling," and architecturally this holds: compile.py's generated eval closures bind np inside the expression-evaluation environment to the inferred xp namespace at call time (env: dict[str, object] = {"np": xp, "t": t_val}, compile.py:796), not to literal numpy — there is no hardcoded backend anywhere in src/op_system/.

But this claim is currently untested beyond NumPy and JAX: only 2 test files touch JAX, zero touch CuPy or PyTorch, and there's no cupy optional-dependency group in pyproject.toml (only jax/jax-inference/data exist).

This matters concretely now: ACCIDDA/op_engine's revised multi-backend plan (op_engine#26, op_engine#69) makes CuPy a first-class Tier-2 backend for CoreSolver's sparse-accelerated implicit solve, justified specifically by CuPy's near-identical API to scipy.sparse. That guarantee is only as trustworthy as op_system's own compiled RHS actually behaving correctly under CuPy arrays end-to-end — which has never been verified.

Scope

  • Add a CuPy-gated contract test (parametrized alongside the existing NumPy/JAX parity tests, or a new dedicated module) that:
    • Compiles a small representative spec (e.g. the SIR example already used elsewhere in the test suite).
    • Calls eval_fn/pytree_eval_fn with CuPy arrays for y/state and asserts numerical parity with the NumPy result within tolerance.
    • Skips gracefully (pytest.importorskip("cupy") and/or a GPU-availability check) when CuPy or a CUDA-capable GPU isn't present — standard ubuntu-latest GitHub Actions runners have no GPU, so this test will not run in default CI without a GPU-capable runner. Document this explicitly (in the test docstring and in docs/development/) so it isn't mistaken for silently-passing coverage.
  • Add an optional cupy dependency group to pyproject.toml for local/GPU-runner verification, mirroring the jax extra.
  • Consider structuring the test harness generically (e.g. a small parity-check helper parametrized by array module) so PyTorch (already claimed as a qualifying backend in compile.py's _namespace_of docstring, also currently untested) can reuse it later without new scope creep here.

Acceptance criteria

  • CuPy contract test exists, passes on a CUDA-capable environment, and skips cleanly (not silently) when unavailable.
  • pyproject.toml documents a cupy extra for running it.
  • README/docs note that the "any Array-API backend" claim's verified surface now explicitly includes CuPy (with the GPU-runner caveat), not just NumPy/JAX.
  • just quality and just test continue to pass in standard CI (no GPU dependency introduced into the default test run).

Relationship to other issues / repos

Directly motivated by ACCIDDA/op_engine#26 (tracking) and ACCIDDA/op_engine#69 (Array-API dense fallback + SciPy/CuPy sparse registry) — that plan's cross-backend contract-test acceptance criterion on the op_engine side is only meaningful if op_system's own compiled RHS is proven correct under CuPy first.

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

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions