Backend-native Implementation

pyproject.toml

@@ -101,7 +101,7 @@ test-min = [
     "pyarrow<21",        # https://github.com/scikit-hep/awkward/issues/3579
     "anndata[dask]",
 ]
-test = [ "anndata[test-min,lazy]" ]
+test = [ "anndata[test-min,lazy]", "jax", "jaxlib" ] # TODO: remove jax? own extra?
 gpu = [ "cupy" ]
 cu12 = [ "cupy-cuda12x" ]
 cu11 = [ "cupy-cuda11x" ]

src/anndata/_core/index.py

@@ -10,7 +10,7 @@
 import pandas as pd
 from scipy.sparse import issparse
 
-from ..compat import AwkArray, CSArray, CSMatrix, DaskArray, XDataArray
+from ..compat import AwkArray, CSArray, CSMatrix, DaskArray, XDataArray, has_xp
 from .xarray import Dataset2D
 
 if TYPE_CHECKING:
@@ -108,8 +108,11 @@ def name_idx(i):
         if isinstance(indexer.data, DaskArray):
             return indexer.data.compute()
         return indexer.data
+    elif has_xp(indexer):
+        msg = "Need to implement array api-based indexing"
+        raise NotImplementedError(msg)
     msg = f"Unknown indexer {indexer!r} of type {type(indexer)}"
-    raise IndexError()
+    raise IndexError(msg)
 
 
 def _fix_slice_bounds(s: slice, length: int) -> slice:

src/anndata/_core/merge.py

@@ -658,7 +658,12 @@ def _apply_to_array(self, el, *, axis, fill_value=None):
 
         indexer = self.idx
 
-        # Indexes real fast, and does outer indexing
+        # Fallback to numpy: keep pandas
+        # Force to NumPy (materializes JAX/Cubed); fine for small tests,
+        # but may be slow or fail on large/lazy arrays
+        if not isinstance(el, np.ndarray):
+            el = np.asarray(el)  # fine for jax-in-cpu tests
+
         return pd.api.extensions.take(
             el, indexer, axis=axis, allow_fill=True, fill_value=fill_value
         )
