autoray.autoray

AUTORAY - backend agnostic array operations.

Copyright 2019-2026 Johnnie Gray

Licensed under the Apache License, Version 2.0 (the “License”); you may not use this file except in compliance with the License. You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an “AS IS” BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.

Attributes

_global_backend
_inferrer_global
_infer_auto
_global_backends_threadaware
_inferrers_threadaware
_importing_thrid
_backend_lock
_CUSTOM_BACKENDS
multi_class_priorities
_backend_device_dtype_dispatchers
_CREATION_ROUTINES
_CREATION_INJECT
_BACKEND_ALIASES
_MODULE_ALIASES
_SUBMODULE_ALIASES
_FUNC_ALIASES
_CUSTOM_WRAPPERS
_FUNCS
_COMPOSED_FUNCTION_GENERATORS
_IS_SCALAR_CACHE
TREE_MAP_REGISTRY
TREE_APPLY_REGISTRY
TREE_ITER_REGISTRY
IS_CONTAINER_CACHE
TREE_MAPPER_CACHE
TREE_ITER_CACHE
TREE_APPLIER_CACHE
LEAF
_BUILTIN_DTYPE_NAMES
_COMPLEX_DTYPES
_DOUBLE_DTYPES
_DTYPE_MAP
_DTYPE_MATCHER
_DEVICE_MATCHER
_DISPATCHERS
_RANDOM_DISTS
_COMPLEX_TO_REAL_DTYPE
_RADEMACHER_SIGNS
_RADEMACHER_ROOTS
_RADEMACHER_SAMPLERS
_RADEMACHER_SHARED
_NAME_SPACE_SUBMODULES
_NAMESPACE_ATTRS
_NAMESPACE_CACHE
numpy
_builtin_dtype_lookup
_JAX_DEFAULT_RNG
_torch_reduce_translation
_paddle_dtype_name_conversion

Classes

DoFunc	Get an automatic dispatch (i.e. backend selection deferred to call time)
ConstantInferrer	A simple dispatch inferrer that always returns the same backend. Used
Leaf	A singleton object to use as a placeholder in a pytree, for
Composed	Compose an autoray.do using function. See the main wrapper
InjectDtypeDevice	Wrapper that possibly injects default dtype and device arguments, if not
AutoNamespace	Mimics a namespace, optionally for a specific backend, device, and
JaxDefaultRNG	Stateful but deterministic random number generator for JAX following
TensorflowDefaultRNG	Stateful random number generator for TensorFlow following numpy's
TorchDefaultRNG
MlxDefaultRNG	Stateful but deterministic random number generator for MLX following
MlxSharedRNG	Draws from mlx's own shared random state, which mx.random.seed

Functions

