autoray

Submodules

Attributes

numpy

Classes

DoFunc

Get an automatic dispatch (i.e. backend selection deferred to call time)

Functions

astype(x, dtype_name, **kwargs)

Cast array as type dtype_name - tries x.astype first.

backend_like(like[, set_globally])

Context manager for setting a default backend. The argument like can

compose(fn, *[, name])

Take a function consisting of multiple autoray.do calls and compose

conj(x)

Array conjugate.

dag(x)

Array Hermitian transpose.

do(fn, *args[, like])

Do function named fn on (*args, **kwargs), peforming single

get_backend([get_globally])

Return the universally set backend, if any.

get_common_dtype(*arrays)

Compute the minimal dtype sufficient for arrays.

get_dtype_name(x)

Find string specifier dtype_name of array x.

get_lib_fn(backend, fn)

Cached retrieval of correct function for backend, all the logic for

get_namespace([like, device, dtype, submodule])

Get an automatic namespace object.

imag(x)

Array imaginary part.

infer_backend(array)

Get the name of the library that defined the class of array - unless

infer_backend_device_dtype(like[, device, dtype])

Infer the backend, device and dtype from like, with optional overrides

infer_backend_multi(*arrays)

Infer which backend should be used for a function that takes multiple

is_array(x)

Is x an array-like object? This simply checks for a shape

is_scalar(x)

Is x a scalar-like object? This checks if x has an ndim

ndim(x)

Get the number of dimensions of an array. This should be preferred to

real(x)

Array real part.

register_backend(cls, name)

Register the name (and by default the module or submodule) of a custom

register_function(backend, name[, fn, wrap, module, ...])

Customize how a single function name is dispatched for backend.

reshape(x, shape)

Array reshaped.

set_backend(like[, set_globally])

Set a default global backend. The argument like can be an explicit

shape(x)

Get the shape of an array as a tuple of int. This should be preferred

size(x)

Get the size, or number of elements, of an array. This should be

to_backend_dtype(dtype_name, like)

Turn string specifier dtype_name into dtype of backend like.

to_numpy(x)

Get a numpy version of array x.

transpose(x, *args)

Array transpose.

tree_apply(f, tree[, is_leaf])

Apply f to all leaves in tree, no new pytree is built.

tree_flatten(tree[, is_leaf, get_ref])

Flatten tree into a list of leaves.

tree_iter(tree[, is_leaf])

Iterate over all leaves in tree.

tree_map(f, tree[, is_leaf])

Map f over all leaves in tree, returning a new pytree.

tree_unflatten(objs, tree[, is_leaf])

Unflatten objs into a pytree of the same structure as tree.

autojit([fn, backend, compiler_opts])

Just-in-time compile an autoray function, automatically choosing

stop_gradient(x)

Stop gradient flow through array x.

Package Contents

class autoray.DoFunc(fn: str)[source]

Get an automatic dispatch (i.e. backend selection deferred to call time) callable for function named fn.

Slightly faster equivalent to functools.partial(do, fn).

Examples

DoFunc objects have a fixed operation but can still be called on any type of array:

>>> sqrt = DoFunc('sqrt')
>>> sqrt
<DoFunc sqrt>

>>> import numpy as np
>>> sqrt(np.random.uniform(size=[5]))
array([0.32464973, 0.90379787, 0.85037325, 0.88729814, 0.46768083])

>>> import cupy as cp
>>> sqrt(cp.random.uniform(size=[5]))
array([0.44541656, 0.88713113, 0.92626237, 0.64080557, 0.69620767])

>>> import tensorflow as tf
>>> sqrt(tf.random.uniform(shape=[5]))
<tf.Tensor: shape=(5,), dtype=float32, numpy=
array([0.3206495 , 0.8056399 , 0.5973012 , 0.13028008, 0.9820518 ],
    dtype=float32)>
__slots__ = ('fn',)
fn
__call__(*args, like=None, **kwargs)[source]
__repr__()[source]
autoray.astype(x, dtype_name, **kwargs)[source]

Cast array as type dtype_name - tries x.astype first.

autoray.backend_like(like, set_globally='auto')[source]

Context manager for setting a default backend. The argument like can be an explicit backend name or an array to infer it from.

