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_xp
function 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 declareAnnData
objects that havearray-api
compatible arrays (in this case,jax
)array-api
compatible objects. In general views are not part of the array-api so our view mechanism will just be "subset the data"as_view
function for thearray-api
-detected objects in b230d11 and added some more tests in 692a270So the todos are likely now (in no particular order):
anndata.concat
will probably need to be updated to perform concatenation/reindexing correctly, so a new method added to theReIndexer
for array-api compatible objects (named_apply_to_array_api
or something as the current ones are named) that aren't already listed (likenumpy.ndarray
)src/anndata/_io/specs/registry.py::Writer.write_elem
will likely need to be updated to use__dlpack__
to convert thearray-api
arrays intonumpy.ndarray
objects 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.ndarray
Migrate creation of
jax.numpy.ndarray
objects (and any other array-api compatible obejcts) into thegen_adata
function 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.,numpy
arrays 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
jax
to our test suite (maybe before doing the above given the current set of tests on this branch as that might be cleaner) - I thinkcubed
is 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_indexer
For each of these, the paradigm should be "create (simple) test case, fix the issue the test raises"