feat: array-api compatibility #2063
Open
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.
@amalia-k510 I will be out for 2.5 weeks starting tomorrow. Please have a look at this branch and work off of it.
There are three commits:
has_xpfunction and uses it to both raise a more informative error for indexing with an array-api compatible object and resolve two of the test fail cases from 21d5882 so that now we can declareAnnDataobjects that havearray-apicompatible arrays (in this case,jax)array-apicompatible objects. In general views are not part of the array-api so our view mechanism will just be "subset the data"as_viewfunction for thearray-api-detected objects in b230d11 and added some more tests in 692a270So the todos are likely now (in no particular order):
anndata.concatwill probably need to be updated to perform concatenation/reindexing correctly, so a new method added to theReIndexerfor array-api compatible objects (named_apply_to_array_apior something as the current ones are named) that aren't already listed (likenumpy.ndarray)src/anndata/_io/specs/registry.py::Writer.write_elemwill likely need to be updated to use__dlpack__to convert thearray-apiarrays intonumpy.ndarrayobjects so that they can be safely written using the existing code. https://numpy.org/doc/stable/reference/generated/numpy.from_dlpack.html will likely be of helpPotentially some tests around views where writing into a view (as of
main) should trigger an update of the parent object, but with thearray-api, we can probably no longer guarantee that for e.g.,jax.numpy.ndarrayMigrate creation of
jax.numpy.ndarrayobjects (and any other array-api compatible obejcts) into thegen_adatafunction so that they are created by default onobsm,varm,layers,obsp, andvarp. At this point a lot of tests will start failing but fear not! This is exactly what you want as it will show you what you need to implement - go one by one and implement the new features (keeping the old behavior as-is i.e.,numpyarrays should continue down their previous codepaths). The above three todo's likely will appear if you do this first!Add another array-api compatible library beyond
jaxto our test suite (maybe before doing the above given the current set of tests on this branch as that might be cleaner) - I thinkcubedis probably the best candidate: https://cubed-dev.github.io/cubed/array-api.htmlSupport indexing by an array-api compatible object as you started in Backend-Agnostic Refactor Using Array API Compatibility #2019 (there is a test I have in this branch for that via
test_jax_indexerFor each of these, the paradigm should be "create (simple) test case, fix the issue the test raises"