autoray
=======

.. py:module:: autoray


Subpackages
-----------

.. toctree::
   :maxdepth: 1

   /autoapi/autoray/experimental/index
   /autoapi/autoray/lazy/index


Submodules
----------

.. toctree::
   :maxdepth: 1

   /autoapi/autoray/_version/index
   /autoapi/autoray/autoray/index
   /autoapi/autoray/compiler/index


Attributes
----------

.. autoapisummary::

   autoray.numpy


Functions
---------

.. autoapisummary::

   autoray.do
   autoray.get_backend
   autoray.set_backend
   autoray.backend_like
   autoray.infer_backend
   autoray.infer_backend_multi
   autoray.get_lib_fn
   autoray.shape
   autoray.ndim
   autoray.size
   autoray.conj
   autoray.transpose
   autoray.dag
   autoray.real
   autoray.imag
   autoray.reshape
   autoray.to_backend_dtype
   autoray.astype
   autoray.get_dtype_name
   autoray.get_common_dtype
   autoray.to_numpy
   autoray.register_backend
   autoray.register_function
   autoray.is_array
   autoray.tree_map
   autoray.tree_iter
   autoray.tree_apply
   autoray.tree_flatten
   autoray.tree_unflatten
   autoray.compose
   autoray.autojit


Package Contents
----------------

.. py:function:: 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.

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


.. py:function:: get_backend(get_globally='auto')

   Return the universally set backend, if any.

   :param get_globally: 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.
   :type get_globally: {"auto", False, True}, optional

   :returns: **backend** -- The name of the backend, or None if no backend is set.
   :rtype: str or None


.. py:function:: set_backend(like, set_globally='auto')

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

   :param like: The backend to set. If an array, the backend of the array's class will
                be set.
   :type like: str or array
   :param set_globally: 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).
   :type set_globally: {"auto", False, True}, optional


.. py:function:: 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.

   :param like: The backend to set. If an array, the backend of the array's class will
                be set.
   :type like: str or array
   :param set_globally: 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).
   :type set_globally: {"auto", False, True}, optional


.. py:function:: 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.


.. py:function:: 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.


.. py:function:: 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.

   :param backend: The module defining the array class to dispatch on.
   :type backend: str
   :param fn: The function to retrieve.
   :type fn: str

   :rtype: callable


.. py:function:: 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.

   :param x: 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.
   :type x: array_like

   :returns: **shape** -- The size of each dimension of the array.
   :rtype: tuple of int


.. py:function:: 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.

   :param x: The array to get the number of dimensions of. It can be an arbitrary
             nested list or tuple of arrays and scalars.
   :type x: array_like

   :returns: **ndim**
   :rtype: int


.. py:function:: 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.

   :param x: The array to get the size of. It can be an arbitrary nested list or
             tuple of arrays and scalars.
   :type x: array_like

   :returns: **size**
   :rtype: int


.. py:function:: conj(x)

   Array conjugate.


.. py:function:: transpose(x, *args)

   Array transpose.


.. py:function:: dag(x)

   Array Hermitian transpose.


.. py:function:: real(x)

   Array real part.


.. py:function:: imag(x)

   Array imaginary part.


.. py:function:: reshape(x, shape)

   Array reshaped.


.. py:function:: to_backend_dtype(dtype_name, like)

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


.. py:function:: astype(x, dtype_name, **kwargs)

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


.. py:function:: get_dtype_name(x)

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


.. py:function:: get_common_dtype(*arrays)

   Compute the minimal dtype sufficient for ``arrays``.


.. py:function:: to_numpy(x)

   Get a numpy version of array ``x``.


.. py:function:: register_backend(cls, name)

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

   :param cls: The array class itself.
   :type cls: type
   :param name: 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.
   :type name: str


.. py:function:: register_function(backend, name, fn, wrap=False)

   Directly provide your own function.

   :param backend: The name of the backend to register the function for.
   :type backend: str
   :param name: Name of the function, e.g. `'sum'` or `'linalg.svd'`.
   :type name: str
   :param fn: The function to register.
   :type fn: callable
   :param wrap: Whether to wrap the old function like ``fn(old_fn)`` rather than
                directly supply the entire new function.
   :type wrap: bool, optional


.. py:function:: is_array(x)

   An alternative leaf tester for addressing only arrays within trees.


.. py:function:: tree_map(f, tree, is_leaf=is_not_container)

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

   :param f: A function to apply to all leaves in ``tree``.
   :type f: callable
   :param tree: A nested sequence of tuples, lists, dicts and other objects.
   :type tree: pytree
   :param is_leaf: A function to determine if an object is a leaf, ``f`` is only applied
                   to objects for which ``is_leaf(x)`` returns ``True``.
   :type is_leaf: callable

   :rtype: pytree


.. py:function:: tree_iter(tree, is_leaf=is_not_container)

   Iterate over all leaves in ``tree``.

   :param f: A function to apply to all leaves in ``tree``.
   :type f: callable
   :param tree: A nested sequence of tuples, lists, dicts and other objects.
   :type tree: pytree
   :param is_leaf: A function to determine if an object is a leaf, ``f`` is only applied
                   to objects for which ``is_leaf(x)`` returns ``True``.
   :type is_leaf: callable


.. py:function:: tree_apply(f, tree, is_leaf=is_not_container)

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

   :param f: A function to apply to all leaves in ``tree``.
   :type f: callable
   :param tree: A nested sequence of tuples, lists, dicts and other objects.
   :type tree: pytree
   :param is_leaf: A function to determine if an object is a leaf, ``f`` is only applied
                   to objects for which ``is_leaf(x)`` returns ``True``.
   :type is_leaf: callable


.. py:function:: tree_flatten(tree, is_leaf=is_not_container, get_ref=False)

   Flatten ``tree`` into a list of leaves.

   :param tree: A nested sequence of tuples, lists, dicts and other objects.
   :type tree: pytree
   :param is_leaf: 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.
   :type is_leaf: callable
   :param get_ref: If ``True``, a reference tree is also returned which can be used to
                   reconstruct the original tree from a flattened list.
   :type get_ref: bool

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


.. py:function:: tree_unflatten(objs, tree, is_leaf=is_leaf_placeholder)

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

   :param objs: A sequence of objects to be unflattened into a pytree.
   :type objs: sequence
   :param tree: A nested sequence of tuples, lists, dicts and other objects, the objs
                will be inserted into a new pytree of the same structure.
   :type tree: pytree
   :param is_leaf: 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)``.
   :type is_leaf: callable

   :rtype: pytree


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

   :param fn: The funtion to compose, and its default implementation.
   :type fn: callable
   :param name: The name of the composed function. If not provided, the name of the
                function will be used.
   :type name: str, optional


.. py:data:: numpy

.. py:function:: 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.

   :param fn: The autoray function to compile.
   :type fn: callable
   :param backend: If set, use this as the default backend.
   :type backend: {None, 'python', 'jax', 'torch', 'tensorflow'}, optional
   :param compiler_opts: Dict of dicts when you can supply options for each compiler backend
                         separately, e.g.:
                         ``@autojit(compiler_opts={'tensorflow': {'jit_compile': True}})``.
   :type compiler_opts: dict[dict], optional

   :returns: **cfn** -- The function with auto compilation.
   :rtype: callable