@@ -1399,6 +1404,37 @@ def concat_dataset2d_on_annot_axis(
     return ds_concat_2d
 
 
+def _is_sparse(x):
+    try:
+        return scipy.sparse.issparse(x)
+    except TypeError:
+        return False
+
+
+def _to_numpy_if_array_api(x):
+    if isinstance(x, np.ndarray | pd.DataFrame | pd.Series | DaskArray) or _is_sparse(
+        x
+    ):
+        return x
+    try:
+        import array_api_compat as aac
+
+        # If this succeeds, it's an array-API array (e.g. JAX, cubed, cupy, dask)
+        aac.array_namespace(x)
+        return np.asarray(x)
+    except TypeError:
+        # Not an array-API object (or lib not available) → return unchanged
+        return x
+
+
+def _normalize_nested(obj):
+    if isinstance(obj, dict):
+        return {k: _normalize_nested(v) for k, v in obj.items()}
+    if isinstance(obj, list | tuple):
+        return type(obj)(_normalize_nested(v) for v in obj)
+    return _to_numpy_if_array_api(obj)
+
+
 def concat(  # noqa: PLR0912, PLR0913, PLR0915
     adatas: Collection[AnnData] | Mapping[str, AnnData],
     *,
@@ -1759,6 +1795,11 @@ def concat(  # noqa: PLR0912, PLR0913, PLR0915
     )
     uns = uns_merge([a.uns for a in adatas])
 
+    # TODO: try pandas extension arrays after concat errors are fixed
+    # converting to numpy since pandas does not support array-API arrays
+    # normalizes uns (handles JAX / array-API arrays nested in dicts/lists)
+    uns = _normalize_nested(uns)
+
     raw = None
     has_raw = [a.raw is not None for a in adatas]
     if all(has_raw):
@@ -1785,6 +1826,13 @@ def concat(  # noqa: PLR0912, PLR0913, PLR0915
             "not concatenating `.raw` attributes."
         )
         warn(msg, UserWarning, stacklevel=2)
+
+    layers = _normalize_nested(layers)
+    concat_mapping = _normalize_nested(concat_mapping)
+    alt_mapping = _normalize_nested(alt_mapping)
+    concat_pairwise = _normalize_nested(concat_pairwise)
+    alt_pairwise = _normalize_nested(alt_pairwise)
+
     return AnnData(
         **{
             "X": X,

src/anndata/_core/storage.py

@@ -10,7 +10,7 @@
 from anndata.compat import CSArray, CSMatrix
 
 from .._warnings import ImplicitModificationWarning
-from ..compat import XDataset
+from ..compat import XDataset, has_xp
 from ..utils import (
     ensure_df_homogeneous,
     join_english,
@@ -67,6 +67,8 @@ def coerce_array(
             return np.array(value)
         except (ValueError, TypeError) as _e:
             e = _e
+    if has_xp(value):
+        return value
     # if value isn’t the right type or convertible, raise an error
     msg = f"{name} needs to be of one of {join_english(map(str, array_data_structure_types))}, not {type(value)}."
     if e is not None:

src/anndata/_core/views.py

@@ -21,6 +21,7 @@
     CupyCSRMatrix,
     DaskArray,
     ZappyArray,
+    has_xp,
 )
 from .access import ElementRef
 from .xarray import Dataset2D
@@ -296,6 +297,9 @@ def __setattr__(self, key: str, value: Any):
 
 @singledispatch
 def as_view(obj, view_args):
+    if has_xp(obj):
+        # TODO: Determine if we need some sort of specific view object for array-api
+        return obj
     msg = f"No view type has been registered for {type(obj)}"
     raise NotImplementedError(msg)
 

src/anndata/_io/specs/registry.py

@@ -8,6 +8,12 @@
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Generic, TypeVar
 
+import numpy as np
+import pandas as pd
+from pandas.api.extensions import ExtensionArray
+from scipy import sparse as sp
+
+from anndata import AnnData
 from anndata._io.utils import report_read_key_on_error, report_write_key_on_error
 from anndata._types import Read, ReadLazy, _ReadInternal, _ReadLazyInternal
 from anndata.compat import DaskArray, ZarrGroup, _read_attr, is_zarr_v2
@@ -34,6 +40,48 @@
     LazyDataStructures = DaskArray | Dataset2D | CategoricalArray | MaskedArray
 
 
+def is_sparse_like(x):
+    try:
+        return sp.issparse(x)
+    except AttributeError:
+        return False
+
+
+def to_numpy_if_array_api(x):
+    if isinstance(
+        x,
+        np.ndarray
+        | np.generic
+        | pd.DataFrame
+        | pd.Series
+        | pd.Index
+        | ExtensionArray
+        | DaskArray
+        | sp.spmatrix
+        | AnnData,
+    ):
+        return x
+
+    # Try array-API detection only for unknown leaves
+    try:
+        import array_api_compat as aac
+
+        # If this succeeds, it's an array-API array (e.g. JAX, CuPy, torch, …)
+        aac.array_namespace(x)
+        return np.asarray(x)
+    except (ImportError, AttributeError, TypeError):
+        # Not an array-API object (or not supported), so return unchanged
+        return x
+
+
+def normalize_nested(obj):
+    if isinstance(obj, dict):
+        return {k: normalize_nested(v) for k, v in obj.items()}
+    if isinstance(obj, list | tuple):
+        return type(obj)(normalize_nested(v) for v in obj)
+    return to_numpy_if_array_api(obj)
+
+
 # TODO: This probably should be replaced by a hashable Mapping due to conversion b/w "_" and "-"
 # TODO: Should filetype be included in the IOSpec if it changes the encoding? Or does the intent that these things be "the same" overrule that?
 @dataclass(frozen=True)
@@ -386,6 +434,10 @@ def write_elem(
         elif k in store:
             del store[k]
 
+        # Normalize array-API (e.g., JAX/CuPy) payloads buried in mappings/lists
+        if not isinstance(elem, AnnData):
+            elem = normalize_nested(elem)
+
         write_func = self.find_write_func(dest_type, elem, modifiers)
 
         if self.callback is None:

src/anndata/compat/__init__.py

@@ -12,6 +12,7 @@
 import numpy as np
 import pandas as pd
 import scipy
+from array_api_compat import get_namespace as array_api_get_namespace
 from numpy.typing import NDArray
 from packaging.version import Version
 from zarr import Array as ZarrArray  # noqa: F401
@@ -436,3 +437,11 @@ def _map_cat_to_str(cat: pd.Categorical) -> pd.Categorical:
         return cat.map(str, na_action="ignore")
     else:
         return cat.map(str)
+
+
+def has_xp(mod):
+    try:
+        array_api_get_namespace(mod)
+        return True
+    except TypeError:
+        return False