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

docs/conf.py

@@ -188,14 +188,15 @@
     "pp.filter_genes": (["np", "sp", "da"], []),
     "pp.harmony_integrate": (["np"], []),
     "pp.hashsolo": (["np", "sp"], []),
-    "pp.highly_variable_genes": (["np", "sp", "da"], ["da[sp[csc]]"]),
-    "pp.log1p": (["np", "sp", "da"], []),
+    # HVG flavor seurat_v3 has no array api support yet if `batch_key` is not None
+    "pp.highly_variable_genes": (["np", "sp", "da", "xp"], ["da[sp[csc]]"]),
+    "pp.log1p": (["np", "sp", "da", "xp"], []),
     "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.scale": (["np", "sp", "da", "xp"], []),
     "pp.scrublet": (["np", "sp"], []),
     "pp.scrublet_simulate_doublets": (["np", "sp"], []),
     "tl.dendrogram": (["np", "sp"], []),
@@ -207,7 +208,7 @@
     "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.rank_genes_groups": (["np", "sp", "xp"], []),
     "tl.score_genes": (["np", "sp"], []),
     "tl.tsne": (["np", "sp"], []),
     "tl.umap": (["np", "sp"], []),

docs/extensions/array_support.py

@@ -18,7 +18,7 @@
     from sphinx.application import Sphinx
 
 
-ALL_INNER = list(_docs.parse(["np", "sp"], inner=True))
+ALL_INNER = list(_docs.parse(["np", "sp", "xp"], inner=True))
 
 
 class ArraySupport(SphinxDirective):

docs/release-notes/4179.feat.md

@@ -0,0 +1 @@
+Add Array-API support, enabling JAX and other array-api backends in `adata.X` {smaller}`A. Karesh`

pyproject.toml