do(fn, *args[, like])	Do function named fn on (*args, **kwargs), peforming single
_default_infer_from_sig(fn, *args, **kwargs)	This is the default backend dispatcher, used if no global backend has
_default_infer_from_sig_threadaware(fn, args, kwargs)
get_backend([get_globally])	Return the universally set backend, if any.
set_backend(like[, set_globally])	Set a default global backend. The argument like can be an explicit
backend_like(like[, set_globally])	Context manager for setting a default backend. The argument like can
register_backend(cls, name)	Register the name (and by default the module or submodule) of a custom
_infer_class_backend_cached(cls)
infer_backend(array)	Get the name of the library that defined the class of array - unless
_infer_class_backend_multi_cached(classes)
infer_backend_multi(*arrays)	Infer which backend should be used for a function that takes multiple
_invalidate_backend_inference_caches()	Remove all cache entries that depend on the backend of a class.
_make_device_dtype_dispatch(like)	Make a dispatcher function that possibly looks up default device and
infer_backend_device_dtype(like[, device, dtype])	Infer the backend, device and dtype from like, with optional overrides
_register_creation_inject(backend, fn, inject_dtype, ...)	Record whether dtype and/or device should be injected into the
register_creation_routine(backend, fn[, inject_dtype, ...])	Register a function that creates a new array, with dtype and possibly
_choose_backend(fn, args, kwargs[, like])	Private function to choose a backend based on function name and
choose_backend(fn, *args[, like])	Choose a backend based on function name, arguments, and the like
import_lib_fn(backend, fn)
get_lib_fn(backend, fn)	Cached retrieval of correct function for backend, all the logic for
register_backend_alias(alias, backend)	Register an alias for a backend, i.e. if the backend alias is
register_module_alias(alias, module)	Register an alias for a module.
register_submodule_alias(backend, fn, module)	Register an alias for a submodule location of a function.
register_func_alias(backend, fn, alias)	Register an alias for a function name.
register_custom_wrapper(backend, fn[, wrapper])	Register a custom wrapper for a function. The wrapper is called lazily
register_function(backend, name[, fn, wrap, module, ...])	Customize how a single function name is dispatched for backend.
is_array(x)	Is x an array-like object? This simply checks for a shape
_is_scalar_ndim(x)
_always_false(x)
_always_true(x)
is_scalar(x)	Is x a scalar-like object? This checks if x has an ndim
tree_register_container(cls, mapper, iterator, applier)	Register a new container type for use with tree_map and
is_not_container(x)	The default function to determine if an object is a leaf. This simply
identity(f, tree, is_leaf)
tree_map(f, tree[, is_leaf])	Map f over all leaves in tree, returning a new pytree.
empty(tree, is_leaf)
tree_iter(tree[, is_leaf])	Iterate over all leaves in tree.
nothing(f, tree, is_leaf)
tree_apply(f, tree[, is_leaf])	Apply f to all leaves in tree, no new pytree is built.
is_leaf_placeholder(x)
tree_flatten(tree[, is_leaf, get_ref])	Flatten tree into a list of leaves.
tree_unflatten(objs, tree[, is_leaf])	Unflatten objs into a pytree of the same structure as tree.
tree_map_tuple(f, tree, is_leaf)
tree_iter_tuple(tree, is_leaf)
tree_apply_tuple(f, tree, is_leaf)
tree_map_list(f, tree, is_leaf)
tree_iter_list(tree, is_leaf)
tree_apply_list(f, tree, is_leaf)
tree_map_dict(f, tree, is_leaf)
tree_iter_dict(tree, is_leaf)
tree_apply_dict(f, tree, is_leaf)
_choose_namespace(backend, args)	Choose the namespace to supply to a composed function, given the
compose([fn, name])	Take a function consisting of multiple autoray.do calls and compose
shape(x)	Get the shape of an array as a tuple of int. This should be preferred
ndim(x)	Get the number of dimensions of an array. This should be preferred to
size(x)	Get the size, or number of elements, of an array. This should be
conj(x)	Array conjugate.
transpose(x, *args)	Array transpose.
dag(x)	Array Hermitian transpose.
real(x)	Array real part.
imag(x)	Array imaginary part.
reshape(x, shape)	Array reshaped.
_to_backend_dtype_from_str_cached(dtype_name, like)
to_backend_dtype(dtype_name, like)	Turn string specifier dtype_name into dtype of backend like.
_dtype_to_name_cached(dtype)
get_dtype_name(x)	Find string specifier dtype_name of array x.
get_common_dtype(*arrays)	Compute the minimal dtype sufficient for arrays.
astype(x, dtype_name, **kwargs)	Cast array as type dtype_name - tries x.astype first.
to_numpy(x)	Get a numpy version of array x, via np.asarray by default.
_parse_compound_backend_spec(spec)	Parse a composite string specifier like "torch-float32-cuda:0"
_dtype_is_inexact(dtype_name)	Whether string dtype_name is a floating point or complex dtype.
to_device(x, device)	Move array x to device, returning it unchanged if device
from_numpy(x[, dtype, device, backend])	Convert a numpy array (or array-like) x into a like backend
to(tree[, like, backend, dtype, device])	Convert an array, or nested collection ("pytree") of arrays, to a
svd_not_full_matrices_wrapper(fn)
svd_sUV_to_UsVH_wrapper(fn)
svd_UsV_to_UsVH_wrapper(fn)
svd_manual_full_matrices_kwarg(fn)
qr_allow_fat(fn)
tril_to_band_part(fn)
triu_to_band_part(fn)
cholesky_lower(fn)	Make a cholesky wrapper that translates upper to lower bool.
cholesky_manual_upper(fn)	Make a cholesky wrapper adding upper for backends that only compute
binary_allow_1d_rhs_wrap(fn)
scale_random_uniform_manually(fn)
scale_random_normal_manually(fn)
with_dtype_wrapper(fn)	Add ability to handle dtype keyword.
translate_wrapper(fn, translator)	Wrap a function to match the api of another according to a translation.
make_translator(t)
complex_add_re_im(re, im)
allclose(x, y[, rtol, atol])
_handle_size_to_shape([size])
wrap_args_kwargs_from_raw(fn)	Take a function with signature (*args, **kwargs) and wrap it to
register_dispatch(fun, dispatcher[, raw_signature])	Register a new dispatcher, a function that takes the arguments and
default_dispatcher(args, kwargs)	Try to infer backend from first argument passed to function.
join_array_dispatcher(args, kwargs)	Dispatcher for functions where first argument is a sequence.
einsum_dispatcher(args, kwargs)	Dispatcher for handling einsum.
binary_dispatcher(args, kwargs)	There are cases when we want to take into account both backends of two
random_array_dispatcher(shape[, rng])	Use the generator's backend when given, or infer it from shape.
_get_rademacher_table(dtype_name, backend)	The lookup table to index for a rademacher sample: the two signs, or
_sample_rademacher(rng, shape, dtype_name, device, backend)	Draw from {-1, +1}, or from the four roots of unity for a complex
_sample_random_array(rng, dist, shape, dtype_name, ...)
random_array(shape[, dist, loc, scale, dtype, device, ...])	Generate an array of random samples.
_namespace_key_part(x)	Cached str of a device or dtype, which normalizes them for the
_reset_namespace(xp)	Drop the cached lookups of xp and of any submodule it made.
_reset_namespaces()	Drop the cached function and submodule lookups of every live namespace,
get_namespace([like, device, dtype, submodule])	Get an automatic namespace object.
builtins_get_dtype_name(x)
_rademacher_numpy(rng, shape, dtype_name, device)
cupy_to_numpy(x)
_cupy_parse_device(device)	Check a device string is valid for cupy, returning the gpu index, or
cupy_to_device(x, device)
cupy_from_numpy(x[, dtype, device])
_rademacher_cupy(rng, shape, dtype_name, device)
get_jax()
_jax_parse_device(device)	Parse a device string like "cuda:0" into (platform, index).
jax_to_device(x, device)
jax_from_numpy(x[, dtype, device])
_warn_jax_generated_seed()	Warn one time only, since the message applies to every seedless call.
jax_default_rng([seed])
jax_random_seed([seed])
_get_jax_default_rng()
jax_random_uniform([low, high, size])
jax_random_normal([loc, scale, size])
_rademacher_jax(rng, shape, dtype_name, device)
aesara_shape(x)
autograd_take(x, indices[, axis])
dask_to_numpy(x)
dask_eye_wrapper(eye_fn)
ctf_array(x)
ctf_to_numpy(x)
ctf_count_nonzero(x)
ctf_get_dtype_name(x)
sparse_array(x)
sparse_to_numpy(x)
sparse_transpose(x[, axes])
sparse_reshape(x, shape)
sparse_sum(x[, axis, keepdims, dtype, out])
sparse_prod(x[, axis, keepdims, dtype, out])
sparse_conj(x)
sparse_real(x)
sparse_imag(x)
sparse_count_nonzero(x)
sparse_complex(re, im)
sparse_random_uniform([low, high, size, dtype])
sparse_random_normal([loc, scale, size, dtype])
get_tensorflow()
tensorflow_pad_wrap(tf_pad)
tensorflow_wrap_norm(tf_norm)
tensorflow_default_rng([seed])
tensorflow_to_numpy(x)
_tensorflow_translate_device(device)	Translate a device string like "cuda:0" to tensorflow form.
tensorflow_to_device(x, device)
tensorflow_from_numpy(x[, dtype, device])
tensorflow_indices(dimensions)
tensorflow_swapaxes(a, axis1, axis2)
tensorflow_cholesky(x[, upper])
tensorflow_solve_triangular(a, b[, lower])
tensorflow_astype(x, dtype)
get_torch()
torch_shape(x)
torch_size(x)
torch_linalg_solve_wrap(fn)
torch_tensordot_wrap(fn)
torch_split_wrap(fn)
torch_maybe_convert_dtype_from_str(fn)
torch_zeros_ones_wrap(fn)
torch_eye_wrap(fn)
torch_sort_wrap(fn)
torch_flip_wrap(torch_flip)
torch_nonzero_wrap(torch_nonzero)
torch_default_rng([seed])
torch_random_seed([seed])
_rademacher_torch(rng, shape, dtype_name, device)
torch_random_normal([loc, scale, size, dtype])
torch_random_uniform([low, high, size, dtype])
torch_to_numpy(x)
torch_to_device(x, device)
torch_from_numpy(x[, dtype, device])
torch_copy(x)
torch_transpose(x[, axes])
torch_astype(x, dtype)
torch_complex(re, im)
torch_real(x)
torch_imag(x)
torch_linalg_eigh(x)
torch_linalg_eigvalsh(x)
torch_scipy_linalg_solve_triangular(a, b[, lower, ...])
torch_pad(array, pad_width[, mode, constant_values])
torch_indices(dimensions)
torch_take(a, indices[, axis])
torch_trace(x[, axis1, axis2])
mxnet_to_numpy(x)
paddle_get_dtype_name(x)
paddle_shape(x)
paddle_wrap_reduction(fn)
paddle_split_wrap(fn)
paddle_imag(x)
paddle_indices(dimensions)
paddle_pad(array, pad_width[, mode, constant_values])
paddle_ravel(x)
paddle_real(x)
paddle_to_numpy(x)
paddle_transpose(a[, axes])
pytensor_shape(x)
pytensor_wrap_qr_with_shapes(fn)
pytensor_wrap_svd_with_shapes(fn)
get_mlx()
mlx_to_device(x, device)
mlx_default_rng([seed])
_get_mlx_shared_rng()
mlx_random_seed([seed])
mlx_random_uniform([low, high, size])
mlx_random_normal([loc, scale, size])
mlx_count_nonzero(x[, axis, keepdims])
mlx_ravel(x, *args, **kwargs)
mlx_zeros_ones_wrap(fn)
mlx_array_asarray_wrap(fn)
mlx_eye_wrap(fn)

