Array API compatibility library#

This is a small wrapper around common array libraries that is compatible with the Array API standard. Currently, NumPy, CuPy, PyTorch, Dask, and JAX are supported. If you want support for other array libraries, or if you encounter any issues, please open an issue.

Note that some of the functionality in this library is backwards incompatible with the corresponding wrapped libraries. The end-goal is to eventually make each array library itself fully compatible with the array API, but this requires making backwards incompatible changes in many cases, so this will take some time.

Currently all libraries here are implemented against the 2022.12 version of the standard.

Installation#

array-api-compat is available on both PyPI

python -m pip install array-api-compat

and conda-forge

conda install --channel conda-forge array-api-compat

Usage#

The typical usage of this library will be to get the corresponding array API compliant namespace from the input arrays using array_namespace(), like

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)

If you wish to have library-specific code-paths, you can import the corresponding wrapped namespace for each library, like

import array_api_compat.numpy as np

import array_api_compat.cupy as cp

import array_api_compat.torch as torch

import array_api_compat.dask as da

Note

There is no array_api_compat.jax submodule. JAX support is contained in JAX itself in the jax.experimental.array_api module. array-api-compat simply wraps that submodule. The main JAX support in this module consists of supporting it in the helper functions.

Each will include all the functions from the normal NumPy/CuPy/PyTorch/dask.array namespace, except that functions that are part of the array API are wrapped so that they have the correct array API behavior. In each case, the array object used will be the same array object from the wrapped library.

Difference between array_api_compat and array_api_strict#

array_api_strict is a strict minimal implementation of the array API standard, formerly known as numpy.array_api (see NEP 47). For example, array_api_strict does not include any functions that are not part of the array API specification, and will explicitly disallow behaviors that are not required by the spec (e.g., cross-kind type promotions). (cupy.array_api is similar to array_api_strict)

array_api_compat, on the other hand, is just an extension of the corresponding array library namespaces with changes needed to be compliant with the array API. It includes all additional library functions not mentioned in the spec, and allows any library behaviors not explicitly disallowed by it, such as cross-kind casting.

In particular, unlike array_api_strict, this package does not use a separate Array object, but rather just uses the corresponding array library array objects (numpy.ndarray, cupy.ndarray, torch.Tensor, etc.) directly. This is because those are the objects that are going to be passed as inputs to functions by end users. This does mean that a few behaviors cannot be wrapped (see below), but most of the array API functional, so this does not affect most things.

Array consuming library authors coding against the array API may wish to test against array_api_strict to ensure they are not using functionality outside of the standard, but prefer this implementation for the default behavior for end-users.

Vendoring#

This library supports vendoring as an installation method. To vendor the library, simply copy array_api_compat into the appropriate place in the library, like

cp -R array_api_compat/ mylib/vendored/array_api_compat

You may also rename it to something else if you like (nowhere in the code references the name “array_api_compat”).

Alternatively, the library may be installed as dependency on PyPI.

Scope#

At this time, the scope of array-api-compat is limited to wrapping array libraries so that they can comply with the array API standard. This includes a small set of helper functions which may be useful to most users of array-api-compat, for instance, functions that provide meta-functionality to aid in supporting the array API, or functions that are necessary to work around wrapping limitations for certain libraries.

Things that are out-of-scope include:

  • functions that have not yet been standardized (although note that functions that are in a draft version of the standard are in scope),

  • functions that are complicated to implement correctly/maintain,

  • anything that requires the use of non-Python code.

If you want a function that is not in array-api-compat that isn’t part of the standard, you should request it either for inclusion in the standard or in specific array libraries.

Why is the scope limited in this way? Firstly, we want to keep array-api-compat as primarily a polyfill compatibility shim. The goal is to let consuming libraries use the array API today, even with array libraries that do not yet fully support it. In an ideal world—one that we hope to eventually see in the future—array-api-compat would be unnecessary, because every array library would fully support the standard.

The inclusion of non-standardized functions in array-api-compat would undermine this goal. But much more importantly, it would also undermine the goals of the Data APIs Consortium. The Consortium creates the array API standard via the consensus of stakeholders from various array libraries and users. If a not-yet-standardized function were included in array-api-compat, it would become de facto standard, bypassing the decision making processes of the Consortium.

Secondly, we want to keep array-api-compat as minimal as possible, so that it is easy for libraries to add as a (possibly vendored) dependency.

Thirdly, array-api-compat has a relatively small development team. Pull requests to array-api-compat would not necessarily receive the same stringent level of scrutiny that changes to established array libraries like NumPy or PyTorch would. For wrapped standard functions, this is fine, since the wrappers typically just clean up a few small inconsistencies from the standard, leaving the complexity of the implementation to the base array library function. Furthermore, standard functions are tested by the rigorous array-api-tests test suite. For this reason, functions that require complex implementations are generally out-of-scope and should be preferred to be implemented in upstream array libraries.