Array API compatibility library for Python

1.14.0 · active · verified Fri Apr 10

array-api-compat is a small wrapper around popular array libraries like NumPy, CuPy, PyTorch, Dask, JAX, ndonnx, and Sparse. Its primary goal is to make these libraries compatible with the Python Array API standard, allowing array-consuming libraries to write array-agnostic code today, even if the underlying array libraries do not yet fully implement the standard. The library is actively maintained by the Consortium for Python Data API Standards. The current version is 1.14.0, with ongoing development.

Warnings

breaking Support for Python 3.8 was dropped in version 1.15.0.dev0 (March 29, 2024). Users on Python 3.8 will need to upgrade their Python version to use recent `array-api-compat` releases.
Fix: Upgrade to Python 3.9 or newer.
gotcha `array-api-compat` is a wrapper and does not monkey-patch or modify the underlying array objects. This means array object *methods* or *operators* that behave differently from the Array API standard might remain non-compliant. Users should prefer elementwise functions (e.g., `xp.add(x, y)` instead of `x + y`) and provided helper functions (e.g., `array_api_compat.size(x)` instead of `x.size` for PyTorch) for consistent behavior.
Fix: Rely on functions from the `xp` (array_api_compat) namespace rather than direct array object methods/operators for Array API compliant behavior. Use helper functions like `array_api_compat.size()` where native library methods are incompatible (e.g., `torch.Tensor.size`).
gotcha The behavior of `array_api_compat` with NumPy 2.0 has changed. Initially, NumPy 2.0 was unwrapped due to its inherent Array API support, but later versions of `array-api-compat` re-wrapped it for 2023.12 standard compatibility. This could lead to subtle behavioral differences or regressions depending on the specific versions of `array-api-compat` and NumPy in use.
Fix: Consult the `array-api-compat` changelog and test thoroughly when upgrading NumPy 2.x or `array-api-compat` versions to ensure expected behavior.
gotcha `array-api-compat` explicitly avoids implementing functions that are not part of the Array API standard. This design choice keeps the library minimal and prevents it from inadvertently becoming a de facto standard. Consequently, if a function is commonly available in a specific array library but not in the Array API standard, it will not be exposed via `array-api-compat`.
Fix: Be aware that the `xp` namespace will only contain functions specified by the Array API standard. For non-standard functions, you may need to use library-specific imports or `array-api-extra`.
gotcha Dask support within `array-api-compat` has known limitations. Specifically, `sort`, `argsort`, and full support for the `linalg` and `fft` extensions may be missing or non-compliant with the Array API standard. Functions within `dask.array.fft` are typically the original, unwrapped Dask functions.
Fix: Exercise caution when using sorting, linear algebra, or FFT operations with Dask arrays via `array-api-compat`. Verify behavior against Dask's native implementation if unexpected results occur.

Install

pip install array-api-compat Install with pip

Imports

array_namespace
```
from array_api_compat import array_namespace
```
The primary entry point to get an Array API compliant namespace from input arrays.
numpy_namespace
```
import array_api_compat.numpy as xp
```
Directly imports the Array API compatible wrapper for NumPy. The native NumPy namespace (especially pre-2.0) may not be fully compliant with the Array API standard.
torch_namespace
```
import array_api_compat.torch as xp
```
Directly imports the Array API compatible wrapper for PyTorch. The native PyTorch namespace may not be fully compliant with the Array API standard.

Quickstart

This quickstart demonstrates the recommended usage of `array-api-compat` by obtaining an Array API compliant namespace using `array_namespace()`. This namespace (`xp`) can then be used with any supported array type (like NumPy, PyTorch, CuPy, etc.), enabling generic array-agnostic code.

import numpy as np
from array_api_compat import array_namespace

def process_arrays(x, y):
    xp = array_namespace(x, y)
    # Now use xp as the array API compliant namespace
    result = xp.mean(x, axis=0) + 2 * xp.std(y, axis=0)
    return result

# Example usage with NumPy arrays
a = np.array([1.0, 2.0, 3.0])
b = np.array([4.0, 5.0, 6.0])

output = process_arrays(a, b)
print(output)
# Expected output type is the same as input (e.g., numpy.ndarray)
print(type(output))

view raw JSON →