Module Contents

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

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>

class autoray.autoray.DoFunc(fn: str)

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)

__repr__()

autoray.autoray._default_infer_from_sig(fn, *args, **kwargs)

This is the default backend dispatcher, used if no global backend has been set. Hot swapping this function out as below avoids having to check manually for a global backend or worse, a thread aware global backend, on every call to do.

autoray.autoray._global_backend = None

autoray.autoray._inferrer_global

autoray.autoray._infer_auto

autoray.autoray._global_backends_threadaware

autoray.autoray._inferrers_threadaware

autoray.autoray._importing_thrid

autoray.autoray._backend_lock

autoray.autoray._default_infer_from_sig_threadaware(fn, args, kwargs)

class autoray.autoray.ConstantInferrer(backend)

A simple dispatch inferrer that always returns the same backend. Used with set_backend to set a uniform constant backend for all calls.

__slots__ = ('backend',)

backend

__call__(fn, args, kwargs)

autoray.autoray.get_backend(get_globally='auto')

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.autoray.set_backend(like, set_globally='auto')

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.autoray.backend_like(like, set_globally='auto')

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.autoray._CUSTOM_BACKENDS

autoray.autoray.register_backend(cls, name)

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.autoray._infer_class_backend_cached(cls)

autoray.autoray.infer_backend(array)

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.autoray.multi_class_priorities

autoray.autoray._infer_class_backend_multi_cached(classes)

autoray.autoray.infer_backend_multi(*arrays)

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.autoray._backend_device_dtype_dispatchers

autoray.autoray._invalidate_backend_inference_caches()

Remove all cache entries that depend on the backend of a class.

autoray.autoray._make_device_dtype_dispatch(like)

Make a dispatcher function that possibly looks up default device and dtype if those are not given. Whether the dispatcher should look up those attributes is cached on like.__class__.

autoray.autoray.infer_backend_device_dtype(like, device=None, dtype=None)

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.autoray._CREATION_ROUTINES

autoray.autoray._CREATION_INJECT

autoray.autoray._register_creation_inject(backend, fn, inject_dtype, inject_device)

Record whether dtype and/or device should be injected into the call to creation routine fn for backend, based on the like argument. See register_function.

autoray.autoray.register_creation_routine(backend, fn, inject_dtype=True, inject_device=False)

Register a function that creates a new array, with dtype and possibly device kwargs, that should be inferred from the like argument. This is not necessary for array creation routines that don’t accept either.

Deprecated since version 0.8.12: Prefer register_function(backend, fn, inject_dtype=..., inject_device=...), which can register the location, name, wrapper and creation-injection behaviour of a function in a single call.

