feat: array-api compatibility

[ilan-gold]

@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:

1. 21d5882 adds some tests that should (and do) fail because we have no handling.
2. 7cd6d69 adds the 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 declare AnnData objects that have array-api compatible arrays (in this case, jax)
3. eff9dde builds off the solved problem in 7cd6d69 and now adds a test indexing into such a declared object. This test now fails because we have no way of creating views of array-api compatible objects. In general views are not part of the array-api so our view mechanism will just be "subset the data"
4. I implemented the as_view function for the array-api-detected objects in b230d11 and added some more tests in 692a270

So 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 the ReIndexer for array-api compatible objects (named _apply_to_array_api or something as the current ones are named) that aren't already listed (like numpy.ndarray)

• src/anndata/_io/specs/registry.py::Writer.write_elem will likely need to be updated to use __dlpack__ to convert the array-api arrays into numpy.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 help

• Potentially some tests around views where writing into a view (as of main) should trigger an update of the parent object, but with the array-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 the gen_adata function so that they are created by default on obsm, varm, layers, obsp, and varp. 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 think cubed is probably the best candidate: https://cubed-dev.github.io/cubed/array-api.html

• Support 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"

• Closes #
• Tests added
• Release note added (or unnecessary)

[codecov]

Codecov Report

❌ Patch coverage is 0% with 17 lines in your changes missing coverage. Please review.
✅ Project coverage is 31.71%. Comparing base (7ff7e3b) to head (b52e5b7).
⚠️ Report is 137 commits behind head on main.

Files with missing lines	Patch %	Lines
src/anndata/compat/__init__.py	0.00%	7 Missing ⚠️
src/anndata/_core/index.py	0.00%	5 Missing ⚠️
src/anndata/_core/storage.py	0.00%	3 Missing ⚠️
src/anndata/_core/views.py	0.00%	2 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (7ff7e3b) and HEAD (b52e5b7). Click for more details.

HEAD has 4 uploads less than BASE

Flag	BASE (7ff7e3b)	HEAD (b52e5b7)
	5	1

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #2063       +/-   ##
===========================================
- Coverage   86.81%   31.71%   -55.10%     
===========================================
  Files          46       46               
  Lines        7052     7066       +14     
===========================================
- Hits         6122     2241     -3881     
- Misses        930     4825     +3895     

Files with missing lines	Coverage Δ
src/anndata/_core/views.py	28.87% <0.00%> (-62.27%)	⬇️
src/anndata/_core/storage.py	29.54% <0.00%> (-68.08%)	⬇️
src/anndata/_core/index.py	44.44% <0.00%> (-48.22%)	⬇️
src/anndata/compat/__init__.py	28.44% <0.00%> (-52.29%)	⬇️

... and 40 files with indirect coverage changes