Parameters:

  • like (str or array) – The backend to set. If an array, the backend of the array’s class will be set.

  • set_globally ({"auto", False, True}, optional) –

    Whether to set the backend globally or for the current thread:

    • True: set the backend globally.

    • False: set the backend for the current thread.

    • ”auto”: set the backend globally if this thread is the thread that imported autoray. Otherwise set the backend for the current thread.

    Only one thread should ever call this function with set_globally=True, (by default this is importing thread).

autoray.compose(fn, *, name=None)[source]

Take a function consisting of multiple autoray.do calls and compose it into a new, single, named function, registered with autoray.do.

This creates a default implementation of this function for each new backend encountered without explicitly having to write each out, but also allows for specific implementations to be overridden for specific backends.

If the function takes a backend argument, it will be supplied with the backend name, to save having to re-choose the backend.

Specific implementations can be provided by calling the register method of the composed function, or it can itself be used like a decorator:

@compose
def foo(x):
    ...

@foo.register("numpy")
@numba.njit
def foo_numba(x):
    ...
Parameters:

  • fn (callable) – The funtion to compose, and its default implementation.

  • name (str, optional) – The name of the composed function. If not provided, the name of the function will be used.

autoray.conj(x)[source]

Array conjugate.

autoray.dag(x)[source]

Array Hermitian transpose.

autoray.do(fn: str, *args, like=None, **kwargs)[source]

Do function named fn on (*args, **kwargs), peforming single dispatch to retrieve fn based on whichever library defines the class of the args[0], or the like keyword argument if specified.

Parameters:

  • fn (str) – Name of the function to do, e.g. ‘sum’ or ‘linalg.svd’.

  • args – Positional arguments to pass to the function.

  • like (str or array, optional) – Backend to use, either as an explicit backend name or an example array to infer the backend from. If not specified, the backend is inferred from the first argument, or from a globally set backend if any.

  • kwargs – Keyword arguments to pass to the function.

Examples

Works on numpy arrays:

>>> import numpy as np
>>> x_np = np.random.uniform(size=[5])
>>> y_np = do('sqrt', x_np)
>>> y_np
array([0.32464973, 0.90379787, 0.85037325, 0.88729814, 0.46768083])

>>> type(y_np)
numpy.ndarray

Works on cupy arrays:

>>> import cupy as cp
>>> x_cp = cp.random.uniform(size=[5])
>>> y_cp = do('sqrt', x_cp)
>>> y_cp
array([0.44541656, 0.88713113, 0.92626237, 0.64080557, 0.69620767])

>>> type(y_cp)
cupy.core.core.ndarray

Works on tensorflow arrays:

>>> import tensorflow as tf
>>> x_tf = tf.random.uniform(shape=[5])
>>> y_tf = do('sqrt', x_tf)
>>> y_tf
<tf.Tensor 'Sqrt_1:0' shape=(5,) dtype=float32>

>>> type(y_tf)
tensorflow.python.framework.ops.Tensor

You get the idea.

For functions that don’t dispatch on the first argument you can use the like keyword:

>>> do('eye', 3, like=x_tf)
<tf.Tensor: id=91, shape=(3, 3), dtype=float32>
autoray.get_backend(get_globally='auto')[source]

Return the universally set backend, if any.

Parameters:

get_globally ({"auto", False, True}, optional) –

Which backend to return:

  • True: return the globally set backend, if any.

  • False: return the backend set for the current thread, if any.

  • ”auto”: return the globally set backend, if this thread is the thread that imported autoray. Otherwise return the backend set for the current thread, if any.

Returns:

backend – The name of the backend, or None if no backend is set.

Return type:

str or None

autoray.get_common_dtype(*arrays)[source]

Compute the minimal dtype sufficient for arrays.

autoray.get_dtype_name(x)[source]

Find string specifier dtype_name of array x.

autoray.get_lib_fn(backend, fn)[source]

Cached retrieval of correct function for backend, all the logic for finding the correct funtion only runs the first time.

Parameters:

  • backend (str) – The module defining the array class to dispatch on.

  • fn (str) – The function to retrieve.

Return type:

callable

autoray.get_namespace(like=None, device=None, dtype=None, submodule=None)[source]

Get an automatic namespace object.