Parameters:

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

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

• inject_dtype (bool, optional) – Whether to inject a dtype argument based on the like argument.

• inject_device (bool, optional) – Whether to inject a device argument based on the like argument.

autoray.autoray._choose_backend(fn, args, kwargs, like=None)

Private function to choose a backend based on function name and signature, which passes args and kwargs by reference for performance and also to allow injection of dtype and device arguments for array creation routines.

autoray.autoray.choose_backend(fn, *args, like=None, **kwargs)

Choose a backend based on function name, arguments, and the like keyword argument. The default, if like is not specified, is to infer the backend from the function call, the default of which is simply to use the first argument, if no custom dispatcher is found. Otherwise the backend is chosen based on the like argument - which can be an explicit backend name or an arbitrary object.

autoray.autoray._BACKEND_ALIASES

autoray.autoray._MODULE_ALIASES

autoray.autoray._SUBMODULE_ALIASES

autoray.autoray._FUNC_ALIASES

autoray.autoray._CUSTOM_WRAPPERS

autoray.autoray._FUNCS

autoray.autoray._COMPOSED_FUNCTION_GENERATORS

autoray.autoray.import_lib_fn(backend, fn)

autoray.autoray.get_lib_fn(backend, fn)

Cached retrieval of correct function for backend, all the logic for finding the correct function 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.autoray.register_backend_alias(alias, backend)

Register an alias for a backend, i.e. if the backend alias is inferred for an array, use functions from backend instead.

Parameters:

• alias (str) – The alias name.

• backend (str) – The actual backend name.

autoray.autoray.register_module_alias(alias, module)

Register an alias for a module.

Parameters:

• alias (str) – The alias name.

• module (str) – The actual module name.

autoray.autoray.register_submodule_alias(backend, fn, module)

Register an alias for a submodule location of a function.

Deprecated since version 0.8.12: Prefer register_function(backend, fn, module=module).

Parameters:

• backend (str) – The name of the backend.

• fn (str) – The name of the function.

• module (str) – The module where the function is located.

autoray.autoray.register_func_alias(backend, fn, alias)

Register an alias for a function name.

Deprecated since version 0.8.12: Prefer register_function(backend, fn, alias=alias).

Parameters:

• backend (str) – The name of the backend.

• fn (str) – The name of the function.

• alias (str) – The name of the function in the backend.

autoray.autoray.register_custom_wrapper(backend, fn, wrapper=None)

Register a custom wrapper for a function. The wrapper is called lazily so that no imports are done until the function is actually used.

Deprecated since version 0.8.12: Prefer register_function(backend, fn, wrapper=wrapper), or as a decorator register_function(backend, fn, wrapper=True).

Parameters:

• backend (str) – The name of the backend.

• fn (str) – The name of the function.

• wrapper (callable, optional) – The wrapper function. It should take a function as input and return a function. If not supplied, this function can be used as a decorator with backend and fn only.

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

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.autoray.is_array(x)

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.autoray._IS_SCALAR_CACHE

autoray.autoray._is_scalar_ndim(x)

autoray.autoray._always_false(x)

autoray.autoray._always_true(x)

autoray.autoray.is_scalar(x)

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.autoray.TREE_MAP_REGISTRY

autoray.autoray.TREE_APPLY_REGISTRY

autoray.autoray.TREE_ITER_REGISTRY

autoray.autoray.tree_register_container(cls, mapper, iterator, applier)

Register a new container type for use with tree_map and tree_apply.

Parameters:

• cls (type) – The container type to register.

• mapper (callable) – A function that takes f, tree and is_leaf and returns a new tree of type cls with f applied to all leaves.

• applier (callable) – A function that takes f, tree and is_leaf and applies f to all leaves in tree.

autoray.autoray.IS_CONTAINER_CACHE

autoray.autoray.is_not_container(x)

The default function to determine if an object is a leaf. This simply checks if the object is an instance of any of the registered container types.

autoray.autoray.identity(f, tree, is_leaf)

autoray.autoray.TREE_MAPPER_CACHE

autoray.autoray.tree_map(f, tree, is_leaf=is_not_container)

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.autoray.empty(tree, is_leaf)

autoray.autoray.TREE_ITER_CACHE

autoray.autoray.tree_iter(tree, is_leaf=is_not_container)

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.autoray.nothing(f, tree, is_leaf)

autoray.autoray.TREE_APPLIER_CACHE

autoray.autoray.tree_apply(f, tree, is_leaf=is_not_container)

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.

class autoray.autoray.Leaf

A singleton object to use as a placeholder in a pytree, for unflattening.

__slots__ = ()

__repr__()

autoray.autoray.LEAF

autoray.autoray.is_leaf_placeholder(x)

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

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.autoray.tree_unflatten(objs, tree, is_leaf=is_leaf_placeholder)

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.autoray.tree_map_tuple(f, tree, is_leaf)

autoray.autoray.tree_iter_tuple(tree, is_leaf)

autoray.autoray.tree_apply_tuple(f, tree, is_leaf)

autoray.autoray.tree_map_list(f, tree, is_leaf)

autoray.autoray.tree_iter_list(tree, is_leaf)

autoray.autoray.tree_apply_list(f, tree, is_leaf)

autoray.autoray.tree_map_dict(f, tree, is_leaf)

autoray.autoray.tree_iter_dict(tree, is_leaf)

autoray.autoray.tree_apply_dict(f, tree, is_leaf)

autoray.autoray._choose_namespace(backend, args)

Choose the namespace to supply to a composed function, given the already chosen backend and the positional args of the call. The first argument supplies the dtype and device defaults if it belongs to backend, otherwise the namespace has none.

class autoray.autoray.Composed(fn, name=None)

