autoray.autoray
===============

.. py:module:: autoray.autoray

.. autoapi-nested-parse::

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

.. autoapisummary::

   autoray.autoray._global_backend
   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._CUSTOM_BACKENDS
   autoray.autoray.multi_class_priorities
   autoray.autoray._CREATION_ROUTINES
   autoray.autoray._CREATION_INJECT
   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.TREE_MAP_REGISTRY
   autoray.autoray.TREE_APPLY_REGISTRY
   autoray.autoray.TREE_ITER_REGISTRY
   autoray.autoray.IS_CONTAINER_CACHE
   autoray.autoray.TREE_MAPPER_CACHE
   autoray.autoray.TREE_ITER_CACHE
   autoray.autoray.TREE_APPLIER_CACHE
   autoray.autoray.LEAF
   autoray.autoray._COMPLEX_DTYPES
   autoray.autoray._DOUBLE_DTYPES
   autoray.autoray._DTYPE_MAP
   autoray.autoray._DISPATCHERS
   autoray.autoray.numpy
   autoray.autoray._builtin_dtype_lookup
   autoray.autoray._JAX_RANDOM_KEY


Classes
-------

.. autoapisummary::

   autoray.autoray.Leaf
   autoray.autoray.Composed
   autoray.autoray.NumpyMimic


Functions
---------