If like is None, the namespace essentially provides an alternative syntax to do, dispatching each function at calltime, and allowing the backend and function implementations to be dynamically updated.

If like is supplied however, the backend is eagerly dispatched and functions are loaded and cached specifically for that backend. In this case, default device and dtype can also be specified for various array creation routines, or if like is an array, inferred from that.

Parameters:

  • like (array-like, str or None, optional) – An array-like object to dispatch on, an explicit backend name, or None.

  • device (str or None, optional) – The device to use for array creation, or None to infer from like.

  • dtype (str or None, optional) – The data type to use for array creation, or None to infer from like.

Returns:

An automatic namespace object.

Return type:

AutoNamespace

autoray.imag(x)[source]

Array imaginary part.

autoray.infer_backend(array)[source]

Get the name of the library that defined the class of array - unless array is directly a subclass of numpy.ndarray, in which case assume numpy is the desired backend.

autoray.infer_backend_device_dtype(like, device=None, dtype=None)[source]

Infer the backend, device and dtype from like, with optional overrides for device and dtype. The dispatcher is cached on like.__class__ to avoid repeated lookups of the same attributes for the same type of array.

Parameters:

  • like (array-like or str or None) – The array to infer the backend, device and dtype from. If str, an explicit backend name. If None, the backend None is simply returned.

  • device (str or device_like, optional) – If given, an explicit device to use. If None, and like is an array with a device attribute, that is used.

  • dtype (str or dtype_like, optional) – If given, an explicit dtype to use. If None, and like is an array with a dtype attribute, that is used.

Returns:

  • backend (str or None) – The inferred backend name, or None if like is None.

  • device (str or device_like or None) – The inferred device, or None if not given and not found on like.

  • dtype (str or dtype_like or None) – The inferred dtype, or None if not given and not found on like.

autoray.infer_backend_multi(*arrays)[source]

Infer which backend should be used for a function that takes multiple arguments. This assigns a priority to each backend, and returns the backend with the highest priority. By default, the priority is:

  • builtins: -2

  • numpy: -1

  • other backends: 0

  • autoray.lazy: 1

I.e. when mixing with numpy, other array libraries are preferred, when mixing with autoray.lazy, autoray.lazy is preferred. This has quite low overhead due to caching.

autoray.is_array(x)[source]

Is x an array-like object? This simply checks for a shape attribute, thus 0-dimensional arrays are also considered arrays, but lists and tuples are not.

Parameters:

x (object) – Object to check.

Return type:

bool

See also

is_scalar

autoray.is_scalar(x)[source]

Is x a scalar-like object? This checks if x has an ndim attribute equal to 0. If x has no ndim attribute, it checks if x is iterable - if it is not iterable, it is considered a scalar.

Parameters:

x (object) – Object to check.

Return type:

bool

See also

is_array

autoray.ndim(x)[source]

Get the number of dimensions of an array. This should be preferred to calling x.ndim, since not all backends implement that, and it can also be called on nested lists and tuples.

Parameters:

x (array_like) – The array to get the number of dimensions of. It can be an arbitrary nested list or tuple of arrays and scalars.

Returns:

ndim

Return type:

int

autoray.numpy
autoray.real(x)[source]

Array real part.

autoray.register_backend(cls, name)[source]

Register the name (and by default the module or submodule) of a custom array class.

Parameters:

  • cls (type) – The array class itself.

  • name (str) – The name of the backend that should be used for this class. By default this wil be assumed to be the location of the relevant functions for this class, but this can be overridden.

autoray.register_function(backend, name, fn=None, *, wrap=False, module=None, alias=None, wrapper=None, inject_dtype=None, inject_device=None)[source]

Customize how a single function name is dispatched for backend.

This is the unified entry point for all function-level registration. It can set where the function lives (module), what it is called in the backend (alias), a lazy wrapper to apply on import, creation-routine dtype/device injection, and/or a direct implementation fn.