Compose an autoray.do using function. See the main wrapper compose.

_default_fn

_name = None

_supply_backend

_supply_namespace

register(backend, fn=None)

Register a different implementation for backend.

_make_default_function(backend, namespace=None)

make_function(backend)

Make a new function for the specific backend.

__call__(*args, like=None, **kwargs)

__repr__()

autoray.autoray.compose(fn=None, *, name=None)

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. If it takes a namespace argument, it will similarly be supplied with an AutoNamespace. Calling through a namespace supplies that namespace, otherwise it is taken from the first argument if that matches the backend, and has no dtype or device defaults if not.

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):
    ...

Supply name to register the function under a name other than its own, which requires calling compose first:

@compose(name="linalg.qr")
def qr(x):
    ...

Parameters:

• fn (callable, optional) – The function to compose, and its default implementation. Omitting it returns a decorator that takes it, which is how the second form above supplies name.

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

autoray.autoray.shape(x)

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.autoray.ndim(x)

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.autoray.size(x)

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.autoray.conj(x)

Array conjugate.

autoray.autoray.transpose(x, *args)

Array transpose.

autoray.autoray.dag(x)

Array Hermitian transpose.

autoray.autoray.real(x)

Array real part.

autoray.autoray.imag(x)

Array imaginary part.

autoray.autoray.reshape(x, shape)

Array reshaped.

autoray.autoray._to_backend_dtype_from_str_cached(dtype_name, like)

autoray.autoray.to_backend_dtype(dtype_name, like)

Turn string specifier dtype_name into dtype of backend like.

autoray.autoray._BUILTIN_DTYPE_NAMES

autoray.autoray._dtype_to_name_cached(dtype)

autoray.autoray.get_dtype_name(x)

Find string specifier dtype_name of array x.

autoray.autoray._COMPLEX_DTYPES

autoray.autoray._DOUBLE_DTYPES

autoray.autoray._DTYPE_MAP

autoray.autoray.get_common_dtype(*arrays)

Compute the minimal dtype sufficient for arrays.

autoray.autoray.astype(x, dtype_name, **kwargs)

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

autoray.autoray.to_numpy(x)

Get a numpy version of array x, via np.asarray by default.

autoray.autoray._DTYPE_MATCHER

autoray.autoray._DEVICE_MATCHER

autoray.autoray._parse_compound_backend_spec(spec)

Parse a composite string specifier like "torch-float32-cuda:0" into a (backend, dtype, device) tuple, each a string or None. Token order is not important: dtype and device tokens are recognized by pattern, and a single remaining token, if any, is taken as the backend.

autoray.autoray._dtype_is_inexact(dtype_name)

Whether string dtype_name is a floating point or complex dtype.

autoray.autoray.to_device(x, device)

Move array x to device, returning it unchanged if device is None. A bare device type without an index, e.g. "gpu" or "cuda", means ‘ensure on this type of device’: arrays already on such a device are not migrated between indices. The default implementation tries x.to(device), treating backends without any device concept as ‘cpu’.

Parameters:

• x (array) – The array to move.

• device (str or device-like or None) – The device to move to, e.g. "cuda:0".

Return type:

array

autoray.autoray.from_numpy(x, dtype=None, device=None, backend=None)

Convert a numpy array (or array-like) x into a like backend array, directly with the given dtype and on the given device where possible. It is registered as a creation routine, so if like is an example array, unspecified dtype and device default to matching it. The default implementation is asarray then to_device, but backends can register more direct routes, e.g. a single torch.as_tensor call.

Parameters:

• x (array-like) – The numpy array (or nested iterable) to convert.

• dtype (str or dtype, optional) – The target dtype.

• device (str or device-like, optional) – The target device, e.g. "cuda:0".

• like (str or array, optional) – The target backend, as an explicit name, or an example array to also infer default dtype and device from. Handled by the dispatch layer.

Return type:

array

See also

to, to_numpy, to_device

autoray.autoray.to(tree, like=None, *, backend=None, dtype=None, device=None)

Convert an array, or nested collection (“pytree”) of arrays, to a target backend, dtype and/or device. All three can be specified together in a single string such as "torch-float32-cuda:0", in any order, or explicitly via the keyword arguments, which take precedence. Unspecified properties are left unchanged, and non-array leaves are passed through untouched. Repeated references to the same input array are converted once and share the same output array. Note that, matching torch.nn.Module.to semantics, only floating point and complex arrays are cast when a dtype is given, so that e.g. integer index arrays are preserved.

Parameters:

• tree (array or pytree of arrays) – The array or nested collection (tuple, list, dict, or any registered container) of arrays to convert.

• like (str or array, optional) – The conversion target. If a string, a dash separated specifier like "backend-dtype-device", with each part optional. If an array, the backend, dtype and device to target are inferred from it.

• backend (str, optional) – Explicit target backend, taking precedence over like.

• dtype (str or dtype, optional) – Explicit target dtype, taking precedence over like. Only applied to floating point and complex arrays.

• device (str or device-like, optional) – Explicit target device, taking precedence over like.

Returns:

The converted array or collection, matching the structure of tree.

Return type:

array or pytree of arrays

Examples

>>> import numpy as np
>>> xs = {"a": np.random.rand(2, 3), "b": np.arange(3)}
>>> ys = to(xs, "torch-float32")
>>> ys["a"].dtype
torch.float32
>>> ys["b"].dtype  # integer arrays are not cast
torch.int64

See also

to_device, astype, to_numpy, tree_map

autoray.autoray.svd_not_full_matrices_wrapper(fn)

autoray.autoray.svd_sUV_to_UsVH_wrapper(fn)

autoray.autoray.svd_UsV_to_UsVH_wrapper(fn)

autoray.autoray.svd_manual_full_matrices_kwarg(fn)

