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

# autoray

## Submodules

- autoray._version
- autoray.autoray
- autoray.compiler
- autoray.experimental
- autoray.lazy

## Attributes

`numpy` |

## 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.
`is_array`(x) | An alternative leaf tester for addressing only arrays within trees.
`infer_backend_multi`(*arrays) | Infer which backend should be used for a function that takes multiple
`infer_backend`(array) | Get the name of the library that defined the class of `array` - unless
`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]) | Directly provide your own function.
`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

## Package Contents

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

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

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

  Array conjugate.

autoray.dag(x)

  Array Hermitian transpose.

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.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.get_common_dtype(*arrays)

  Compute the minimal dtype sufficient for `arrays`.

autoray.get_dtype_name(x)

  Find string specifier `dtype_name` of array `x`.

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

  Array imaginary part.

autoray.is_array(x)

  An alternative leaf tester for addressing only arrays within trees.

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

  Array real part.

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.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.reshape(x, shape)

  Array reshaped.

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

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

autoray.to_numpy(x)

  Get a numpy version of array `x`.

autoray.transpose(x, *args)

  Array transpose.

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

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

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

  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