Helper Functions¶

In addition to the wrapped library namespaces and functions in the array API specification, there are several helper functions included here that aren’t part of the specification but which are useful for using the array API:

Entry-point Helpers¶

The array_namespace() function is the primary entry-point for array API consuming libraries.

array_api_compat.array_namespace(*xs: object | complex | None, api_version: str | None = None, use_compat: bool | None = None) → ModuleType¶

Get the array API compatible namespace for the arrays xs.

Parameters:

xs (arrays) – one or more arrays. xs can also be Python scalars (bool, int, float, complex, or None), which are ignored.
api_version (str) – The newest version of the spec that you need support for (currently the compat library wrapped APIs support v2024.12).
use_compat (bool or None) – If None (the default), the native namespace will be returned if it is already array API compatible, otherwise a compat wrapper is used. If True, the compat library wrapped library will be returned. If False, the native library namespace is returned.

Returns:

out (namespace) – The array API compatible namespace corresponding to the arrays in xs.

Raises:

TypeError – If xs contains arrays from different array libraries or contains a non-array.

Typical usage is to pass the arguments of a function to array_namespace() at the top of a function to get the corresponding array API namespace:

def your_function(x, y):
    xp = array_api_compat.array_namespace(x, y)
    # Now use xp as the array library namespace
    return xp.mean(x, axis=0) + 2*xp.std(y, axis=0)

Wrapped array namespaces can also be imported directly. For example, array_namespace(np.array(...)) will return array_api_compat.numpy. This function will also work for any array library not wrapped by array-api-compat if it explicitly defines __array_namespace__ (the wrapped namespace is always preferred if it exists).

array_api_compat.is_array_api_obj(x: object) → TypeGuard[_ArrayApiObj]¶: Return True if x is an array API compatible array object.

See also

array_namespace, is_numpy_array, is_cupy_array, is_torch_array, is_ndonnx_array, is_dask_array, is_jax_array

Array Method Helpers¶

array-api-compat does not attempt to wrap or monkey patch the array object for any library. Consequently, any API differences for the array object cannot be directly wrapped. Some libraries do not define some of these methods or define them differently. For these, helper functions are provided which can be used instead.

Note that if you have a compatibility issue with an operator method (like __add__, i.e., +) you can prefer to use the corresponding elementwise function instead, which would be wrapped.

array_api_compat.device(x: _ArrayApiObj, /) → Device¶

Hardware device the array data resides on.

This is equivalent to x.device according to the standard. This helper is included because some array libraries either do not have the device attribute or include it with an incompatible API.

Parameters:: x (array) – array instance from an array API compatible library.
Returns:: out (device) – a device object (see the Device Support section of the array API specification).

Notes

For NumPy the device is always "cpu". For Dask, the device is always a special DASK_DEVICE object.

See also

to_device: Move array data to a different device.

array_api_compat.to_device(x: object, device: object, /, *, stream: int | Any | None = None) → object¶

Copy the array from the device on which it currently resides to the specified device.

This is equivalent to x.to_device(device, stream=stream) according to the standard. This helper is included because some array libraries do not have the to_device method.

Parameters:

x (array) – array instance from an array API compatible library.
device (device) – a device object (see the Device Support section of the array API specification).
stream (int | Any | None) – stream object to use during copy. In addition to the types supported in array.__dlpack__, implementations may choose to support any library-specific stream object with the caveat that any code using such an object would not be portable.

Returns:

out (array) – an array with the same data and data type as x and located on the specified device.

Notes

For NumPy, this function effectively does nothing since the only supported device is the CPU. For CuPy, this method supports CuPy CUDA Device and Stream objects. For PyTorch, this is the same as x.to(device) (the stream argument is not supported in PyTorch).

See also

device: Hardware device the array data resides on.

array_api_compat.size(x: HasShape[Collection[SupportsIndex]]) → int¶

array_api_compat.size(x: HasShape[Collection[SupportsIndex | None]]) → int | None

Return the total number of elements of x.

This is equivalent to x.size according to the standard.

This helper is included because PyTorch defines size in an incompatible way. It also fixes dask.array’s behaviour which returns nan for unknown sizes, whereas the standard requires None.

Inspection Helpers¶