autoray.autoray.qr_allow_fat(fn)

autoray.autoray.tril_to_band_part(fn)

autoray.autoray.triu_to_band_part(fn)

autoray.autoray.cholesky_lower(fn)

Make a cholesky wrapper that translates upper to lower bool.

autoray.autoray.cholesky_manual_upper(fn)

Make a cholesky wrapper adding upper for backends that only compute the lower factor.

autoray.autoray.binary_allow_1d_rhs_wrap(fn)

autoray.autoray.scale_random_uniform_manually(fn)

autoray.autoray.scale_random_normal_manually(fn)

autoray.autoray.with_dtype_wrapper(fn)

Add ability to handle dtype keyword. If not None, dtype should be specified as a string, otherwise conversion will happen regardless.

autoray.autoray.translate_wrapper(fn, translator)

Wrap a function to match the api of another according to a translation. The translator entries in the form of an ordered dict should have entries like:

(desired_kwarg: (backend_kwarg, default_value))

with the order defining the args of the function.

autoray.autoray.make_translator(t)

autoray.autoray.complex_add_re_im(re, im)

autoray.autoray.allclose(x, y, rtol=1e-05, atol=1e-08)

autoray.autoray._handle_size_to_shape(size=None)

autoray.autoray.wrap_args_kwargs_from_raw(fn)

Take a function with signature (*args, **kwargs) and wrap it to accept a single tuple of args and a dict of kwargs.

autoray.autoray.register_dispatch(fun, dispatcher, raw_signature=True)

Register a new dispatcher, a function that takes the arguments and keyword arguments of a function and returns the backend to use, when the backend is not explicitly given.

This is useful in case the backend to be used by a function cannot be inferred from the first argument.

Parameters:

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

• dispatcher (callable) – The dispatcher function to use. This should take the arguments and keyword arguments of the function and return the backend to use.

• raw_signature (bool, optional) – The dispatcher has signature (*args, **kwargs) if True, otherwise it has signature (args, kwargs).

autoray.autoray.default_dispatcher(args, kwargs)

Try to infer backend from first argument passed to function.

autoray.autoray._DISPATCHERS

autoray.autoray.join_array_dispatcher(args, kwargs)

Dispatcher for functions where first argument is a sequence.

autoray.autoray.einsum_dispatcher(args, kwargs)

Dispatcher for handling einsum.

einsum can be called with a str equation as the first argument, or with ‘interleaved’ inputs. This dispatcher handles both cases and also takes into account all arrays.

autoray.autoray.binary_dispatcher(args, kwargs)

There are cases when we want to take into account both backends of two arguments, e.g. a lazy variable and a constant array.

autoray.autoray.random_array_dispatcher(shape, rng=None, **kwargs)

Use the generator’s backend when given, or infer it from shape.

autoray.autoray._RANDOM_DISTS = ('normal', 'uniform', 'rademacher')

autoray.autoray._COMPLEX_TO_REAL_DTYPE

autoray.autoray._RADEMACHER_SIGNS

autoray.autoray._RADEMACHER_ROOTS

autoray.autoray._RADEMACHER_SAMPLERS

autoray.autoray._RADEMACHER_SHARED

autoray.autoray._get_rademacher_table(dtype_name, backend)

The lookup table to index for a rademacher sample: the two signs, or the four roots of unity for a complex dtype_name.

autoray.autoray._sample_rademacher(rng, shape, dtype_name, device, backend)

Draw from {-1, +1}, or from the four roots of unity for a complex dtype_name. Both have modulus one and mean zero.

autoray.autoray._sample_random_array(rng, dist, shape, dtype_name, device, backend)

autoray.autoray.random_array(shape, dist='normal', loc=0.0, scale=1.0, dtype=None, device=None, rng=None, backend=None)

Generate an array of random samples.

Parameters:

• shape (tuple[int]) – Shape of the generated array.

• dist ({"normal", "uniform", "rademacher"}, optional) – Distribution to sample before applying loc and scale. "rademacher" draws each entry from {-1, +1}, or from the four roots of unity for a complex dtype, with equal probability.

• loc (float or complex, optional) – Location applied after sampling.

• scale (float or complex, optional) – Scale applied after sampling.

• dtype (str or dtype_like, optional) – Output dtype, defaulting to that inferred from like, or float64 if there is none.

• device (str or device_like, optional) – Output device. Defaults to that inferred from like.

• rng (int or random number generator, optional) – None uses the backend’s shared random state where one is available. An integer makes a new generator for this call. A backend-specific generator uses and advances its own state, and also supplies the backend, so that like is not needed. Each backend additionally accepts its own seed and state objects, such as a numpy SeedSequence or BitGenerator, or a jax key.

Returns:

Random samples with x = loc + scale * z. Before this transform, a complex normal or rademacher z has total variance one, and a complex uniform z fills the unit square. A rademacher z has modulus exactly one.

Return type:

array

class autoray.autoray.InjectDtypeDevice(fn, device=None, dtype=None)

Wrapper that possibly injects default dtype and device arguments, if not None, into the kwargs of function fn.

__slots__ = ('_device', '_dtype', '_fn')

_fn

_device = None

_dtype = None

__call__(*args, **kwargs)

__repr__()

autoray.autoray._NAME_SPACE_SUBMODULES

class autoray.autoray.AutoNamespace(like=None, device=None, dtype=None, submodule=None)

Mimics a namespace, optionally for a specific backend, device, and dtype, caching the lookup of functions, and injecting default device and dtype arguments for certain creation routines.

Parameters:

• like (array_like, str, or None) – The backend to use, or an object to infer the backend from. If None, the default behavior is to use autoray.do and auto dispatch backend at function call time. If given, the functions are cached at first call.

• device (str, optional) – The device to use for the backend. If None, it will be inferred from the like paramater is that is array-like or set to None.

