Adopt the Python array API standard for backend dispatch and array typing #45

🤖 AI text below 🤖

Summary

Adopt the Python array API standard through array-api-compat. The array backend would then be inferred from the arrays themselves (array_namespace(x)) instead of being passed around as xp, and "an array of the current backend" would get an honest type, replacing today's loose use of numpy.typing.NDArray.

Current state

• The algorithms receive the array module as an xp argument and call xp.asarray, xp.linalg.qr, xp.empty and so on directly. This keeps the CPU (NumPy) and GPU (CuPy) paths in sync, but xp has to be threaded through every function.
• Arrays of either backend are annotated as numpy.typing.NDArray. For CuPy arrays that's incorrect. It goes unnoticed only because CuPy ships no type stubs, so the type checker sees Any.
• Backend-specific features (CUDA streams and events, page-locked memory, CuPy's memory pool, the NCCL communicator planned for the multi-GPU work in feat(sweep): ✨ split the SRC sweep among the GPUs of a node #44) already live behind src_method.utils._backend.

Proposal

1. Namespace from the arrays. Use array_api_compat.array_namespace(*arrays) in library code instead of an xp parameter, wherever the arrays are at hand. The public device= argument stays; it only chooses where inputs are placed.
2. One array type. Introduce a single alias for arrays of the current backend, based on an array API typing protocol or a documented Array alias, and use it everywhere NDArray is used today.
3. Standard names where they exist. Move array creation, linalg.qr/linalg.svd, conj, reshaping, transposition and asarray to the standard API.
4. Keep what the standard doesn't cover behind _backend.py: streams and events, pinned memory, the memory-pool cap and queries, and the host-side random draws, which stay on NumPy's default_rng so a seed gives the same draws on every device.
5. Host transfers. Use from_dlpack for zero-copy exchange where it applies. to_numpy stays as the single device-to-host transfer.

Benefits

• Fewer parameters, and backend dispatch that follows the data.
• Types that are correct for CuPy, and useful once a stricter type checker is in place (see Adopt ty as type checker #28).
• A route to further backends (PyTorch, JAX, Dask) at little extra cost, as far as opt_einsum and the linear algebra allow.

Open questions and risks

• opt_einsum already dispatches on the operand types. Check that it accepts the arrays array-api-compat wraps or returns, without an extra copy.
• truncated_qr does part of its work on the host (the SVD of R). Keep that explicit.
• Performance must not regress: the per-operation overhead of the compatibility layer should be negligible next to the contractions. Measure with tests/test_stack.py::test_benchmark_src_stack_depth3 and tests/test_package.py::test_benchmark_src_mpo_mpo.
• Scheduling: land after the out-of-core (feat(sweep): ✨ out-of-core SRC sweep on one GPU #42) and multi-GPU (feat(sweep): ✨ split the SRC sweep among the GPUs of a node #44) work, so the refactor doesn't collide with them.

Acceptance

• No xp parameters left in the algorithm modules, apart from the _backend.py seam.
• One array alias used consistently; no numpy.typing.NDArray on backend-generic arrays.
• uv run prek run --all-files and uv run pytest, including the GPU tests on a GPU machine, pass; the benchmarks above within noise.