These convenience functions can be used to test if an array or namespace comes from a specific library without importing that library if it hasn’t been imported yet.

array_api_compat.is_numpy_array(x: object) → TypeIs[npt.NDArray[Any]]¶

Return True if x is a NumPy array.

This function does not import NumPy if it has not already been imported and is therefore cheap to use.

This also returns True for ndarray subclasses and NumPy scalar objects.

array_api_compat.is_cupy_array(x: object) → bool¶

Return True if x is a CuPy array.

This function does not import CuPy if it has not already been imported and is therefore cheap to use.

This also returns True for cupy.ndarray subclasses and CuPy scalar objects.

array_api_compat.is_torch_array(x: object) → TypeIs[torch.Tensor]¶

Return True if x is a PyTorch tensor.

This function does not import PyTorch if it has not already been imported and is therefore cheap to use.

array_api_compat.is_dask_array(x: object) → TypeIs[da.Array]¶

Return True if x is a dask.array Array.

This function does not import dask if it has not already been imported and is therefore cheap to use.

array_api_compat.is_jax_array(x: object) → TypeIs[jax.Array]¶

Return True if x is a JAX array.

This function does not import JAX if it has not already been imported and is therefore cheap to use.

array_api_compat.is_pydata_sparse_array(x: object) → TypeIs[sparse.SparseArray]¶

Return True if x is an array from the sparse package.

This function does not import sparse if it has not already been imported and is therefore cheap to use.

array_api_compat.is_ndonnx_array(x: object) → TypeIs[ndx.Array]¶

Return True if x is a ndonnx Array.

This function does not import ndonnx if it has not already been imported and is therefore cheap to use.

array_api_compat.is_writeable_array(x: object) → TypeGuard[_ArrayApiObj]¶: Return False if x.__setitem__ is expected to raise; True otherwise. Return False if x is not an array API compatible object.

Warning

As there is no standard way to check if an array is writeable without actually writing to it, this function blindly returns True for all unknown array types.

array_api_compat.is_lazy_array(x: object) → TypeGuard[_ArrayApiObj]¶

Return True if x is potentially a future or it may be otherwise impossible or expensive to eagerly read its contents, regardless of their size, e.g. by calling bool(x) or float(x).

Return False otherwise; e.g. bool(x) etc. is guaranteed to succeed and to be cheap as long as the array has the right dtype and size.

Note

This function errs on the side of caution for array types that may or may not be lazy, e.g. JAX arrays, by always returning True for them.

array_api_compat.is_numpy_namespace(xp: ModuleType) → bool¶

Returns True if xp is a NumPy namespace.

This includes both NumPy itself and the version wrapped by array-api-compat.

array_api_compat.is_cupy_namespace(xp: ModuleType) → bool¶

Returns True if xp is a CuPy namespace.

This includes both CuPy itself and the version wrapped by array-api-compat.

array_api_compat.is_torch_namespace(xp: ModuleType) → bool¶

Returns True if xp is a PyTorch namespace.

This includes both PyTorch itself and the version wrapped by array-api-compat.

array_api_compat.is_dask_namespace(xp: ModuleType) → bool¶

Returns True if xp is a Dask namespace.

This includes both dask.array itself and the version wrapped by array-api-compat.

array_api_compat.is_jax_namespace(xp: ModuleType) → bool¶

Returns True if xp is a JAX namespace.

This includes jax.numpy and jax.experimental.array_api which existed in older versions of JAX.

array_api_compat.is_pydata_sparse_namespace(xp: ModuleType) → bool¶: Returns True if xp is a pydata/sparse namespace.

See also

array_namespace, is_numpy_namespace, is_cupy_namespace, is_torch_namespace, is_ndonnx_namespace, is_dask_namespace, is_jax_namespace, is_array_api_strict_namespace

array_api_compat.is_ndonnx_namespace(xp: ModuleType) → bool¶: Returns True if xp is an NDONNX namespace.

See also

array_namespace, is_numpy_namespace, is_cupy_namespace, is_torch_namespace, is_dask_namespace, is_jax_namespace, is_pydata_sparse_namespace, is_array_api_strict_namespace

array_api_compat.is_array_api_strict_namespace(xp: ModuleType) → bool¶: Returns True if xp is an array-api-strict namespace.

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