• dtype (str, optional) – The dtype to use for the backend. If None, it will be inferred from the like parameter if that is array-like or set to None.

• submodule (str, optional) – This is used internally when nesting attribute lookups, e.g. xp.random.normal, xp.linalg.eigh.

_submodule = None

_get_submodule(name)

_get_root_namespace()

_get_fn(name)

__getattr__(name)

__repr__()

autoray.autoray._NAMESPACE_ATTRS = ('_backend', '_device', '_dtype', '_submodule')

autoray.autoray._NAMESPACE_CACHE

autoray.autoray._namespace_key_part(x)

Cached str of a device or dtype, which normalizes them for the namespace cache key. Cached because numpy.dtype.__str__ is slow, and a composed function taking a namespace looks one up on every call.

autoray.autoray._reset_namespace(xp)

Drop the cached lookups of xp and of any submodule it made.

autoray.autoray._reset_namespaces()

Drop the cached function and submodule lookups of every live namespace, keeping the namespace objects themselves, so that a namespace held by a caller stays valid when functions are registered.

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

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.autoray.numpy

autoray.autoray._builtin_dtype_lookup

autoray.autoray.builtins_get_dtype_name(x)

autoray.autoray._rademacher_numpy(rng, shape, dtype_name, device)

autoray.autoray.cupy_to_numpy(x)

autoray.autoray._cupy_parse_device(device)

Check a device string is valid for cupy, returning the gpu index, or None for a bare ‘gpu’ / ‘cuda’ (meaning any gpu).

autoray.autoray.cupy_to_device(x, device)

autoray.autoray.cupy_from_numpy(x, dtype=None, device=None)

autoray.autoray._rademacher_cupy(rng, shape, dtype_name, device)

autoray.autoray.get_jax()

autoray.autoray._jax_parse_device(device)

Parse a device string like “cuda:0” into (platform, index).

autoray.autoray.jax_to_device(x, device)

autoray.autoray.jax_from_numpy(x, dtype=None, device=None)

autoray.autoray._warn_jax_generated_seed()

Warn one time only, since the message applies to every seedless call.

class autoray.autoray.JaxDefaultRNG(seed=None, **kwargs)

Stateful but deterministic random number generator for JAX following numpy’s Generator API.

Create this generator inside a jax.jit function from a seed or key passed to that function. A compiled function that captures a generator created outside it reuses the same values. Later use of that generator can also fail. seed=None warns one time, because compilation chooses the generated seed once.

jax

key

classmethod from_key(key)

binomial(n, p, size=None, **kwargs)

choice(a, size=None, replace=True, p=None, axis=0, **kwargs)

exponential(scale=1.0, size=None, **kwargs)

gumbel(loc=0.0, scale=1.0, size=None, **kwargs)

integers(low, high=None, size=None, **kwargs)

normal(loc=0.0, scale=1.0, size=None, **kwargs)

permutation(x, **kwargs)

poisson(lam=1.0, size=None, **kwargs)

random(size=None, **kwargs)

uniform(low=0.0, high=1.0, size=None, **kwargs)

autoray.autoray.jax_default_rng(seed=None, **kwargs)

autoray.autoray._JAX_DEFAULT_RNG = None

autoray.autoray.jax_random_seed(seed=None)

autoray.autoray._get_jax_default_rng()

autoray.autoray.jax_random_uniform(low=0.0, high=1.0, size=None, **kwargs)

autoray.autoray.jax_random_normal(loc=0.0, scale=1.0, size=None, **kwargs)

autoray.autoray._rademacher_jax(rng, shape, dtype_name, device)

autoray.autoray.aesara_shape(x)

autoray.autoray.autograd_take(x, indices, axis=None)

autoray.autoray.dask_to_numpy(x)

autoray.autoray.dask_eye_wrapper(eye_fn)

autoray.autoray.ctf_array(x)

autoray.autoray.ctf_to_numpy(x)

autoray.autoray.ctf_count_nonzero(x)

autoray.autoray.ctf_get_dtype_name(x)

autoray.autoray.sparse_array(x)

autoray.autoray.sparse_to_numpy(x)

autoray.autoray.sparse_transpose(x, axes=None)

autoray.autoray.sparse_reshape(x, shape)

autoray.autoray.sparse_sum(x, axis=None, keepdims=False, dtype=None, out=None)

autoray.autoray.sparse_prod(x, axis=None, keepdims=False, dtype=None, out=None)

autoray.autoray.sparse_conj(x)

autoray.autoray.sparse_real(x)

autoray.autoray.sparse_imag(x)

autoray.autoray.sparse_count_nonzero(x)

autoray.autoray.sparse_complex(re, im)

autoray.autoray.sparse_random_uniform(low=0.0, high=1.0, size=None, dtype=None, **kwargs)

autoray.autoray.sparse_random_normal(loc=0.0, scale=1.0, size=None, dtype=None, **kwargs)

autoray.autoray.get_tensorflow()

autoray.autoray.tensorflow_pad_wrap(tf_pad)

autoray.autoray.tensorflow_wrap_norm(tf_norm)

class autoray.autoray.TensorflowDefaultRNG(seed=None, **kwargs)

Stateful random number generator for TensorFlow following numpy’s Generator API, compatible with tf.function.

tf

uniform(low=0.0, high=1.0, size=None, **kwargs)

random(size=None, **kwargs)

normal(loc=0.0, scale=1.0, size=None, **kwargs)

integers(low, high=None, size=None, **kwargs)

autoray.autoray.tensorflow_default_rng(seed=None, **kwargs)

autoray.autoray.tensorflow_to_numpy(x)

autoray.autoray._tensorflow_translate_device(device)

Translate a device string like “cuda:0” to tensorflow form.