.. autoapisummary::

   autoray.autoray.do
   autoray.autoray._default_infer_from_sig
   autoray.autoray._default_infer_from_sig_threadaware
   autoray.autoray._always_the_same
   autoray.autoray.get_backend
   autoray.autoray.set_backend
   autoray.autoray.backend_like
   autoray.autoray.register_backend
   autoray.autoray._infer_class_backend_cached
   autoray.autoray.infer_backend
   autoray.autoray._infer_class_backend_multi_cached
   autoray.autoray.infer_backend_multi
   autoray.autoray.register_creation_routine
   autoray.autoray._maybe_inject_dtype_device
   autoray.autoray._choose_backend
   autoray.autoray.choose_backend
   autoray.autoray.import_lib_fn
   autoray.autoray.get_lib_fn
   autoray.autoray.register_function
   autoray.autoray.tree_register_container
   autoray.autoray.is_not_container
   autoray.autoray.is_array
   autoray.autoray.identity
   autoray.autoray.tree_map
   autoray.autoray.empty
   autoray.autoray.tree_iter
   autoray.autoray.nothing
   autoray.autoray.tree_apply
   autoray.autoray.is_leaf_placeholder
   autoray.autoray.tree_flatten
   autoray.autoray.tree_unflatten
   autoray.autoray.tree_map_tuple
   autoray.autoray.tree_iter_tuple
   autoray.autoray.tree_apply_tuple
   autoray.autoray.tree_map_list
   autoray.autoray.tree_iter_list
   autoray.autoray.tree_apply_list
   autoray.autoray.tree_map_dict
   autoray.autoray.tree_iter_dict
   autoray.autoray.tree_apply_dict
   autoray.autoray.compose
   autoray.autoray.shape
   autoray.autoray.ndim
   autoray.autoray.size
   autoray.autoray.conj
   autoray.autoray.transpose
   autoray.autoray.dag
   autoray.autoray.real
   autoray.autoray.imag
   autoray.autoray.reshape
   autoray.autoray.to_backend_dtype
   autoray.autoray.get_dtype_name
   autoray.autoray.get_common_dtype
   autoray.autoray.astype
   autoray.autoray.to_numpy
   autoray.autoray.svd_not_full_matrices_wrapper
   autoray.autoray.svd_sUV_to_UsVH_wrapper
   autoray.autoray.svd_UsV_to_UsVH_wrapper
   autoray.autoray.svd_manual_full_matrices_kwarg
   autoray.autoray.qr_allow_fat
   autoray.autoray.tril_to_band_part
   autoray.autoray.triu_to_band_part
   autoray.autoray.cholesky_lower
   autoray.autoray.binary_allow_1d_rhs_wrap
   autoray.autoray.scale_random_uniform_manually
   autoray.autoray.scale_random_normal_manually
   autoray.autoray.with_dtype_wrapper
   autoray.autoray.translate_wrapper
   autoray.autoray.make_translator
   autoray.autoray.complex_add_re_im
   autoray.autoray.allclose
   autoray.autoray.wrap_args_kwargs_from_raw
   autoray.autoray.register_dispatch
   autoray.autoray.default_dispatcher
   autoray.autoray.join_array_dispatcher
   autoray.autoray.einsum_dispatcher
   autoray.autoray.binary_dispatcher
   autoray.autoray._get_mimic_function_or_attribute
   autoray.autoray.builtins_get_dtype_name
   autoray.autoray.numpy_to_numpy
   autoray.autoray.cupy_to_numpy
   autoray.autoray.jax_random_seed
   autoray.autoray.jax_random_get_key
   autoray.autoray.jax_random_uniform
   autoray.autoray.jax_random_normal
   autoray.autoray.jax_to_numpy
   autoray.autoray.aesara_shape
   autoray.autoray.dask_to_numpy
   autoray.autoray.dask_eye_wrapper
   autoray.autoray.mars_to_numpy
   autoray.autoray.ctf_array
   autoray.autoray.ctf_to_numpy
   autoray.autoray.ctf_count_nonzero
   autoray.autoray.ctf_get_dtype_name
   autoray.autoray.sparse_array
   autoray.autoray.sparse_to_numpy
   autoray.autoray.sparse_transpose
   autoray.autoray.sparse_reshape
   autoray.autoray.sparse_sum
   autoray.autoray.sparse_prod
   autoray.autoray.sparse_conj
   autoray.autoray.sparse_real
   autoray.autoray.sparse_imag
   autoray.autoray.sparse_count_nonzero
   autoray.autoray.sparse_random_uniform
   autoray.autoray.sparse_random_normal
   autoray.autoray.tensorflow_to_numpy
   autoray.autoray.tensorflow_indices
   autoray.autoray.tensorflow_pad_wrap
   autoray.autoray.torch_shape
   autoray.autoray.torch_size
   autoray.autoray.torch_to_numpy
   autoray.autoray.torch_transpose
   autoray.autoray.torch_count_nonzero
   autoray.autoray.torch_astype
   autoray.autoray._torch_get_dtype_name
   autoray.autoray.torch_get_dtype_name
   autoray.autoray.torch_real
   autoray.autoray.torch_imag
   autoray.autoray.torch_linalg_solve_wrap
   autoray.autoray.torch_linalg_eigh
   autoray.autoray.torch_linalg_eigvalsh
   autoray.autoray.torch_tensordot_wrap
   autoray.autoray.torch_pad
   autoray.autoray.torch_split_wrap
   autoray.autoray.torch_zeros_ones_wrap
   autoray.autoray.torch_eye_wrap
   autoray.autoray.torch_sort_wrap
   autoray.autoray.torch_indices
   autoray.autoray.mxnet_to_numpy


Module 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:: _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``.


.. py:data:: _global_backend
   :value: None


.. py:data:: _inferrer_global

.. py:data:: _infer_auto

.. py:data:: _global_backends_threadaware

.. py:data:: _inferrers_threadaware

.. py:data:: _importing_thrid

.. py:data:: _backend_lock

.. py:function:: _default_infer_from_sig_threadaware(fn, args, kwargs)

.. py:function:: _always_the_same(fn, args, kwargs, backend)

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

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

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

.. py:function:: _infer_class_backend_multi_cached(classes)

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

.. py:data:: _CREATION_INJECT

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

   :param backend: The backend to register the function for.
   :type backend: str
   :param fn: The name of the function to register.
   :type fn: str
   :param inject_dtype: Whether to inject a `dtype` argument based on the `like` argument.
   :type inject_dtype: bool, optional
   :param inject_device: Whether to inject a `device` argument based on the `like` argument.
   :type inject_device: bool, optional


