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

Custom backends and functions can be dynamically registered with:

• register_backend
• register_function

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