Hide navigation sidebar
Hide table of contents sidebar
Skip to content
Back to top
View this page
Toggle table of contents sidebar

# autoray.autoray

AUTORAY - backend agnostic array operations.

Copyright 2019-2023 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` |
`_CREATION_ROUTINES` |
`_CREATION_INJECT` |
`_BACKEND_ALIASES` |
`_MODULE_ALIASES` |
`_SUBMODULE_ALIASES` |
`_FUNC_ALIASES` |
`_CUSTOM_WRAPPERS` |
`_FUNCS` |
`_COMPOSED_FUNCTION_GENERATORS` |
`TREE_MAP_REGISTRY` |
`TREE_APPLY_REGISTRY` |
`TREE_ITER_REGISTRY` |
`IS_CONTAINER_CACHE` |
`TREE_MAPPER_CACHE` |
`TREE_ITER_CACHE` |
`TREE_APPLIER_CACHE` |
`LEAF` |
`_COMPLEX_DTYPES` |
`_DOUBLE_DTYPES` |
`_DTYPE_MAP` |
`_DISPATCHERS` |
`numpy` |
`_builtin_dtype_lookup` |
`_JAX_RANDOM_KEY` |

## Classes

`Leaf` | A singleton object to use as a placeholder in a pytree, for
`Composed` | Compose an `autoray.do` using function. See the main wrapper
`NumpyMimic` | A class to mimic the syntax of using numpy directly.

## 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) |
`_always_the_same`(fn, args, kwargs, backend) |
`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
`register_creation_routine`(backend, fn[, inject_dtype, ...]) | Register a function that creates a new array, with dtype and possibly
`_maybe_inject_dtype_device`(backend, fn, args, kwargs, like) |
`_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_function`(backend, name, fn[, wrap]) | Directly provide your own function.
`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
`is_array`(x) | An alternative leaf tester for addressing only arrays within trees.
`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) |
`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`(dtype_name, like) | Turn string specifier `dtype_name` into dtype of backend `like`.
`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`.
`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) |
`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]) |
`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
`_get_mimic_function_or_attribute`(self, fn) |
`builtins_get_dtype_name`(x) |
`numpy_to_numpy`(x) |
`cupy_to_numpy`(x) |
`jax_random_seed`([seed]) |
`jax_random_get_key`() |
`jax_random_uniform`([low, high, size]) |
`jax_random_normal`([loc, scale, size]) |
`jax_to_numpy`(x) |
`aesara_shape`(x) |
`dask_to_numpy`(x) |
`dask_eye_wrapper`(fn) |
`mars_to_numpy`(x) |
`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_random_uniform`([low, high, size, dtype]) |
`sparse_random_normal`([loc, scale, size, dtype]) |
`tensorflow_to_numpy`(x) |
`tensorflow_indices`(dimensions) |
`tensorflow_pad_wrap`(tf_pad) |
`torch_shape`(x) |
`torch_size`(x) |
`torch_to_numpy`(x) |
`torch_transpose`(x[, axes]) |
`torch_count_nonzero`(x) |
`torch_astype`(x, dtype) |
`_torch_get_dtype_name`(dtype) |
`torch_get_dtype_name`(x) |
`torch_real`(x) |
`torch_imag`(x) |
`torch_linalg_solve_wrap`(fn) |
`torch_linalg_eigh`(x) |
`torch_linalg_eigvalsh`(x) |
`torch_tensordot_wrap`(fn) |
`torch_pad`(array, pad_width[, mode, constant_values]) |
`torch_split_wrap`(fn) |
`torch_zeros_ones_wrap`(fn) |
`torch_eye_wrap`(fn) |
`torch_sort_wrap`(fn) |
`torch_indices`(dimensions) |
`mxnet_to_numpy`(x) |

## Module Contents

autoray.autoray.do(fn, *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.

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

autoray.autoray._always_the_same(fn, args, kwargs, backend)

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

autoray.autoray._CREATION_INJECT

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.

  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._maybe_inject_dtype_device(backend, fn, args, kwargs, like)

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 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.autoray.register_function(backend, name, fn, wrap=False)

  Directly provide your own function.

  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) – The function to register.

    - wrap (bool, optional) – Whether to wrap the old function like `fn(old_fn)` rather than directly supply the entire new function.

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

  An alternative leaf tester for addressing only arrays within trees.

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__()

    Return repr(self).

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)

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

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

  register(backend, fn=None)

    Register a different implementation for `backend`.

  make_function(backend)

    Make a new function for the specific `backend`.

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

    Return repr(self).

autoray.autoray.compose(fn, *, 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.

  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.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(dtype_name, like)

  Turn string specifier `dtype_name` into dtype of backend `like`.

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`.

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)

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.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._get_mimic_function_or_attribute(self, fn)

class autoray.autoray.NumpyMimic(submodule=None)

  A class to mimic the syntax of using numpy directly.

  __getattribute__(attr)

    Return getattr(self, name).

  static __repr__()

    Return repr(self).

autoray.autoray.numpy

autoray.autoray._builtin_dtype_lookup

autoray.autoray.builtins_get_dtype_name(x)

autoray.autoray.numpy_to_numpy(x)

autoray.autoray.cupy_to_numpy(x)

autoray.autoray._JAX_RANDOM_KEY = None

autoray.autoray.jax_random_seed(seed=None)

autoray.autoray.jax_random_get_key()

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

autoray.autoray.aesara_shape(x)

autoray.autoray.dask_to_numpy(x)

autoray.autoray.dask_eye_wrapper(fn)

autoray.autoray.mars_to_numpy(x)

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_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.tensorflow_to_numpy(x)

autoray.autoray.tensorflow_indices(dimensions)

autoray.autoray.tensorflow_pad_wrap(tf_pad)

autoray.autoray.torch_shape(x)

autoray.autoray.torch_size(x)

autoray.autoray.torch_to_numpy(x)

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

autoray.autoray.torch_count_nonzero(x)

autoray.autoray.torch_astype(x, dtype)

autoray.autoray._torch_get_dtype_name(dtype)

autoray.autoray.torch_get_dtype_name(x)

autoray.autoray.torch_real(x)

autoray.autoray.torch_imag(x)

autoray.autoray.torch_linalg_solve_wrap(fn)

autoray.autoray.torch_linalg_eigh(x)

autoray.autoray.torch_linalg_eigvalsh(x)

autoray.autoray.torch_tensordot_wrap(fn)

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

autoray.autoray.torch_split_wrap(fn)

autoray.autoray.torch_zeros_ones_wrap(fn)

autoray.autoray.torch_eye_wrap(fn)

autoray.autoray.torch_sort_wrap(fn)

autoray.autoray.torch_indices(dimensions)

autoray.autoray.mxnet_to_numpy(x)