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

docs/conf.py

@@ -191,11 +191,11 @@
     "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.normalize_total": (["np", "sp[csr]", "da", "aa"], []),
     "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", "aa"], []),
     "pp.scrublet": (["np", "sp"], []),
     "pp.scrublet_simulate_doublets": (["np", "sp"], []),
     "tl.dendrogram": (["np", "sp"], []),

docs/extensions/array_support.py

@@ -54,11 +54,22 @@ def run(self) -> list[nodes.Node]:  # noqa: D102
             ))
 
         title = nodes.title("", "", *self.parse_inline(":ref:`array-support`")[0])
-        rows = self._render_support_data(data)
+        rows = [
+            *self._render_support_data(data),
+            self._render_row(
+                self._render_array_type(_docs.ArrayApi()),
+                support=_docs.ArrayApi() in array_types,
+                in_dask=False,
+            ),
+        ]
         return self._render_table(headers, rows, title=title)
 
     def _render_overview(self) -> list[nodes.Node]:
-        headers = ["Function", *(at.rst(short=True) for at in ALL_INNER)]
+        headers = [
+            "Function",
+            *(at.rst(short=True) for at in ALL_INNER),
+            _docs.ArrayApi().rst(short=True),
+        ]
         rows: list[nodes.row] = []
         for fn, (include, exclude) in self._array_support.items():
             row_header, _ = self.parse_inline(f":func:`scanpy.{fn}`")
@@ -71,6 +82,7 @@ def _render_overview(self) -> list[nodes.Node]:
                         ALL_INNER, map(_docs.DaskArray, ALL_INNER), strict=True
                     )
                 ),
+                self._render_support(_docs.ArrayApi() in ats),
             ]
             rows.append(
                 nodes.row(

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."""
@@ -772,6 +824,16 @@ 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 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"]
@@ -79,7 +89,7 @@ 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")]
+    inner_includes = [i for i in include if not i.startswith(("da", "aa"))]
     for t in include:
         if (
             match := re.fullmatch(r"([^\[]+)(?:\[(.+)\])?", t)
@@ -103,6 +113,11 @@ def _parse_mod(
                 msg = f"`np` takes no tags {tags!r}"
                 raise ValueError(msg)
             yield Numpy()
+        case "aa":
+            if tags:  # pragma: no cover
+                msg = f"`aa` takes no tags {tags!r}"
+                raise ValueError(msg)
+            yield ArrayApi()
         case "sp":
             if tags - {"csr", "csc"}:  # pragma: no cover
                 msg = f"invalid tags {tags!r}"

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
@@ -408,12 +410,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