diff --git a/README.md b/README.md index 4b0b0c9c..3d80c708 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ This is a small wrapper around common array libraries that is compatible with the [Array API standard](https://data-apis.org/array-api/latest/). Currently, -NumPy, CuPy, PyTorch, Dask, JAX, ndonnx and `sparse` are supported. If you want +NumPy, CuPy, PyTorch, Dask, JAX, ndonnx, `sparse` and dpnp are supported. If you want support for other array libraries, or if you encounter any issues, please [open an issue](https://github.com/data-apis/array-api-compat/issues). diff --git a/docs/helper-functions.rst b/docs/helper-functions.rst index 155eda9a..d7ed6f6e 100644 --- a/docs/helper-functions.rst +++ b/docs/helper-functions.rst @@ -51,6 +51,7 @@ yet. .. autofunction:: is_jax_array .. autofunction:: is_pydata_sparse_array .. autofunction:: is_ndonnx_array +.. autofunction:: is_dpnp_array .. autofunction:: is_writeable_array .. autofunction:: is_lazy_array .. autofunction:: is_numpy_namespace @@ -60,4 +61,5 @@ yet. .. autofunction:: is_jax_namespace .. autofunction:: is_pydata_sparse_namespace .. autofunction:: is_ndonnx_namespace +.. autofunction:: is_dpnp_namespace .. autofunction:: is_array_api_strict_namespace diff --git a/docs/index.md b/docs/index.md index c5c15174..baf2bee6 100644 --- a/docs/index.md +++ b/docs/index.md @@ -2,7 +2,7 @@ This is a small wrapper around common array libraries that is compatible with the [Array API standard](https://data-apis.org/array-api/latest/). Currently, -NumPy, CuPy, PyTorch, Dask, JAX, ndonnx, and Sparse are supported. If you want +NumPy, CuPy, PyTorch, Dask, JAX, ndonnx, Sparse and dpnp are supported. If you want support for other array libraries, or if you encounter any issues, please [open an issue](https://github.com/data-apis/array-api-compat/issues). diff --git a/docs/supported-array-libraries.md b/docs/supported-array-libraries.md index 46fcdc27..f5d89768 100644 --- a/docs/supported-array-libraries.md +++ b/docs/supported-array-libraries.md @@ -132,6 +132,11 @@ Similar to JAX, `sparse` Array API support is contained directly in `sparse`. Similar to JAX, `ndonnx` Array API support is contained directly in `ndonnx`. +(dpnp-support)= +## [dpnp](https://github.com/IntelPython/dpnp) + +Similar to JAX, `dpnp` Array API support is contained directly in `dpnp`. + (array-api-strict-support)= ## [array-api-strict](https://data-apis.org/array-api-strict/) diff --git a/src/array_api_compat/common/_helpers.py b/src/array_api_compat/common/_helpers.py index d5342658..76397cd4 100644 --- a/src/array_api_compat/common/_helpers.py +++ b/src/array_api_compat/common/_helpers.py @@ -32,6 +32,7 @@ if TYPE_CHECKING: import cupy as cp import dask.array as da + import dpnp import jax import ndonnx as ndx import numpy as np @@ -48,6 +49,8 @@ npt.NDArray[Any] | cp.ndarray | da.Array + | dpnp.ndarray + | dpnp.tensor.usm_ndarray | jax.Array | ndx.Array | sparse.SparseArray @@ -112,6 +115,7 @@ def is_numpy_array(x: object) -> TypeIs[npt.NDArray[Any]]: is_dask_array is_jax_array is_pydata_sparse_array + is_dpnp_array """ # TODO: Should we reject ndarray subclasses? cls = cast(Hashable, type(x)) @@ -141,6 +145,7 @@ def is_cupy_array(x: object) -> bool: is_dask_array is_jax_array is_pydata_sparse_array + is_dpnp_array """ cls = cast(Hashable, type(x)) return _issubclass_fast(cls, "cupy", "ndarray") @@ -163,6 +168,7 @@ def is_torch_array(x: object) -> TypeIs[torch.Tensor]: is_dask_array is_jax_array is_pydata_sparse_array + is_dpnp_array """ cls = cast(Hashable, type(x)) return _issubclass_fast(cls, "torch", "Tensor") @@ -186,6 +192,7 @@ def is_ndonnx_array(x: object) -> TypeIs[ndx.Array]: is_dask_array is_jax_array is_pydata_sparse_array + is_dpnp_array """ cls = cast(Hashable, type(x)) return _issubclass_fast(cls, "ndonnx", "Array") @@ -209,6 +216,7 @@ def is_dask_array(x: object) -> TypeIs[da.Array]: is_ndonnx_array is_jax_array is_pydata_sparse_array + is_dpnp_array """ cls = cast(Hashable, type(x)) return _issubclass_fast(cls, "dask.array", "Array") @@ -233,6 +241,7 @@ def is_jax_array(x: object) -> TypeIs[jax.Array]: is_ndonnx_array is_dask_array is_pydata_sparse_array + is_dpnp_array """ cls = cast(Hashable, type(x)) # We test for jax.core.Tracer here to identify jax arrays during jit tracing. From jax 0.8.2 on, @@ -267,12 +276,41 @@ def is_pydata_sparse_array(x: object) -> TypeIs[sparse.SparseArray]: is_ndonnx_array is_dask_array is_jax_array + is_dpnp_array """ # TODO: Account for other backends. cls = cast(Hashable, type(x)) return _issubclass_fast(cls, "sparse", "SparseArray") +def is_dpnp_array(x: object) -> bool: + """ + Return True if `x` is a dpnp array. + + This function does not import `dpnp` if it has not already been imported + and is therefore cheap to use. + + + See Also + -------- + + array_namespace + is_array_api_obj + is_numpy_array + is_cupy_array + is_torch_array + is_ndonnx_array + is_dask_array + is_jax_array + is_pydata_sparse_array + """ + cls = cast(Hashable, type(x)) + return ( + _issubclass_fast(cls, "dpnp", "ndarray") + or _issubclass_fast(cls, "dpnp.tensor", "usm_ndarray") + ) + + def is_array_api_obj(x: object) -> TypeGuard[_ArrayApiObj]: """ Return True if `x` is an array API compatible array object. @@ -287,6 +325,8 @@ def is_array_api_obj(x: object) -> TypeGuard[_ArrayApiObj]: is_ndonnx_array is_dask_array is_jax_array + is_pydata_sparse_array + is_dpnp_array """ try: # TODO: drop this check after np.matrix is gone @@ -338,6 +378,7 @@ def is_numpy_namespace(xp: Namespace) -> bool: is_dask_namespace is_jax_namespace is_pydata_sparse_namespace + is_dpnp_namespace is_array_api_strict_namespace """ return xp.__name__ in {"numpy", _compat_module_name() + ".numpy"} @@ -360,6 +401,7 @@ def is_cupy_namespace(xp: Namespace) -> bool: is_dask_namespace is_jax_namespace is_pydata_sparse_namespace + is_dpnp_namespace is_array_api_strict_namespace """ return xp.__name__ in {"cupy", _compat_module_name() + ".cupy"} @@ -382,6 +424,7 @@ def is_torch_namespace(xp: Namespace) -> bool: is_dask_namespace is_jax_namespace is_pydata_sparse_namespace + is_dpnp_namespace is_array_api_strict_namespace """ return xp.__name__ in {"torch", _compat_module_name() + ".torch"} @@ -401,6 +444,7 @@ def is_ndonnx_namespace(xp: Namespace) -> bool: is_dask_namespace is_jax_namespace is_pydata_sparse_namespace + is_dpnp_namespace is_array_api_strict_namespace """ return xp.__name__ == "ndonnx" @@ -423,6 +467,7 @@ def is_dask_namespace(xp: Namespace) -> bool: is_ndonnx_namespace is_jax_namespace is_pydata_sparse_namespace + is_dpnp_namespace is_array_api_strict_namespace """ return xp.__name__ in {"dask.array", _compat_module_name() + ".dask.array"} @@ -445,6 +490,7 @@ def is_jax_namespace(xp: Namespace) -> bool: is_ndonnx_namespace is_dask_namespace is_pydata_sparse_namespace + is_dpnp_namespace is_array_api_strict_namespace """ return xp.__name__ in {"jax.numpy", "jax.experimental.array_api"} @@ -464,11 +510,34 @@ def is_pydata_sparse_namespace(xp: Namespace) -> bool: is_ndonnx_namespace is_dask_namespace is_jax_namespace + is_dpnp_namespace is_array_api_strict_namespace """ return xp.__name__ == "sparse" +def is_dpnp_namespace(xp: Namespace) -> bool: + """ + Returns True if `xp` is a dpnp namespace. + + This includes `dpnp` and `dpnp.tensor`. + + See Also + -------- + + array_namespace + is_numpy_namespace + is_cupy_namespace + is_torch_namespace + is_ndonnx_namespace + is_dask_namespace + is_jax_namespace + is_pydata_sparse_namespace + is_array_api_strict_namespace + """ + return xp.__name__ in {"dpnp", "dpnp.tensor"} + + def is_array_api_strict_namespace(xp: Namespace) -> bool: """ Returns True if `xp` is an array-api-strict namespace. @@ -484,6 +553,7 @@ def is_array_api_strict_namespace(xp: Namespace) -> bool: is_dask_namespace is_jax_namespace is_pydata_sparse_namespace + is_dpnp_namespace """ return xp.__name__ == "array_api_strict" @@ -645,6 +715,7 @@ def your_function(x, y): is_dask_array is_jax_array is_pydata_sparse_array + is_dpnp_array """ namespaces: list[Namespace] = [] @@ -1077,6 +1148,8 @@ def is_lazy_array(x: object) -> TypeGuard[_ArrayApiObj]: "is_cupy_namespace", "is_dask_array", "is_dask_namespace", + "is_dpnp_array", + "is_dpnp_namespace", "is_jax_array", "is_jax_namespace", "is_numpy_array", diff --git a/tests/_helpers.py b/tests/_helpers.py index 17865aa0..b12ff720 100644 --- a/tests/_helpers.py +++ b/tests/_helpers.py @@ -4,7 +4,7 @@ wrapped_libraries = ["numpy", "cupy", "torch", "dask.array"] all_libraries = wrapped_libraries + [ - "array_api_strict", "jax.numpy", "ndonnx", "sparse" + "array_api_strict", "jax.numpy", "ndonnx", "sparse", "dpnp" ] def import_(library, wrapper=False): diff --git a/tests/test_array_namespace.py b/tests/test_array_namespace.py index 7f196387..e56b8eb4 100644 --- a/tests/test_array_namespace.py +++ b/tests/test_array_namespace.py @@ -25,7 +25,8 @@ def test_array_namespace(request, library, api_version, use_compat): return if (library == "sparse" and api_version in ("2023.12", "2024.12")) or ( library == "jax.numpy" and - api_version in ("2021.12", "2022.12", "2023.12", "2024.12") + api_version in ("2021.12", "2022.12", "2023.12", "2024.12")) or ( + library == "dpnp" and api_version in ("2021.12", "2022.12", "2023.12") ): xfail(request, "Unsupported API version") diff --git a/tests/test_common.py b/tests/test_common.py index 1ebf6e7a..ff7c3d3f 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -8,10 +8,11 @@ from array_api_compat import ( # noqa: F401 is_numpy_array, is_cupy_array, is_torch_array, is_dask_array, is_jax_array, is_pydata_sparse_array, - is_ndonnx_array, + is_ndonnx_array, is_dpnp_array, is_numpy_namespace, is_cupy_namespace, is_torch_namespace, is_dask_namespace, is_jax_namespace, is_pydata_sparse_namespace, is_array_api_strict_namespace, is_ndonnx_namespace, + is_dpnp_namespace, ) from array_api_compat import ( @@ -29,6 +30,7 @@ 'jax.numpy': 'is_jax_array', 'sparse': 'is_pydata_sparse_array', 'ndonnx': 'is_ndonnx_array', + 'dpnp': 'is_dpnp_array', } is_namespace_functions = { @@ -40,6 +42,7 @@ 'sparse': 'is_pydata_sparse_namespace', 'array_api_strict': 'is_array_api_strict_namespace', 'ndonnx': 'is_ndonnx_namespace', + 'dpnp': 'is_dpnp_namespace', } @@ -195,6 +198,8 @@ def test_device_to_device(library, request): xfail(request, reason="Stub raises ValueError") if library == "sparse": xfail(request, reason="No __array_namespace_info__()") + if library == "dpnp": + xfail(request, reason="`dev` not in `devices`") if library == "array_api_strict": if np.__version__ < "2": xfail(request, reason="no copy argument of np.asarray") @@ -259,6 +264,8 @@ def test_asarray_cross_library(source_library, target_library, request): pytest.skip(reason="cupy does not support implicit conversion to CPU") elif source_library == "sparse" and target_library != "sparse": pytest.skip(reason="`sparse` does not allow implicit densification") + elif source_library == "dpnp" and target_library != "dpnp": + pytest.skip(reason="dpnp does not allow implicit conversion") src_lib = import_(source_library, wrapper=True) tgt_lib = import_(target_library, wrapper=True)