Parameters:

  • backend (str) – The name of the backend to register the function for.

  • name (str) – Name of the function, e.g. ‘sum’ or ‘linalg.svd’.

  • fn (callable, optional) – A direct implementation to use. If not supplied, and no other keyword argument is given, this function can be used as a decorator with backend and name only.

  • wrap (bool, optional) – Whether to wrap the old function like fn(old_fn) rather than directly supply the entire new function. This wrapper is eagerly called when registering, unlike wrapper.

  • module (str, optional) – Register the submodule location of the function, for when it is found somewhere other than the expected backend namespace, e.g. 'scipy.linalg'.

  • alias (str, optional) – Register a different name that the function is called in the backend, e.g. 'absolute' for 'abs'.

  • wrapper (callable, optional) – Register a custom wrapper, called lazily as wrapper(old_fn) the first time the function is imported, for when kwargs need translating or results modifying. Pass wrapper=True with fn=None to use this as a decorator that captures the wrapper.

  • inject_dtype (bool, optional) – Mark name as a creation routine that should have a dtype argument injected based on the like argument. Defaults to True when inject_device is given.

  • inject_device (bool, optional) – Mark name as a creation routine that should have a device argument injected based on the like argument.

Examples

Register a relocated, renamed and wrapped function in a single call:

register_function(
    "paddle", "random.normal",
    module="paddle", alias="randn", wrapper=scale_normal_manually,
)

Supply a direct implementation:

register_function("numpy", "complex", complex_add_re_im)

Use as a decorator for a direct implementation:

@register_function("torch", "to_numpy")
def torch_to_numpy(x):
    return x.detach().cpu().numpy()
autoray.reshape(x, shape)[source]

Array reshaped.

autoray.set_backend(like, set_globally='auto')[source]

Set a default global backend. The argument like can be an explicit backend name or an array.

Parameters:

  • like (str or array) – The backend to set. If an array, the backend of the array’s class will be set.

  • set_globally ({"auto", False, True}, optional) –

    Whether to set the backend globally or for the current thread:

    • True: set the backend globally.

    • False: set the backend for the current thread.

    • ”auto”: set the backend globally if this thread is the thread that imported autoray. Otherwise set the backend for the current thread.

    Only one thread should ever call this function with set_globally=True, (by default this is importing thread).

autoray.shape(x)[source]

Get the shape of an array as a tuple of int. This should be preferred to calling x.shape directly, as it:

  1. Allows customization (e.g. for torch and aesara which return different types for shape - use @shape.register(backend) to customize the behavior from this default implementation).

  2. Can be used on nested lists and tuples, without calling numpy.

Parameters:

x (array_like) – The array to get the shape of. It can be an arbitrary nested list or tuple of arrays and scalars, but is assumed not to be ragged.

Returns:

shape – The size of each dimension of the array.

Return type:

tuple of int

autoray.size(x)[source]

Get the size, or number of elements, of an array. This should be preferred to calling x.size, since not all backends implement that, and it can also be called on nested lists and tuples.

Parameters:

x (array_like) – The array to get the size of. It can be an arbitrary nested list or tuple of arrays and scalars.

Returns:

size

Return type:

int

autoray.to_backend_dtype(dtype_name, like)[source]

Turn string specifier dtype_name into dtype of backend like.

autoray.to_numpy(x)[source]

Get a numpy version of array x.

autoray.transpose(x, *args)[source]

Array transpose.

autoray.tree_apply(f, tree, is_leaf=is_not_container)[source]

Apply f to all leaves in tree, no new pytree is built.

Parameters:

  • f (callable) – A function to apply to all leaves in tree.

  • tree (pytree) – A nested sequence of tuples, lists, dicts and other objects.

  • is_leaf (callable) – A function to determine if an object is a leaf, f is only applied to objects for which is_leaf(x) returns True.

autoray.tree_flatten(tree, is_leaf=is_not_container, get_ref=False)[source]

Flatten tree into a list of leaves.

Parameters:

  • tree (pytree) – A nested sequence of tuples, lists, dicts and other objects.

  • is_leaf (callable) – A function to determine if an object is a leaf, only objects for which is_leaf(x) returns True are returned in the flattened list.

  • get_ref (bool) – If True, a reference tree is also returned which can be used to reconstruct the original tree from a flattened list.

Returns:

  • objs (list) – The flattened list of leaf objects.

  • (ref_tree) (pytree) – If get_ref is True, a reference tree, with leaves of Leaf, is returned which can be used to reconstruct the original tree.

autoray.tree_iter(tree, is_leaf=is_not_container)[source]