.. py:function:: _maybe_inject_dtype_device(backend, fn, args, kwargs, like)

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


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


.. py:data:: _BACKEND_ALIASES

.. py:data:: _MODULE_ALIASES

.. py:data:: _SUBMODULE_ALIASES

.. py:data:: _FUNC_ALIASES

.. py:data:: _CUSTOM_WRAPPERS

.. py:data:: _FUNCS

.. py:data:: _COMPOSED_FUNCTION_GENERATORS

.. py:function:: import_lib_fn(backend, fn)

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

.. py:data:: TREE_APPLY_REGISTRY

.. py:data:: TREE_ITER_REGISTRY

.. py:function:: tree_register_container(cls, mapper, iterator, applier)

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

   :param cls: The container type to register.
   :type cls: type
   :param mapper: A function that takes ``f``, ``tree`` and ``is_leaf`` and returns a new
                  tree of type ``cls`` with ``f`` applied to all leaves.
   :type mapper: callable
   :param applier: A function that takes ``f``, ``tree`` and ``is_leaf`` and applies ``f``
                   to all leaves in ``tree``.
   :type applier: callable


.. py:data:: IS_CONTAINER_CACHE

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


.. py:function:: is_array(x)

   An alternative leaf tester for addressing only arrays within trees.


.. py:function:: identity(f, tree, is_leaf)

.. py:data:: TREE_MAPPER_CACHE

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

.. py:data:: TREE_ITER_CACHE

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

.. py:data:: TREE_APPLIER_CACHE

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

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


   .. py:attribute:: __slots__
      :value: ()



   .. py:method:: __repr__()

      Return repr(self).



.. py:data:: LEAF

.. py:function:: is_leaf_placeholder(x)

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

.. py:function:: tree_iter_tuple(tree, is_leaf)

.. py:function:: tree_apply_tuple(f, tree, is_leaf)

.. py:function:: tree_map_list(f, tree, is_leaf)

.. py:function:: tree_iter_list(tree, is_leaf)

.. py:function:: tree_apply_list(f, tree, is_leaf)

.. py:function:: tree_map_dict(f, tree, is_leaf)

.. py:function:: tree_iter_dict(tree, is_leaf)

.. py:function:: tree_apply_dict(f, tree, is_leaf)

.. py:class:: Composed(fn, name=None)

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


   .. py:method:: register(backend, fn=None)

      Register a different implementation for ``backend``.



   .. py:method:: make_function(backend)

      Make a new function for the specific ``backend``.



   .. py:method:: __call__(*args, like=None, **kwargs)


   .. py:method:: __repr__()

      Return repr(self).



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

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


.. py:data:: _COMPLEX_DTYPES

.. py:data:: _DOUBLE_DTYPES

.. py:data:: _DTYPE_MAP

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

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


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

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


.. py:function:: to_numpy(x)

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


.. py:function:: svd_not_full_matrices_wrapper(fn)

.. py:function:: svd_sUV_to_UsVH_wrapper(fn)

.. py:function:: svd_UsV_to_UsVH_wrapper(fn)

.. py:function:: svd_manual_full_matrices_kwarg(fn)

.. py:function:: qr_allow_fat(fn)

.. py:function:: tril_to_band_part(fn)

.. py:function:: triu_to_band_part(fn)

.. py:function:: cholesky_lower(fn)

.. py:function:: binary_allow_1d_rhs_wrap(fn)

.. py:function:: scale_random_uniform_manually(fn)

.. py:function:: scale_random_normal_manually(fn)

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


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


.. py:function:: make_translator(t)

.. py:function:: complex_add_re_im(re, im)

.. py:function:: allclose(x, y, rtol=1e-05, atol=1e-08)

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


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

   :param fun: The name of the function to register the dispatcher for.
   :type fun: str
   :param dispatcher: The dispatcher function to use. This should take the arguments and
                      keyword arguments of the function and return the backend to use.
   :type dispatcher: callable
   :param raw_signature: The ``dispatcher`` has signature ``(*args, **kwargs)`` if ``True``,
                         otherwise it has signature ``(args, kwargs)``.
   :type raw_signature: bool, optional


