WIP at

Author: crusaderky

src/array_api_extra/_lib/_compat.py

@@ -6,14 +6,23 @@
     from ..._array_api_compat_vendor import (  # pyright: ignore[reportMissingImports]
         array_namespace,  # pyright: ignore[reportUnknownVariableType]
         device,  # pyright: ignore[reportUnknownVariableType]
+        is_dask_array,  # pyright: ignore[reportUnknownVariableType]
+        is_jax_array,  # pyright: ignore[reportUnknownVariableType]
+        is_writeable_array,  # pyright: ignore[reportUnknownVariableType]
     )
 except ImportError:
     from array_api_compat import (  # pyright: ignore[reportMissingTypeStubs]
         array_namespace,  # pyright: ignore[reportUnknownVariableType]
         device,
+        is_dask_array,  # pyright: ignore[reportUnknownVariableType]
+        is_jax_array,  # pyright: ignore[reportUnknownVariableType]
+        is_writeable_array,  # pyright: ignore[reportUnknownVariableType,reportAttributeAccessIssue]
     )
 
 __all__ = [
     "array_namespace",
     "device",
+    "is_dask_array",
+    "is_jax_array",
+    "is_writeable_array",
 ]

src/array_api_extra/_lib/_compat.pyi

@@ -11,3 +11,6 @@ def array_namespace(
     use_compat: bool | None = None,
 ) -> ArrayModule: ...
 def device(x: Array, /) -> Device: ...
+def is_dask_array(x: object, /) -> bool: ...
+def is_jax_array(x: object, /) -> bool: ...
+def is_writeable_array(x: object, /) -> bool: ...