Iterate over all leaves in tree.

Parameters:

  • f (callable) – A function to apply to all leaves in tree.

  • tree (pytree) – A nested sequence of tuples, lists, dicts and other objects.

  • is_leaf (callable) – A function to determine if an object is a leaf, f is only applied to objects for which is_leaf(x) returns True.

autoray.tree_map(f, tree, is_leaf=is_not_container)[source]

Map f over all leaves in tree, returning a new pytree.

Parameters:

  • f (callable) – A function to apply to all leaves in tree.

  • tree (pytree) – A nested sequence of tuples, lists, dicts and other objects.

  • is_leaf (callable) – A function to determine if an object is a leaf, f is only applied to objects for which is_leaf(x) returns True.

Return type:

pytree

autoray.tree_unflatten(objs, tree, is_leaf=is_leaf_placeholder)[source]

Unflatten objs into a pytree of the same structure as tree.

Parameters:

  • objs (sequence) – A sequence of objects to be unflattened into a pytree.

  • tree (pytree) – A nested sequence of tuples, lists, dicts and other objects, the objs will be inserted into a new pytree of the same structure.

  • is_leaf (callable) – A function to determine if an object is a leaf, only objects for which is_leaf(x) returns True will have the next item from objs inserted. By default checks for the Leaf object inserted by tree_flatten(..., get_ref=True).

Return type:

pytree

autoray.autojit(fn=None, *, backend=None, compiler_opts=None)[source]

Just-in-time compile an autoray function, automatically choosing the backend based on the input arrays, or via keyword argument.

The backend used to do the compilation can be set in three ways:

  1. Automatically based on the arrays the function is called with, i.e. cfn(*torch_arrays) will use torch.jit.trace.

  2. In this wrapper, @autojit(backend='jax'), to provide a specific default instead.

  3. When you call the function cfn(*arrays, backend='torch') to override on a per-call basis.

If the arrays supplied are of a different backend type to the compiler, then the returned array will also be converted back, i.e. cfn(*numpy_arrays, backend='tensorflow') will return a numpy array.

The 'python' backend simply extracts and unravels all the do calls into a code object using compile which is then run with exec. This makes use of shared intermediates and constant folding, strips away any python scaffoliding, and is compatible with any library, but the resulting function is not ‘low-level’ in the same way as the other backends.

Parameters:

  • fn (callable) – The autoray function to compile.

  • backend (str, optional) –

    If set, use this as the default backend. The options are:

    • 'python': extract and unravel all the do calls into a code object using compile which is then run with exec.

    • 'jax': use jax.jit to compile the function.

    • 'tensorflow': use tf.function to compile the function.

    • 'torch': use torch.jit.trace to compile the function.

    • 'pytensor': use pytensor.function to compile the function.

    If not set, the backend will be inferred from the input arrays, or it can be set at call time with the backend keyword argument.

  • compiler_opts (dict[dict], optional) – Dict of dicts when you can supply options for each compiler backend separately, e.g.: @autojit(compiler_opts={'tensorflow': {'jit_compile': True}}).

Returns:

cfn – The function with auto compilation.

Return type:

callable

autoray.stop_gradient(x)[source]

Stop gradient flow through array x.

In autodiff backends (JAX, PyTorch, TensorFlow, etc.), this detaches x from the computational graph so that no gradients are propagated through it. For non-autodiff backends (NumPy, CuPy, etc.), this is a no-op and returns x unchanged.

Parameters:

x (array) – The array to stop gradient flow through.

Returns:

An array with the same value as x but detached from the autodiff computational graph.

Return type:

array

Examples

With JAX:

>>> import jax, jax.numpy as jnp, autoray as ar
>>> x = jnp.array([1.0, 2.0, 3.0])
>>> ar.stop_gradient(x)  # equivalent to jax.lax.stop_gradient(x)

With PyTorch:

>>> import torch, autoray as ar
>>> x = torch.tensor([1.0, 2.0], requires_grad=True)
>>> y = ar.stop_gradient(x)  # equivalent to x.detach()
>>> y.requires_grad
False

With NumPy (no-op):

>>> import numpy as np, autoray as ar
>>> x = np.array([1.0, 2.0])
>>> ar.stop_gradient(x) is x
True