Abstract your array operations (autoray)

0.8.10 · active · verified Sun Apr 12

Autoray is a lightweight Python library designed to abstract array and tensor operations, enabling users to write backend-agnostic numeric code. It provides an automatic dispatch mechanism that works across various array libraries such as NumPy, PyTorch, JAX, TensorFlow, CuPy, Dask, and more, as long as they provide a NumPy-ish API. This allows for swapping custom functions, lazy computation tracing, and unified compilation interfaces. The library is actively maintained, with its current version being 0.8.10, and sees a continuous release cadence.

Warnings

breaking The iteration behavior of `LazyArray.__iter__` changed in version `0.8.0`. It now iterates over slices of the array rather than the computational graph nodes, which can break code relying on previous lazy graph inspection.
Fix: Review any code iterating directly over `LazyArray` instances and adapt it to the new slicing behavior or use `LazyArray.nodes` if access to graph nodes is explicitly required.
breaking The minimum required Python version was bumped to `3.10` in version `0.8.3`. Users on older Python versions (e.g., 3.9 or earlier) will need to upgrade their Python environment to use `autoray` versions `0.8.3` and higher.
Fix: Upgrade your Python environment to version 3.10 or newer, or pin `autoray` to a version less than `0.8.3` if an upgrade is not feasible.
gotcha When using `autoray.do('linalg.svd', x)`, the `full_matrices` argument defaults to `False`. This differs from NumPy's default behavior, where `full_matrices` is `True`. This can lead to unexpected output shapes if not explicitly handled.
Fix: Always explicitly specify `full_matrices=True` or `full_matrices=False` when calling `ar.do('linalg.svd', ...)` to ensure consistent behavior across backends and avoid surprises, aligning with the desired output.
gotcha Autoray performs internal translations for certain functions to match backend-specific APIs (e.g., NumPy's `sum` becomes TensorFlow's `tf.reduce_sum`). While designed for compatibility, these implicit translations can lead to subtle differences in behavior or performance if not fully understood.
Fix: Familiarize yourself with the `autoray` documentation regarding backend deviations and translations, especially when encountering unexpected results or performance characteristics with specific backend libraries. For critical operations, consider direct backend calls if fine-grained control is necessary.

Install

pip install autoray Install stable version

Imports

autoray
```
import autoray as ar
```
The most common and recommended way to import autoray, typically aliased as `ar` for convenience, to access functions like `ar.do` or `ar.get_namespace`.
numpy
```
from autoray import numpy as np
```
Imports a NumPy-like API that dispatches through autoray, allowing for a drop-in replacement for existing NumPy code.
do
```
from autoray import do
```
While `ar.do` is preferred with the alias, `do` can be imported directly. `autoray.do` is not the correct import path for the function itself.

Quickstart

This quickstart demonstrates the core functionalities of `autoray`: automatic backend dispatch using `ar.do()`, explicit backend selection with the `like` argument, and obtaining a backend-specific API using `ar.get_namespace()`. It also shows how to define a function that works generically across different array backends.

import autoray as ar
import numpy as np

# Basic usage with automatic dispatch (inferred from array type)
x_np = np.random.uniform(size=)
y_np = ar.do('sqrt', x_np)
print(f"Numpy sqrt: {y_np}, type: {type(y_np)}")

# Using 'like' argument for explicit backend or inference
# If torch is not installed, this will silently fall back to numpy behavior
# For full torch functionality, ensure 'torch' is installed.
x_torch_like = ar.do('random.uniform', size=(10, 10), like="torch")
print(f"Array generated with 'like="torch"': {type(x_torch_like)}")

# Using get_namespace for a backend-specific API (Python Array API style)
try:
    # Attempt to get a torch namespace
    xp = ar.get_namespace(like="torch") # Requires 'torch' to be installed for actual torch arrays
    z = xp.ones((3, 4), dtype=xp.float32)
    result = xp.exp(z)
    print(f"Torch-like exp result shape: {result.shape}, type: {type(result)}")
except (ImportError, TypeError):
    # Fallback if torch is not installed, get numpy namespace
    xp = ar.get_namespace(like="numpy")
    z = xp.ones((3, 4), dtype=xp.float32)
    result = xp.exp(z)
    print(f"Numpy-like exp result shape (fallback): {result.shape}, type: {type(result)}")

# Example of a more complex operation with automatic dispatch
def noised_svd(x):
    U, s, VH = ar.do('linalg.svd', x)
    sn = s + 0.1 * ar.do('random.normal', size=ar.shape(s), like=s)
    return ar.do('einsum', 'ij,j,jk->ik', U, sn, VH)

# Use a numpy array for demonstration
x_complex_op = np.random.rand(10, 10)
y_complex_op = noised_svd(x_complex_op)
print(f"Complex operation result shape: {y_complex_op.shape}")

view raw JSON →