autoray on Pypi

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
dask
autograd
tensorflow
sparse
mars
... 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:

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