Repository navigation
refactor: ♻️ adopt the array API standard for backend dispatch and array typing - #69
Merged
Merged
Conversation
…ray typing The algorithms take the array namespace from their arrays through array-api-compat instead of receiving an `xp` module. The namespace is picked once from `device=` by `DeviceGroup`, which now also resolves the device ids; `EnvironmentStore` follows the arrays it is given, and `truncated_qr` and the splitting kernels follow their inputs. Arrays of either backend are annotated with one alias, `src_method.utils.Array`; host-only arrays (the results, the exact path, the sketches) are `np.ndarray`. Array creation, `linalg.qr`, `conj`, `reshape`, `moveaxis`, `concat` and `asarray` use the standard names. Streams, pinned memory and CuPy's pool stay behind `utils._backend`, and the Gaussian sketches are always drawn on the host. Closes #45 Assisted-by: GitHub Copilot:Claude Opus 5.5
Contributor
Contributor
There was a problem hiding this comment.
🔵 Needs a closer look
Cross-cutting GPU dispatch and lazy staging changes warrant final human review despite the reported passing tests and benchmarks.
0 open findings
What changed in this PR
Adopts array API compatibility for SRC’s CPU/GPU dispatch and array typing, addressing issue #45.
Changes:
- Infers namespaces from arrays and centralizes device selection in
DeviceGroup. - Introduces backend-generic
Arrayannotations and lazy staging-buffer allocation. - Updates dependencies, existing tests, and contributor documentation.
| File | Description |
|---|---|
| uv.lock | Locks the compatibility dependency. |
| tests/test_store.py | Updates store construction. |
| tests/test_stack.py | Updates sweep calls. |
| tests/test_sites.py | Updates device-group setup. |
| tests/test_plan.py | Updates budget-resolution tests. |
| tests/test_package.py | Updates sketch calls. |
| tests/test_group.py | Updates group construction. |
| tests/test_backend.py | Checks namespace detection and updated copying. |
| src/src_method/utils/linalg.py | Infers QR backend; keeps SVD host-side. |
| src/src_method/utils/_backend.py | Provides compatibility namespaces and array alias. |
| src/src_method/utils/__init__.py | Exports Array. |
| src/src_method/stack.py | Passes device selection to the sweep. |
| src/src_method/compress.py | Clarifies host-array return typing. |
| src/src_method/apply.py | Clarifies host-array return typing. |
| src/src_method/_tensor_train.py | Distinguishes generic and host arrays. |
| src/src_method/_sweep.py | Removes namespace parameters; uses standard operations. |
| src/src_method/_store.py | Infers backend and allocates staging rings lazily. |
| src/src_method/_sites.py | Gets its namespace from DeviceGroup. |
| src/src_method/_plan.py | Resolves budgets through device groups. |
| src/src_method/_kernels.py | Uses backend-generic array annotations. |
| src/src_method/_group.py | Centralizes namespace and device selection. |
| pyproject.toml | Adds array-api-compat at runtime. |
| docs/content/docs/getting-started.mdx | Updates runtime dependency guidance. |
| docs/content/docs/contributing/index.mdx | Documents namespace and typing conventions. |
| docs/content/docs/contributing/dependencies.mdx | Lists the new dependency. |
| AGENTS.md | Updates backend-development guidance. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
…rc-method into refactor/array-api
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
🤖 AI text below 🤖
Closes #45.
What changes
array_api_compat.array_namespace(truncated_qr, the out-of-memory splitting in_sweep,EnvironmentStore). Nothing outsideutils/_backend.pytakes anxpparameter any more.DeviceGroup(device, devices)resolves the namespace from the publicdevice=string and the device ids fromResources.devices, soresolve_devicesfolds into it.resolve_budgets(resources, group),SiteSource(..., group=group)andsweep(..., device=...)follow.EnvironmentStoreno longer needs a backend: it follows the arrays it is given, and allocates its page-locked rings on the first stagedput.src_method.utils.Arrayis the documented alias for an array of the active backend (numpy.ndarrayorcupy.ndarray). CuPy ships no stubs, so it isAnyto a type checker. Host-only arrays (the results ofsrc/apply/compress, the exact two-site path, the sketches) are now annotatednp.ndarray, and noNDArrayremains on backend-generic arrays.linalg.qr,conj,reshape,moveaxis,concatandasarrayuse the standard API.truncated_qris shorter and keeps the SVD ofRexplicitly on the host.gaussian_sketch, which now always returns a host array (it was only ever called with NumPy).array-api-compat>=1.12becomes a runtime dependency. The docs (getting started, contributing, dependencies) andAGENTS.mddescribe the new convention.Not done, on purpose:
from_dlpack: there is no zero-copy exchange to make. Host→device and device→host are real copies, andto_numpystays the single device→host transfer.ascontiguousarray(the scatter needs C-contiguous blocks) and the byte-level.view(dtype)of the staging buffers have no standard equivalent.xp.cuda, whicharray_api_compat.cupyre-exports.Tests
On Leonardo Booster (4× A100-64GB, CUDA 13 via forward compat):
pytest(the whole suite, slow and GPU tests included): 265 passed, 0 skipped, including the two-GPU tests.pytest -m "not slow"at the lowest direct versions (Python 3.11, array-api-compat 1.12.0, NumPy 2.0.0, opt_einsum 3.4.0): 245 passed, 1 skipped (the CuPy module).prek run --all-files: clean.Performance
origin/mainvs. this branch, in the same job, alternating processes. Each case was run 5–7 times per process, over 2–3 rounds, and the table gives the best per implementation. Outputs agree bit for bit.gpu_memory="1GB"(environments on the host tier)Single rounds of the CPU depth-3 stack differ by up to 3% between processes, whichever code runs. A rerun with the branch first gave 14.770 s (branch) against 14.802 s (main), and the per-pass sweep timings match.
These cases are dominated by per-call overhead, so I also timed them at steady state (30 back-to-back calls, two rounds):
The compat wrappers cost a few µs per call on CuPy once warm. A first version paid about 1.5 ms per small sweep for
array_namespaceandis_numpy_namespacecalls in the per-batch path. Those now run only where they are needed, andis_hostis an identity check again.pytest-benchmark (min, ms):
test_benchmark_src_stack_depth3test_benchmark_src_mpo_mpoSize
Python lines: 8344 → 8334 (src: 3555 → 3550). Whole diff, docs and lockfile included: +289 / −281.