autoray.autoray.tensorflow_to_device(x, device)

autoray.autoray.tensorflow_from_numpy(x, dtype=None, device=None)

autoray.autoray.tensorflow_indices(dimensions)

autoray.autoray.tensorflow_swapaxes(a, axis1, axis2)

autoray.autoray.tensorflow_cholesky(x, upper=False, **kwargs)

autoray.autoray.tensorflow_solve_triangular(a, b, lower=False, **kwargs)

autoray.autoray.tensorflow_astype(x, dtype)

autoray.autoray.get_torch()

autoray.autoray.torch_shape(x)

autoray.autoray.torch_size(x)

autoray.autoray.torch_linalg_solve_wrap(fn)

autoray.autoray.torch_tensordot_wrap(fn)

autoray.autoray.torch_split_wrap(fn)

autoray.autoray.torch_maybe_convert_dtype_from_str(fn)

autoray.autoray.torch_zeros_ones_wrap(fn)

autoray.autoray.torch_eye_wrap(fn)

autoray.autoray.torch_sort_wrap(fn)

autoray.autoray.torch_flip_wrap(torch_flip)

autoray.autoray.torch_nonzero_wrap(torch_nonzero)

class autoray.autoray.TorchDefaultRNG(seed=None, device=None, dtype=None)

_torch

_set_default_dtype(kwargs)

_set_default_device(kwargs)

integers(low, high=None, size=None, **kwargs)

normal(loc=0.0, scale=1.0, size=None, **kwargs)

random(size=None, **kwargs)

permutation(x, **kwargs)

uniform(low=0.0, high=1.0, size=None, **kwargs)

choice(a, size=None, replace=True, p=None)

autoray.autoray.torch_default_rng(seed=None, **kwargs)

autoray.autoray.torch_random_seed(seed=None)

autoray.autoray._rademacher_torch(rng, shape, dtype_name, device)

autoray.autoray.torch_random_normal(loc=0.0, scale=1.0, size=None, dtype=None, **kwargs)

autoray.autoray.torch_random_uniform(low=0.0, high=1.0, size=None, dtype=None, **kwargs)

autoray.autoray._torch_reduce_translation = [('a', ('input',)), ('axis', ('dim',)), ('keepdims', ('keepdim',))]

autoray.autoray.torch_to_numpy(x)

autoray.autoray.torch_to_device(x, device)

autoray.autoray.torch_from_numpy(x, dtype=None, device=None)

autoray.autoray.torch_copy(x)

autoray.autoray.torch_transpose(x, axes=None)

autoray.autoray.torch_astype(x, dtype)

autoray.autoray.torch_complex(re, im)

autoray.autoray.torch_real(x)

autoray.autoray.torch_imag(x)

autoray.autoray.torch_linalg_eigh(x)

autoray.autoray.torch_linalg_eigvalsh(x)

autoray.autoray.torch_scipy_linalg_solve_triangular(a, b, lower=False, unit_diagonal=False, **kwargs)

autoray.autoray.torch_pad(array, pad_width, mode='constant', constant_values=0)

autoray.autoray.torch_indices(dimensions)

autoray.autoray.torch_take(a, indices, axis=None)

autoray.autoray.torch_trace(x, axis1=0, axis2=1)

autoray.autoray.mxnet_to_numpy(x)

autoray.autoray._paddle_dtype_name_conversion

autoray.autoray.paddle_get_dtype_name(x)

autoray.autoray.paddle_shape(x)

autoray.autoray.paddle_wrap_reduction(fn)

autoray.autoray.paddle_split_wrap(fn)

autoray.autoray.paddle_imag(x)

autoray.autoray.paddle_indices(dimensions)

autoray.autoray.paddle_pad(array, pad_width, mode='constant', constant_values=0)

autoray.autoray.paddle_ravel(x)

autoray.autoray.paddle_real(x)

autoray.autoray.paddle_to_numpy(x)

autoray.autoray.paddle_transpose(a, axes=None)

autoray.autoray.pytensor_shape(x)

autoray.autoray.pytensor_wrap_qr_with_shapes(fn)

autoray.autoray.pytensor_wrap_svd_with_shapes(fn)

autoray.autoray.get_mlx()

autoray.autoray.mlx_to_device(x, device)

class autoray.autoray.MlxDefaultRNG(seed=None, **kwargs)

Stateful but deterministic random number generator for MLX following numpy’s Generator API.

mx

key

_split_key()

_resolve_dtype(dtype)

gumbel(loc=0.0, scale=1.0, size=None, dtype=None, **kwargs)

integers(low, high=None, size=None, dtype=None, **kwargs)

normal(loc=0.0, scale=1.0, size=None, dtype=None, **kwargs)

permutation(x, **kwargs)

random(size=None, **kwargs)

uniform(low=0.0, high=1.0, size=None, dtype=None, **kwargs)

autoray.autoray.mlx_default_rng(seed=None, **kwargs)

class autoray.autoray.MlxSharedRNG

Bases: MlxDefaultRNG

Draws from mlx’s own shared random state, which mx.random.seed sets, rather than from a key of its own.

mx

_split_key()

autoray.autoray._get_mlx_shared_rng()

autoray.autoray.mlx_random_seed(seed=None)

autoray.autoray.mlx_random_uniform(low=0.0, high=1.0, size=None, **kwargs)

autoray.autoray.mlx_random_normal(loc=0.0, scale=1.0, size=None, **kwargs)

autoray.autoray.mlx_count_nonzero(x, axis=None, keepdims=False)

autoray.autoray.mlx_ravel(x, *args, **kwargs)

autoray.autoray.mlx_zeros_ones_wrap(fn)

autoray.autoray.mlx_array_asarray_wrap(fn)

autoray.autoray.mlx_eye_wrap(fn)