@@ -53,7 +53,8 @@ classifiers = [
 dynamic = [ "version" ]
 dependencies = [
     "anndata>=0.12.14",
-    "fast-array-utils[accel,sparse]>=1.4",
+    "array-api-compat",
+    "fast-array-utils[accel,sparse]>=1.5",
     "h5py>=3.11",
     "joblib",
     "matplotlib>=3.10",
@@ -103,10 +104,12 @@ scanpy2 = [ "anndata>=0.13.3", "hv-anndata>=0.0.3a5", "igraph>=0.10.8", "scanpy[
 
 [dependency-groups]
 dev = [
-    "scipy-stubs", # static typing and IDE support
-    "towncrier",   # release note management
+    "scipy-stubs",     # static typing and IDE support
+    "towncrier",       # release note management
+    "types-array-api",
 ]
 test = [
+    "jax",                          # Array API tests
     "scanpy[dask-ml]",
     "scanpy[dask]",
     "scanpy[illico]",

src/scanpy/_utils/__init__.py

@@ -30,6 +30,8 @@
 import numpy as np
 import pandas as pd
 from anndata._core.sparse_dataset import BaseCompressedSparseDataset
+from array_api_compat import array_namespace
+from fast_array_utils.types import HasArrayNamespace
 
 from .. import logging as logg
 from .._compat import CSBase, DaskArray, SpBase, warn
@@ -609,6 +611,20 @@ def axis_mul_or_truediv(
     *,
     allow_divide_by_zero: bool = True,
     out: ArrayLike | None = None,
+) -> np.ndarray:
+    raise NotImplementedError
+
+
+@axis_mul_or_truediv.register(np.ndarray)
+def _(
+    x: np.ndarray,
+    /,
+    scaling_array: np.ndarray,
+    axis: Literal[0, 1],
+    op: Callable[[Any, Any], Any],
+    *,
+    allow_divide_by_zero: bool = True,
+    out: ArrayLike | None = None,
 ) -> np.ndarray:
     _check_op(op)
     scaling_array = _broadcast_axis(scaling_array, axis)
@@ -731,8 +747,38 @@ def _[T: (DaskArray, np.ndarray)](
     )
 
 
+@axis_mul_or_truediv.register(HasArrayNamespace)
+def _(
+    x: HasArrayNamespace,
+    /,
+    scaling_array: np.ndarray,
+    axis: Literal[0, 1],
+    op: Callable[[Any, Any], Any],
+    *,
+    allow_divide_by_zero: bool = True,
+    out: ArrayLike | None = None,
+) -> Any:
+
+    _check_op(op)
+    scaling_array = _broadcast_axis(scaling_array, axis)
+    xp = array_namespace(x)
+    scaling_array = xp.asarray(scaling_array)
+    if op is mul:
+        return x * scaling_array
+    if not allow_divide_by_zero:
+        scaling_array = xp.where(
+            scaling_array == 0, xp.ones_like(scaling_array), scaling_array
+        )
+    return x / scaling_array
+
+
 @singledispatch
 def axis_nnz(x: ArrayLike, /, axis: Literal[0, 1]) -> np.ndarray:
+    raise NotImplementedError
+
+
+@axis_nnz.register(np.ndarray)
+def _(x: np.ndarray, /, axis: Literal[0, 1]) -> np.ndarray:
     return np.count_nonzero(x, axis=axis)
 
 
@@ -751,6 +797,12 @@ def _(x: DaskArray, /, axis: Literal[0, 1]) -> DaskArray:
     )
 
 
+@axis_nnz.register(HasArrayNamespace)
+def _(x: HasArrayNamespace, /, axis: Literal[0, 1]) -> Any:
+    xp = array_namespace(x)
+    return xp.count_nonzero(x, axis=axis)
+
+
 @singledispatch
 def check_nonnegative_integers(x: _SupportedArray, /) -> bool | DaskArray:
     """Check values of X to ensure it is count data."""
@@ -763,6 +815,8 @@ def _check_nonnegative_integers_in_mem(x: _MemoryArray, /) -> bool:
     from numbers import Integral
 
     data = x if isinstance(x, np.ndarray) else x.data
+    if np.issubdtype(data.dtype, np.unsignedinteger):
+        return True
     # Check no negatives
     if np.signbit(data).any():
         return False
@@ -772,6 +826,18 @@ def _check_nonnegative_integers_in_mem(x: _MemoryArray, /) -> bool:
     return not np.any((data % 1) != 0)
 
 
+@check_nonnegative_integers.register(HasArrayNamespace)
+def _check_nonnegative_integers_array_api(x: HasArrayNamespace, /) -> bool:
+    xp = array_namespace(x)
+    if xp.isdtype(x.dtype, "unsigned integer"):
+        return True
+    if bool(xp.any(x < 0)):
+        return False
+    if xp.isdtype(x.dtype, "integral"):
+        return True
+    return not bool(xp.any((x % 1) != 0))
+
+
 @check_nonnegative_integers.register(DaskArray)
 def _check_nonnegative_integers_dask(x: DaskArray, /) -> DaskArray:
     return x.map_blocks(check_nonnegative_integers, dtype=bool, drop_axis=(0, 1))

src/scanpy/_utils/_docs.py

@@ -12,7 +12,7 @@
     from typing import Literal
 
 
-__all__ = ["ArrayType", "DaskArray", "Numpy", "ScipySparse", "parse"]
+__all__ = ["ArrayApi", "ArrayType", "DaskArray", "Numpy", "ScipySparse", "parse"]
 
 
 class ArrayType(ABC):
@@ -32,6 +32,16 @@ def rst(self, *, short: bool = False) -> str:  # pragma: no cover
         return f":class:`{'~' if short else ''}{self}`"
 
 
+@dataclass(unsafe_hash=True, frozen=True)
+class ArrayApi(ArrayType):
+    def __str__(self) -> str:  # pragma: no cover
+        return "array-api"
+
+    def rst(self, *, short: bool = False) -> str:  # pragma: no cover
+        # No single class to link to, so link to the standard itself
+        return "`Array API <https://data-apis.org/array-api/latest/>`__"
+
+
 @dataclass(unsafe_hash=True, frozen=True)
 class ScipySparse(ArrayType):
     format: Literal["csr", "csc"]
@@ -46,7 +56,7 @@ def rst(self, *, short: bool = False) -> str:  # pragma: no cover
         )
 
 
-type Inner = Numpy | ScipySparse
+type Inner = Numpy | ScipySparse | ArrayApi
 
 
 @dataclass(unsafe_hash=True, frozen=True)
@@ -79,7 +89,8 @@ def parse(
         yield from (t for t in parse(include) if t not in excluded)
         return
 
-    inner_includes = [i for i in include if not i.startswith("da")]
+    # bare `da` doesn’t imply `da[xp]`, as dask doesn’t support array API chunks yet
+    inner_includes = [i for i in include if not i.startswith(("da", "xp"))]
     for t in include:
         if (
             match := re.fullmatch(r"([^\[]+)(?:\[(.+)\])?", t)
@@ -114,6 +125,11 @@ def _parse_mod(
         case "da":
             for chunk in parse(tags if tags else inner_includes, inner=True):
                 yield DaskArray(chunk=chunk)
+        case "xp":
+            if tags:  # pragma: no cover
+                msg = f"`xp` takes no tags {tags!r}"
+                raise ValueError(msg)
+            yield ArrayApi()
         case _:  # pragma: no cover
             msg = f"invalid module {mod!r}"
             raise ValueError(msg)

src/scanpy/metrics/_common.py

@@ -7,13 +7,12 @@
 
 import numpy as np
 import pandas as pd
+from fast_array_utils.types import HasArrayNamespace
 
 from .._compat import CSRBase, DaskArray, SpBase, fullname, warn
 from .._utils import NeighborsView
 
 if TYPE_CHECKING:
-    from typing import NoReturn
-
     from anndata import AnnData
     from numpy.typing import NDArray
 
@@ -90,10 +89,12 @@ def _resolve_vals[T: NDArray | DaskArray](val: T) -> T: ...
 def _resolve_vals(val: SpBase) -> CSRBase: ...
 @overload
 def _resolve_vals(val: pd.DataFrame | pd.Series) -> NDArray: ...
+@overload
+def _resolve_vals(val: HasArrayNamespace) -> NDArray: ...
 
 
 @singledispatch
-def _resolve_vals(val: object) -> NoReturn:
+def _resolve_vals(val: object):
     msg = f"Unsupported type {type(val)}"
     raise TypeError(msg)
 
@@ -122,6 +123,12 @@ def _(val: pd.DataFrame | pd.Series) -> NDArray:
     return val.to_numpy()
 
 
+@_resolve_vals.register(HasArrayNamespace)
+def _resolve_vals_array_api(val: HasArrayNamespace) -> NDArray:
+    # Moran's I / Geary's C use numba kernels, so convert at the boundary
+    return np.asarray(val)
+
+
 def _vals_heterogeneous[V: NDArray | CSRBase](
     vals: V,
 ) -> tuple[V, NDArray[np.bool] | slice, NDArray[np.float64]]:

src/scanpy/neighbors/__init__.py

@@ -13,6 +13,7 @@
 
 import numpy as np
 import scipy
+from fast_array_utils.types import HasArrayNamespace
 from packaging.version import Version
 from scipy import sparse
 
@@ -592,7 +593,11 @@ def compute_neighbors(
         self._rp_forest = None
         self.n_neighbors = n_neighbors
         self.knn = knn
+
         x = _choose_representation_compat(self._adata, use_rep=use_rep, n_pcs=n_pcs)
+        if isinstance(x, HasArrayNamespace):
+            # sklearn transformers require numpy, so need to convert at boundary
+            x = np.asarray(x)
         self._distances = transformer.fit_transform(x)
         knn_indices, knn_distances = _get_indices_distances_from_sparse_matrix(
             self._distances, n_neighbors

src/scanpy/preprocessing/_highly_variable_genes.py

@@ -10,7 +10,9 @@
 import numpy as np
 import pandas as pd
 from anndata import AnnData
+from array_api_compat import array_namespace
 from fast_array_utils import stats
+from fast_array_utils.types import HasArrayNamespace
 
 from .. import logging as logg
 from .._compat import CSBase, CSRBase, DaskArray, warn
@@ -94,6 +96,18 @@ def _(data_batch: CSBase, clip_val: np.ndarray) -> tuple[np.ndarray, np.ndarray]
     )
 
 
+@clip_square_sum.register(HasArrayNamespace)
+def _[A: HasArrayNamespace](data_batch: A, clip_val: np.ndarray) -> tuple[A, A]:
+    xp = array_namespace(data_batch)
+    batch_counts = xp.astype(data_batch, xp.float64, copy=True)
+    clip_val_broad = xp.broadcast_to(clip_val, batch_counts.shape)
+    # immutability issue, since jax cannot be mutated in place, we just create a new array with where
+    batch_counts = xp.where(batch_counts > clip_val_broad, clip_val_broad, batch_counts)
+    squared_batch_counts_sum = xp.sum(xp.square(batch_counts), axis=0)
+    batch_counts_sum = xp.sum(batch_counts, axis=0)
+    return squared_batch_counts_sum, batch_counts_sum
+
+
 # parallel=False needed for accuracy
 @numba.njit(cache=True, parallel=False)  # noqa: TID251
 def _sum_and_sum_squares_clipped(
@@ -408,12 +422,17 @@ def _highly_variable_genes_single_batch(
         # use out if possible. only possible since we copy the data matrix
         if isinstance(x, np.ndarray):
             np.expm1(x, out=x)
+        elif isinstance(x, HasArrayNamespace):
+            xp = array_namespace(x)
+            x = xp.expm1(x)
         else:
             x = np.expm1(x)
 
     mean, var = materialize_as_ndarray(stats.mean_var(x, axis=0, correction=1))
     # now actually compute the dispersion
-    mean[mean == 0] = 1e-12  # set entries equal to zero to small value
+    # JAX arrays are immutable, so in-place assignment (mean[mean == 0] = ...)
+    # fails; np.where allocates a fresh array instead
+    mean = np.where(mean == 0, 1e-12, mean)  # set zero entries to a small value
     dispersion = var / mean
     if flavor == "seurat":  # logarithmized mean as in Seurat
         dispersion[dispersion == 0] = np.nan

src/scanpy/preprocessing/_normalization.py

@@ -5,6 +5,7 @@
 
 import numba
 import numpy as np
+from array_api_compat import array_namespace
 from fast_array_utils import stats
 from fast_array_utils.numba import njit
 
@@ -15,14 +16,19 @@
 
 if TYPE_CHECKING:
     from anndata import AnnData
+    from fast_array_utils.types import HasArrayNamespace
 
 
-def _compute_nnz_median(counts: np.ndarray | DaskArray) -> np.floating:
+def _compute_nnz_median(
+    counts: np.ndarray | DaskArray | HasArrayNamespace,
+) -> np.floating:
     """Given a 1D array of counts, compute the median of the non-zero counts."""
     if isinstance(counts, DaskArray):
         counts = counts.compute()
+
+    xp = array_namespace(counts)
     counts_greater_than_zero = counts[counts > 0]
-    median = np.median(counts_greater_than_zero)
+    median = xp.median(counts_greater_than_zero)
     return median