jcmgray/autoray

autoray is a lightweight python AUTOmatic-arRAY library for abstracting your tensor operations. Primarily it provides an automatic dispatch mechanism that means you can write backend agnostic code that works for:

• numpy
• pytorch
• jax
• cupy
• mlx
• dask
• autograd
• tensorflow
• sparse
• paddle
• ... and indeed any library that provides a numpy-ish api, even if it knows nothing about autoray.

Beyond that, abstracting the array interface allows you to:

• swap custom versions of functions for specific backends
• trace through computations lazily without actually running them
• automatically share intermediates and fold constants in computations
• compile functions with a unified interface for different backends

Basic usage

The main function of autoray is do, which takes a function name followed by *args and **kwargs, and automatically looks up (and caches) the correct function to match the equivalent numpy call:

import autoray as ar


def noised_svd(x):
    # automatic dispatch based on supplied array
    U, s, VH = ar.do("linalg.svd", x)

    # automatic dispatch based on different array
    sn = s + 0.1 * ar.do("random.normal", size=ar.shape(s), like=s)

    # automatic dispatch for multiple arrays for certain functions
    return ar.do("einsum", "ij,j,jk->ik", U, sn, VH)


# explicit backend given by string
x = ar.do("random.uniform", size=(100, 100), like="torch")

# this function now works for any backend
y = noised_svd(x)

# explicit inference of backend from array
ar.infer_backend(y)
# 'torch'

If you don't like the explicit do syntax, or simply want a drop-in replacement for existing code, you can also import the autoray.numpy module:

from autoray import numpy as np

# set a temporary default backend
with ar.backend_like("cupy"):
    z = np.ones((3, 4), dtype="float32")

np.exp(z)
# array([[2.7182817, 2.7182817, 2.7182817, 2.7182817],
#        [2.7182817, 2.7182817, 2.7182817, 2.7182817],
#        [2.7182817, 2.7182817, 2.7182817, 2.7182817]], dtype=float32)

Alternatively you can use autoray.get_namespace to get a backend specific (with optional default device and dtype) namespace object, (c.f. the Python Array Api):

xp = ar.get_namespace(z)
xp.einsum("ii->i", z)

Custom backends and functions can be dynamically registered with:

• register_backend
• register_function

The main documentation is available at autoray.readthedocs.io.