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]

❌ 1 Tests Failed:

Tests completed	Failed	Passed	Skipped
1922	1	1921	1190

View the top 1 failed test(s) by shortest run time

tests/test_pca.py::tests.test_pca

Stack Traces | 0s run time

#x1B[1m#x1B[31mtests/test_pca.py#x1B[0m:195: in <module>
    #x1B[0m#x1B[94mfor#x1B[39;49;00m svd_solver, warn_pat_expected, xfail_reason #x1B[95min#x1B[39;49;00m gen_pca_params(#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mtests/test_pca.py#x1B[0m:110: in gen_pca_params
    #x1B[0msvd_solvers, warn_pat_expected = possible_solvers(#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mtests/test_pca.py#x1B[0m:158: in possible_solvers
    #x1B[0msvd_solvers = {#x1B[33m"#x1B[39;49;00m#x1B[33mauto#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m, #x1B[33m"#x1B[39;49;00m#x1B[33mfull#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m, #x1B[33m"#x1B[39;49;00m#x1B[33marpack#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m, #x1B[33m"#x1B[39;49;00m#x1B[33mrandomized#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m} | SKLEARN_ADDITIONAL#x1B[90m#x1B[39;49;00m
                                                             ^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mE   NameError: name 'SKLEARN_ADDITIONAL' is not defined#x1B[0m

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

[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]

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

[amalia-k510]

I think it is just a leftover one that I forgot to change. I ended up switching to using HasArrayNamespace from fast-array-utils for x, since that's what we already dispatch on everywhere. So my question is, would it be better to stick to fast-array-utils or switch to a proper ArrayNamespace instead?