fix: allow creating objects with array-api

Author: ilan-gold

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:
@@ -115,6 +115,9 @@ 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(msg)
 

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/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 packaging.version import Version
 from zarr import Array as ZarrArray  # noqa: F401
 from zarr import Group as ZarrGroup
@@ -415,3 +416,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