.. py:function:: default_dispatcher(args, kwargs)

   Try to infer backend from first argument passed to function.


.. py:data:: _DISPATCHERS

.. py:function:: join_array_dispatcher(args, kwargs)

   Dispatcher for functions where first argument is a sequence.


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


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


.. py:function:: _get_mimic_function_or_attribute(self, fn)

.. py:class:: NumpyMimic(submodule=None)

   A class to mimic the syntax of using `numpy` directly.


   .. py:method:: __getattribute__(attr)

      Return getattr(self, name).



   .. py:method:: __repr__()
      :staticmethod:


      Return repr(self).



.. py:data:: numpy

.. py:data:: _builtin_dtype_lookup

.. py:function:: builtins_get_dtype_name(x)

.. py:function:: numpy_to_numpy(x)

.. py:function:: cupy_to_numpy(x)

.. py:data:: _JAX_RANDOM_KEY
   :value: None


.. py:function:: jax_random_seed(seed=None)

.. py:function:: jax_random_get_key()

.. py:function:: jax_random_uniform(low=0.0, high=1.0, size=None, **kwargs)

.. py:function:: jax_random_normal(loc=0.0, scale=1.0, size=None, **kwargs)

.. py:function:: jax_to_numpy(x)

.. py:function:: aesara_shape(x)

.. py:function:: dask_to_numpy(x)

.. py:function:: dask_eye_wrapper(fn)

.. py:function:: mars_to_numpy(x)

.. py:function:: ctf_array(x)

.. py:function:: ctf_to_numpy(x)

.. py:function:: ctf_count_nonzero(x)

.. py:function:: ctf_get_dtype_name(x)

.. py:function:: sparse_array(x)

.. py:function:: sparse_to_numpy(x)

.. py:function:: sparse_transpose(x, axes=None)

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

.. py:function:: sparse_sum(x, axis=None, keepdims=False, dtype=None, out=None)

.. py:function:: sparse_prod(x, axis=None, keepdims=False, dtype=None, out=None)

.. py:function:: sparse_conj(x)

.. py:function:: sparse_real(x)

.. py:function:: sparse_imag(x)

.. py:function:: sparse_count_nonzero(x)

.. py:function:: sparse_random_uniform(low=0.0, high=1.0, size=None, dtype=None, **kwargs)

.. py:function:: sparse_random_normal(loc=0.0, scale=1.0, size=None, dtype=None, **kwargs)

.. py:function:: tensorflow_to_numpy(x)

.. py:function:: tensorflow_indices(dimensions)

.. py:function:: tensorflow_pad_wrap(tf_pad)

.. py:function:: torch_shape(x)

.. py:function:: torch_size(x)

.. py:function:: torch_to_numpy(x)

.. py:function:: torch_transpose(x, axes=None)

.. py:function:: torch_count_nonzero(x)

.. py:function:: torch_astype(x, dtype)

.. py:function:: _torch_get_dtype_name(dtype)

.. py:function:: torch_get_dtype_name(x)

.. py:function:: torch_real(x)

.. py:function:: torch_imag(x)

.. py:function:: torch_linalg_solve_wrap(fn)

.. py:function:: torch_linalg_eigh(x)

.. py:function:: torch_linalg_eigvalsh(x)

.. py:function:: torch_tensordot_wrap(fn)

.. py:function:: torch_pad(array, pad_width, mode='constant', constant_values=0)

.. py:function:: torch_split_wrap(fn)

.. py:function:: torch_zeros_ones_wrap(fn)

.. py:function:: torch_eye_wrap(fn)

.. py:function:: torch_sort_wrap(fn)

.. py:function:: torch_indices(dimensions)

.. py:function:: mxnet_to_numpy(x)

