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, api_version=None, use_compat=None)¶

Get the array API compatible namespace for the arrays xs.

Parameters:

xs (arrays) – one or more arrays.
api_version (str) – The newest version of the spec that you need support for (currently the compat library wrapped APIs support v2022.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)¶: 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_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: Array, /) → 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: Array, device: Device, /, *, stream: int | Any | None = None) → Array¶

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 (Optional[Union[int, Any]]) – 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)¶

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.

Inspection Helpers¶

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

array_api_compat.is_numpy_array(x)¶

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)¶

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)¶

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)¶

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)¶

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.