capytaine.tools.array_backend module

Helpers to write code that works with any array library implementing the array API standard.

class capytaine.tools.array_backend.IterableTensor(tensor)

Bases: object

Sequence of the sub-arrays along the first axis of an array, or of a list of arrays.

It can be built from a list of arrays of the same shape or from a single array of shape (n, …), which is then kept as it is (no copy, e.g. to apply a FFT across the first axis) and the items are the views array[i, …].

Arrays of the array API standard can not always be iterated, sliced, unpacked or measured with len (unlike the ones of NumPy, PyTorch or JAX), and an index has to cover all the dimensions of the array. This class gives the usual behavior of a sequence in all cases.

as_array(xp)

All the items in a single array of shape (n, …)

class capytaine.tools.array_backend.MatrixLike(*args, **kwargs)

Bases: Protocol

Any 2D (or batched) matrix: an array of any library supported by array-api-compat, or a custom data-sparse matrix such as BlockCirculantMatrix.

Only used for type annotations. __array_namespace__ is not required, since some arrays (e.g. torch tensors) get it from array-api-compat instead of defining it.

property dtype: Any

property shape: Tuple[int, ...]

capytaine.tools.array_backend.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 v2025.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 – The array API compatible namespace corresponding to the arrays in xs.

Return type:

namespace

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

See also

is_array_api_obj, is_numpy_array, is_cupy_array, is_torch_array, is_dask_array, is_jax_array, is_pydata_sparse_array

capytaine.tools.array_backend.complex_dtype(xp, dtype)

Complex dtype with the same precision as dtype.

capytaine.tools.array_backend.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 – a device object (see the Device Support section of the array API specification).

Return type:

device

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.

capytaine.tools.array_backend.ending_dimensions_at_the_beginning(a)

Transform an array of shape (…, n, m) into (n, m, …). Invert of leading_dimensions_at_the_end

capytaine.tools.array_backend.is_array(x)

True if x is an array of any library supported by array-api-compat.

capytaine.tools.array_backend.is_numpy_namespace(xp)

True if xp is the NumPy namespace. (array_api_compat.is_numpy_namespace is only available since array-api-compat 1.9.)

capytaine.tools.array_backend.leading_dimensions_at_the_end(a)

Transform an array of shape (n, m, …) into (…, n, m). Invert of ending_dimensions_at_the_beginning

capytaine.tools.array_backend.split(x, n)

Split an array in n parts of equal sizes along its first axis (replaces np.split, that is not in the array API standard).

capytaine.tools.array_backend.to_backend_of(x, like, *, complex_=False)

Convert x to the array library, device and (complex) precision of like.

capytaine.tools.array_backend.to_numpy(x)

Convert any array (possibly on GPU) to a NumPy array.