feat: add Array API support via array-api-compat

[amalia-k510]

This adds Array API support to scanpy's preprocessing pipeline, where I am planning to use JAX as the test case. The goal is for someone to put a JAX array into adata.X and run the usual pipeline without the array getting silently pulled to CPU partway through. It builds on the recent fast-array-utils work, which means a lot of the heavy lifting (sums, means, variances) already works across backends.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 0.00%. Comparing base (27bba90) to head (e3cdf5e).
✅ All tests successful. No failed tests found.

❌ Your project check has failed because the head coverage (0.00%) is below the target coverage (75.00%). You can increase the head coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #4179       +/-   ##
==========================================
- Coverage   78.98%       0   -78.99%     
==========================================
  Files         128       0      -128     
  Lines       13402       0    -13402     
==========================================
- Hits        10586       0    -10586     
+ Misses       2816       0     -2816     

Flag	Coverage Δ
hatch-test.low-vers	?
hatch-test.pre	?

Flags with carried forward coverage won't be shown. Click here to find out more.
see 128 files with indirect coverage changes

[flying-sheep]

OK, please remember the comments from the fast-array-utils PR, e.g. we agreed to handle singledispatch functions like this: https://github.com/scverse/fast-array-utils/blob/febaf245ecd6da849c2331638219f649935d5a21/src/fast_array_utils/stats/_power.py#L38-L39

We also need tests. There’s MAP_ARRAY_TYPES which should gain an entry, which will make many many tests run with the type you add. We can discuss how to best exclude certain tests, for starters you can just add it and make some of the tests fail – that’ll show you where things aren’t working yet (and might or might not be fixed in this PR)

Please also add a release note (hatch run towncrier:create 4179.feat.md)

[flying-sheep]

this handles numpy with equal performace too no?

[amalia-k510]

xp.count_nonzero resolves to a wrapper over np.count_nonzero. The difference between the two in negligible. I decided to keep it explicit due to the code consistency as that is what I have been doing for other fuctions. I can drop it, if you'd rather have less code.

[flying-sheep]

no median for array api?

[amalia-k510]

median isn't in the array API spec (it's not in the statistical functions list), but array-api-compat namespace expose it, so xp.median works for the backends we actually support.

[flying-sheep]

Looking good! Just nitpicks left in the code that’s there.

One big change we should do is updating the array type metadata so it lists Array API support.

The docs currently have one central table and each function has its own little info box.

Since you know what’s going to be supported, you should update the metadata by adding an aa tag and filling out the support here (so far, np means numpy, sp means scipy sparse, da means dask, da[...] means “... in dask”, and sp[csc] means “scipy sparse with csc format”):

scanpy/docs/conf.py

Lines 178 to 210 in fc0b419

	array_support: dict[str, tuple[list[str], list[str]]] = {
	"experimental.pp.highly_variable_genes": (["np", "sp"], []),
	"get.aggregate": (["np", "sp", "da"], []),
	"pp.calculate_qc_metrics": (["np", "sp", "da"], []),
	"pp.combat": (["np"], []),
	"pp.downsample_counts": (["np", "sp[csr]"], []),
	"pp.filter_cells": (["np", "sp", "da"], []),
	"pp.filter_genes": (["np", "sp", "da"], []),
	"pp.harmony_integrate": (["np"], []),
	"pp.highly_variable_genes": (["np", "sp", "da"], ["da[sp[csc]]"]),
	"pp.log1p": (["np", "sp", "da"], []),
	"pp.neighbors": (["np", "sp"], []),
	"pp.normalize_total": (["np", "sp[csr]", "da"], []),
	"pp.pca": (["np", "sp", "da"], ["da[sp[csc]]"]),
	"pp.regress_out": (["np"], []),
	"pp.sample": (["np", "sp", "da"], []),
	"pp.scale": (["np", "sp", "da"], []),
	"pp.scrublet": (["np", "sp"], []),
	"pp.scrublet_simulate_doublets": (["np", "sp"], []),
	"tl.dendrogram": (["np", "sp"], []),
	"tl.diffmap": (["np", "sp"], []),
	"tl.dpt": (["np", "sp"], []),
	"tl.draw_graph": (["np", "sp"], []), # only uses graph in obsp
	"tl.embedding_density": (["np"], []),
	"tl.ingest": (["np", "sp"], []),
	"tl.leiden": (["np", "sp"], []), # only uses graph in obsp
	"tl.louvain": (["np", "sp"], []), # only uses graph in obsp
	"tl.paga": (["np", "sp"], []),
	"tl.rank_genes_groups": (["np", "sp"], []),
	"tl.score_genes": (["np", "sp"], []),
	"tl.tsne": (["np", "sp"], []),
	"tl.umap": (["np", "sp"], []),
	}

If you want, you can also update the Sphinx code, but you don’t need to learn how to do that if you don’t want to. I’ll happily do that myself.

[flying-sheep]

Did I forget something or why aren’t we using types from here? https://array-api.readthedocs.io/en/latest/index.html

[flying-sheep]

why?