"""
AUTORAY - backend agnostic array operations.


Copyright 2019-2026 Johnnie Gray

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

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

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

import contextlib
import functools
import importlib
import itertools
import math
import numbers
import operator
import threading
from collections import OrderedDict, defaultdict
from inspect import signature


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

    Parameters
    ----------
    fn : str
        Name of the function to do, e.g. 'sum' or 'linalg.svd'.
    args
        Positional arguments to pass to the function.
    like : str or array, optional
        Backend to use, either as an explicit backend name or an example array
        to infer the backend from. If not specified, the backend is inferred
        from the first argument, or from a globally set backend if any.
    kwargs
        Keyword arguments to pass to the function.

    Examples
    --------

    Works on numpy arrays:

        >>> import numpy as np
        >>> x_np = np.random.uniform(size=[5])
        >>> y_np = do('sqrt', x_np)
        >>> y_np
        array([0.32464973, 0.90379787, 0.85037325, 0.88729814, 0.46768083])

        >>> type(y_np)
        numpy.ndarray

    Works on cupy arrays:

        >>> import cupy as cp
        >>> x_cp = cp.random.uniform(size=[5])
        >>> y_cp = do('sqrt', x_cp)
        >>> y_cp
        array([0.44541656, 0.88713113, 0.92626237, 0.64080557, 0.69620767])

        >>> type(y_cp)
        cupy.core.core.ndarray

    Works on tensorflow arrays:

        >>> import tensorflow as tf
        >>> x_tf = tf.random.uniform(shape=[5])
        >>> y_tf = do('sqrt', x_tf)
        >>> y_tf
        <tf.Tensor 'Sqrt_1:0' shape=(5,) dtype=float32>

        >>> type(y_tf)
        tensorflow.python.framework.ops.Tensor

    You get the idea.

    For functions that don't dispatch on the first argument you can use the
    ``like`` keyword:

        >>> do('eye', 3, like=x_tf)
        <tf.Tensor: id=91, shape=(3, 3), dtype=float32>
    """
    backend = _choose_backend(fn, args, kwargs, like=like)
    func = get_lib_fn(backend, fn)
    return func(*args, **kwargs)


# ------------------------- efficiently dispatching ------------------------- #


def _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``.
    """
    return _DISPATCHERS[fn](*args, **kwargs)


_global_backend = None
_inferrer_global = _default_infer_from_sig

# this is the function that autoray uses when `do` is called without an
# explicit like/backend argument. It is set to `_default_infer_from_sig` by
# default, but can be set to `_always_the_same` if a global backend is set e.g.
_infer_auto = _inferrer_global

# if a thread that isn't the 'importing' thread tries to set a backend, this
# by default turns on thread aware dispatching, but once such custom sub
# backends have been reset, the global values above are used again.
_global_backends_threadaware = {}
_inferrers_threadaware = {}
_importing_thrid = threading.get_ident()
_backend_lock = threading.Lock()


def _default_infer_from_sig_threadaware(fn, args, kwargs):
    # check for a thread aware inferrer, default to the global inferrer
    thrid = threading.get_ident()
    return _inferrers_threadaware.get(thrid, _inferrer_global)(
        fn, args, kwargs
    )


def _always_the_same(fn, args, kwargs, backend):
    return backend


def 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 : str or None
        The name of the backend, or None if no backend is set.
    """
    if get_globally == "auto":
        get_globally = threading.get_ident() == _importing_thrid

    if get_globally:
        backend = _global_backend
    else:
        thrid = threading.get_ident()
        backend = _global_backends_threadaware.get(thrid, None)

    return backend


def 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).
    """
    global _global_backend
    global _infer_auto
    global _inferrer_global

    if like is None:
        backend = None
        inferrer = _default_infer_from_sig
    elif isinstance(like, str):
        backend = like
        inferrer = functools.partial(_always_the_same, backend=backend)
    else:
        backend = _infer_class_backend_cached(like.__class__)
        inferrer = functools.partial(_always_the_same, backend=backend)

    if set_globally == "auto":
        set_globally = threading.get_ident() == _importing_thrid

    if set_globally:
        _global_backend = backend
        _inferrer_global = inferrer
        if not _inferrers_threadaware:
            # only revert the actual function if no subthread backends set
            _infer_auto = inferrer
    else:
        thrid = threading.get_ident()
        _backend_lock.acquire()
        if backend is None:
            _global_backends_threadaware.pop(thrid)
            _inferrers_threadaware.pop(thrid)
        else:
            _global_backends_threadaware[thrid] = backend
            _inferrers_threadaware[thrid] = inferrer

        if _inferrers_threadaware:
            # a subthread backend has been set, so we need to be thread aware
            _infer_auto = _default_infer_from_sig_threadaware
        else:
            # no subthread backend has been set anymore, so we can ignore
            # threads and just use the global inferrer
            _infer_auto = _inferrer_global
        _backend_lock.release()


@contextlib.contextmanager
def 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).
    """
    if set_globally == "auto":
        set_globally = threading.get_ident() == _importing_thrid

    old_backend = get_backend(get_globally=set_globally)
    try:
        set_backend(like, set_globally)
        yield
    finally:
        set_backend(old_backend, set_globally)


_CUSTOM_BACKENDS = {}


def 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.
    """
    if not isinstance(cls, type):
        raise TypeError("The array class itself should be supplied.")

    global _CUSTOM_BACKENDS
    _CUSTOM_BACKENDS[cls] = name


@functools.lru_cache(None)
def _infer_class_backend_cached(cls):
    try:
        import numpy as _numpy

        if issubclass(cls, _numpy.ndarray):
            return "numpy"
    except ImportError:
        # numpy not installed
        pass

    if cls in _CUSTOM_BACKENDS:
        return _CUSTOM_BACKENDS[cls]

    lib = cls.__module__.split(".")[0]

    # check if lib should mapped entirely to another lib
    backend = _BACKEND_ALIASES.get(lib, lib)

    return backend


def 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.
    """
    return _infer_class_backend_cached(array.__class__)


multi_class_priorities = {
    "builtins": -2,
    "numpy": -1,
    "autoray.lazy": 1,
}


@functools.lru_cache(None)
def _infer_class_backend_multi_cached(classes):
    return max(
        map(_infer_class_backend_cached, classes),
        key=lambda n: multi_class_priorities.get(n, 0),
    )


def 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.
    """
    return _infer_class_backend_multi_cached(
        tuple(array.__class__ for array in arrays)
    )


# the set of functions that create new arrays, with `dtype` and possibly
# `device` kwargs, that should be inferred from the like argument
_CREATION_ROUTINES = {
    "array": (False, False),
    "asarray": (False, False),
    "empty": (True, False),
    "eye": (True, False),
    "full": (True, False),
    "identity": (True, False),
    "ones": (True, False),
    "zeros": (True, False),
    "random.default_rng": (False, False),
    # TODO: should these be included?
    # "arange",
    # "geomspace",
    # "linspace",
    # "logspace",
}

# cache for whether backends have a device attribute
_CREATION_INJECT = {}


def 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.
    """
    if fn not in _CREATION_ROUTINES:
        import warnings

        warnings.warn(
            f"Registering creation routine for backend '{backend}' and "
            f"function '{fn}' which is not in the default set of known "
            f"creation routines: {_CREATION_ROUTINES.keys()}, this will "
            "currently have no effect.",
        )

    _CREATION_INJECT[backend, fn] = (inject_dtype, inject_device)


def _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.
    """
    if like is None:
        # infer from function call (or global backend)
        return _infer_auto(fn, args, kwargs)
    elif isinstance(like, str):
        # explicit backend
        return like
    else:
        # explicit example array
        backend = _infer_class_backend_cached(like.__class__)

        # check if we should set some extra defaults based on the example array
        if fn in _CREATION_ROUTINES:
            try:
                # check for backend specific defaults
                inject_dtype, inject_device = _CREATION_INJECT[backend, fn]
            except KeyError:
                # populate with non backend specific defaults
                inject_dtype, inject_device = _CREATION_ROUTINES[fn]
                _CREATION_INJECT[backend, fn] = (inject_dtype, inject_device)

            if inject_dtype:
                kwargs.setdefault("dtype", getattr(like, "dtype", type(like)))
            if inject_device:
                kwargs.setdefault("device", like.device)

        return backend


def 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.
    """
    return _choose_backend(fn, args, kwargs, like=like)


# ------------------- importing and caching the function -------------------- #

# lookup for mapping entire lib to another
_BACKEND_ALIASES = {}

# global (non function specific) aliases
_MODULE_ALIASES = {}

# lookup for when functions are elsewhere than the expected location
_SUBMODULE_ALIASES = {}

# lookup for when functions are simply called something else
_FUNC_ALIASES = {}

# custom wrappers for when functions don't just have different location or
#     name. For example, when kwargs need to be translated or results modified
_CUSTOM_WRAPPERS = {}

# actual cache of funtions to use - this is populated lazily and can be used
#     to directly set an implementation of a function for a specific backend
_FUNCS = {}

# these are functions where a default implementation can be constructed
#     (composed of other functions), but this is only done lazily
_COMPOSED_FUNCTION_GENERATORS = {}


def import_lib_fn(backend, fn):
    # first check explicitly composed functions -> if the function hasn't been
    # called directly yet, it won't have been loaded into the cache, and needs
    # generating before e.g. the ``do`` verrsion will work
    if fn in _COMPOSED_FUNCTION_GENERATORS:
        return _COMPOSED_FUNCTION_GENERATORS[fn](backend)

    try:
        # submodule where function is found for backend,
        #     e.g. ['tensorflow', trace'] -> 'tensorflow.linalg'
        try:
            full_location = _SUBMODULE_ALIASES[backend, fn]

            # if explicit submodule alias given, don't use prepended location
            #     for example, ('torch', 'linalg.svd') -> torch.svd
            only_fn = fn.split(".")[-1]

        except KeyError:
            full_location = backend

            # move any prepended location into the full module path
            #     e.g. 'fn=linalg.eigh' -> ['linalg', 'eigh']
            split_fn = fn.split(".")
            full_location = ".".join([full_location] + split_fn[:-1])
            only_fn = split_fn[-1]

            # try aliases for global (not function specific) modules and
            # submodules:
            #     e.g. 'decimal' -> 'math'
            #     e.g. 'cupy.scipy' -> 'cupyx.scipy'
            # we don't do this if the function location has been explicitly
            # give in _SUBMODULE_ALIASES, as that is already a full path
            for k, v in _MODULE_ALIASES.items():
                if full_location[: len(k)] == k:
                    full_location = full_location.replace(k, v, 1)
                    break

        # cached lookup of custom name function might take
        #     e.g. ['tensorflow', 'sum'] -> 'reduce_sum'
        fn_name = _FUNC_ALIASES.get((backend, fn), only_fn)

        # import the function into the cache
        try:
            lib = importlib.import_module(full_location)
        except ImportError:
            if "." in full_location:
                # sometimes libraries hack an attribute to look like submodule
                mod, *submods = full_location.split(".")
                lib = importlib.import_module(mod)
                # also need to handle nested submodules
                for submod in submods:
                    lib = getattr(lib, submod)
            else:
                # failed to import library at all -> catch + raise ImportError
                raise AttributeError

        # check for a custom wrapper but default to identity
        wrapper = _CUSTOM_WRAPPERS.get((backend, fn), lambda fn: fn)

        # store the function!
        lib_fn = _FUNCS[backend, fn] = wrapper(getattr(lib, fn_name))

    except AttributeError:
        # check if there is a backup function (e.g. for older library version)
        backend_alt = backend + "[alt]"
        if backend_alt in _MODULE_ALIASES:
            return import_lib_fn(backend_alt, fn)

        raise ImportError(
            f"autoray couldn't find function '{fn}' for "
            f"backend '{backend.replace('[alt]', '')}'."
        )

    return lib_fn


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

    Returns
    -------
    callable
    """

    try:
        lib_fn = _FUNCS[backend, fn]
    except KeyError:
        lib_fn = import_lib_fn(backend, fn)
    return lib_fn


# --------------------------- register your own! ---------------------------- #


def register_backend_alias(alias, backend):
    """Register an alias for a backend, i.e. if the backend `alias` is
    inferred for an array, use functions from `backend` instead.

    Parameters
    ----------
    alias : str
        The alias name.
    backend : str
        The actual backend name.
    """
    _BACKEND_ALIASES[alias] = backend


def register_module_alias(alias, module):
    """Register an alias for a module.

    Parameters
    ----------
    alias : str
        The alias name.
    module : str
        The actual module name.
    """
    _MODULE_ALIASES[alias] = module


def register_submodule_alias(backend, fn, module):
    """Register an alias for a submodule location of a function.

    Parameters
    ----------
    backend : str
        The name of the backend.
    fn : str
        The name of the function.
    module : str
        The module where the function is located.
    """
    _SUBMODULE_ALIASES[backend, fn] = module


def register_func_alias(backend, fn, alias):
    """Register an alias for a function name.

    Parameters
    ----------
    backend : str
        The name of the backend.
    fn : str
        The name of the function.
    alias : str
        The name of the function in the backend.
    """
    _FUNC_ALIASES[backend, fn] = alias


def register_custom_wrapper(backend, fn, wrapper=None):
    """Register a custom wrapper for a function. The wrapper is called lazily
    so that no imports are done until the function is actually used.

    Parameters
    ----------
    backend : str
        The name of the backend.
    fn : str
        The name of the function.
    wrapper : callable, optional
        The wrapper function. It should take a function as input and return a
        function. If not supplied, this function can be used as a decorator
        with ``backend`` and ``fn`` only.
    """
    if wrapper is None:

        def decorator(wrapper):
            register_custom_wrapper(backend, fn, wrapper)
            return wrapper

        return decorator

    _CUSTOM_WRAPPERS[backend, fn] = wrapper


def register_function(backend, name, fn=None, wrap=False):
    """Directly provide a custom implementation.

    Parameters
    ----------
    backend : str
        The name of the backend to register the function for.
    name : str
        Name of the function, e.g. `'sum'` or `'linalg.svd'`.
    fn : callable, optional
        The function to register. If not supplied, this function can be used
        as a decorator with ``backend`` and ``name`` only.
    wrap : bool, optional
        Whether to wrap the old function like ``fn(old_fn)`` rather than
        directly supply the entire new function. This wrapper is eagerly called
        when registering, unlike `register_custom_wrapper`.
    """
    if fn is None:

        def decorator(fn):
            register_function(backend, name, fn, wrap=wrap)
            return fn

        return decorator

    if wrap:
        old = get_lib_fn(backend, name)
        _FUNCS[backend, name] = fn(old)
    else:
        _FUNCS[backend, name] = fn


# -------------------------- array detection utils -------------------------- #


def is_array(x):
    """Is ``x`` an array-like object? This simply checks for a ``shape``
    attribute, thus 0-dimensional arrays are also considered arrays, but lists
    and tuples are not.

    Parameters
    ----------
    x : object
        Object to check.

    Returns
    -------
    bool

    See Also
    --------
    is_scalar
    """
    return hasattr(x, "shape")


_IS_SCALAR_CACHE = {}


def _is_scalar_ndim(x):
    return x.ndim == 0


def _always_false(x):
    return False


def _always_true(x):
    return True


def is_scalar(x):
    """Is ``x`` a scalar-like object? This checks if ``x`` has an ``ndim``
    attribute equal to 0. If ``x`` has no ``ndim`` attribute, it checks if
    ``x`` is iterable - if it is not iterable, it is considered a scalar.

    Parameters
    ----------
    x : object
        Object to check.

    Returns
    -------
    bool

    See Also
    --------
    is_array
    """
    try:
        return _IS_SCALAR_CACHE[x.__class__](x)
    except KeyError:
        if hasattr(x, "ndim"):
            _IS_SCALAR_CACHE[x.__class__] = _is_scalar_ndim
        else:
            try:
                iter(x)
                _IS_SCALAR_CACHE[x.__class__] = _always_false
            except TypeError:
                _IS_SCALAR_CACHE[x.__class__] = _always_true
        return _IS_SCALAR_CACHE[x.__class__](x)


# ------------------------------- tree utils -------------------------------- #

TREE_MAP_REGISTRY = {}
TREE_APPLY_REGISTRY = {}
TREE_ITER_REGISTRY = {}


def 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``.
    """
    TREE_MAP_REGISTRY[cls] = mapper
    TREE_ITER_REGISTRY[cls] = iterator
    TREE_APPLY_REGISTRY[cls] = applier


IS_CONTAINER_CACHE = {}


def 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.
    """
    try:
        return IS_CONTAINER_CACHE[x.__class__]
    except KeyError:
        isleaf = not any(isinstance(x, cls) for cls in TREE_MAP_REGISTRY)
        IS_CONTAINER_CACHE[x.__class__] = isleaf
        return isleaf


def identity(f, tree, is_leaf):
    return tree


TREE_MAPPER_CACHE = {}


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

    Returns
    -------
    pytree
    """
    if is_leaf(tree):
        return f(tree)

    try:
        return TREE_MAPPER_CACHE[tree.__class__](f, tree, is_leaf)
    except KeyError:
        # reverse so later registered classes take precedence
        for cls, mapper in reversed(TREE_MAP_REGISTRY.items()):
            if isinstance(tree, cls):
                break
        else:
            # neither leaf nor container -> simply return it
            mapper = identity
        TREE_MAPPER_CACHE[tree.__class__] = mapper
        return mapper(f, tree, is_leaf)


def empty(tree, is_leaf):
    return iter(())


TREE_ITER_CACHE = {}


def 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``.
    """
    if is_leaf(tree):
        yield tree
        return

    try:
        yield from TREE_ITER_CACHE[tree.__class__](tree, is_leaf)
    except KeyError:
        # reverse so later registered classes take precedence
        for cls, iterator in reversed(TREE_ITER_REGISTRY.items()):
            if isinstance(tree, cls):
                break
        else:
            # neither leaf nor container -> simply ignore it
            iterator = empty
        TREE_ITER_CACHE[tree.__class__] = iterator
        yield from iterator(tree, is_leaf)


def nothing(f, tree, is_leaf):
    pass


TREE_APPLIER_CACHE = {}


def 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``.
    """
    if is_leaf(tree):
        f(tree)
        return

    try:
        TREE_APPLIER_CACHE[tree.__class__](f, tree, is_leaf)
    except KeyError:
        # reverse so later registered classes take precedence
        for cls, applier in reversed(TREE_APPLY_REGISTRY.items()):
            if isinstance(tree, cls):
                break
        else:
            # neither leaf nor container -> simply ignore it
            applier = nothing
        TREE_APPLIER_CACHE[tree.__class__] = applier
        applier(f, tree, is_leaf)


class Leaf:
    """A singleton object to use as a placeholder in a pytree, for
    unflattening.
    """

    __slots__ = ()

    def __repr__(self):
        return "Leaf"


LEAF = Leaf()


def is_leaf_placeholder(x):
    # don't do `x is LEAF` to allow pickling / unpickling
    return x.__class__ is Leaf


def 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.
    """
    objs = []
    if get_ref:
        # return a new tree with Leaf leaves, as well as the flattened list

        def f(x):
            objs.append(x)
            return LEAF

        ref_tree = tree_map(f, tree, is_leaf)
        return objs, ref_tree
    else:
        tree_apply(objs.append, tree, is_leaf)
        return objs


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

    Returns
    -------
    pytree
    """
    objs = iter(objs)
    return tree_map(lambda _: next(objs), tree, is_leaf)


def tree_map_tuple(f, tree, is_leaf):
    return tuple(tree_map(f, x, is_leaf) for x in tree)


def tree_iter_tuple(tree, is_leaf):
    for x in tree:
        yield from tree_iter(x, is_leaf)


def tree_apply_tuple(f, tree, is_leaf):
    for x in tree:
        tree_apply(f, x, is_leaf)


tree_register_container(
    tuple, tree_map_tuple, tree_iter_tuple, tree_apply_tuple
)


def tree_map_list(f, tree, is_leaf):
    return [tree_map(f, x, is_leaf) for x in tree]


def tree_iter_list(tree, is_leaf):
    for x in tree:
        yield from tree_iter(x, is_leaf)


def tree_apply_list(f, tree, is_leaf):
    for x in tree:
        tree_apply(f, x, is_leaf)


tree_register_container(list, tree_map_list, tree_iter_list, tree_apply_list)


def tree_map_dict(f, tree, is_leaf):
    return {k: tree_map(f, v, is_leaf) for k, v in tree.items()}


def tree_iter_dict(tree, is_leaf):
    for v in tree.values():
        yield from tree_iter(v, is_leaf)


def tree_apply_dict(f, tree, is_leaf):
    for v in tree.values():
        tree_apply(f, v, is_leaf)


tree_register_container(dict, tree_map_dict, tree_iter_dict, tree_apply_dict)


# --------------------------- composed functions ---------------------------- #


class Composed:
    """Compose an ``autoray.do`` using function. See the main wrapper
    ``compose``.
    """

    def __init__(self, fn, name=None):
        self._default_fn = fn
        if name is None:
            name = fn.__name__
        self._name = name
        self._supply_backend = "backend" in signature(fn).parameters

        # this registers the fact that when `get_lib_fn` is called, the
        # function can be created even if it doesn't exist for a specific
        # backend yet.
        _COMPOSED_FUNCTION_GENERATORS[self._name] = self.make_function

    def register(self, backend, fn=None):
        """Register a different implementation for ``backend``."""
        if fn is not None:
            register_function(backend, self._name, fn)
        else:
            # wrapper form
            def wrapper(fn):
                register_function(backend, self._name, fn)
                return fn

            return wrapper

    def make_function(self, backend):
        """Make a new function for the specific ``backend``."""
        if self._supply_backend:
            # make sure it inherits __name__ etc
            fn = functools.wraps(self._default_fn)(
                functools.partial(self._default_fn, backend=backend)
            )
        else:
            fn = self._default_fn
        self.register(backend, fn)
        return fn

    def __call__(self, *args, like=None, **kwargs):
        backend = _choose_backend(self._name, args, kwargs, like=like)
        # `get_lib_fn` will call `make_function` if the function doesn't exist
        fn = get_lib_fn(backend, self._name)
        return fn(*args, **kwargs)

    def __repr__(self):
        return f"Composed('{self._name}')"


def 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.
    """
    if fn is None:
        return functools.partial(compose, name=name)
    return functools.wraps(fn)(Composed(fn, name))


# ---------------------- special top level functions ------------------------ #


@compose
def 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 : tuple of int
        The size of each dimension of the array.
    """
    try:
        return x.shape
    except AttributeError:
        # want to handle builtins / nested stuff
        if isinstance(x, (list, tuple)):
            d = len(x)
            if d != 0:
                # NB: slightly different from np.shape, as we do not check for
                # ragged arrays, but that behavior is seemingly deprecated
                return (d,) + shape(x[0])
            return (d,)
        return ()


@compose
def 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 : int
    """
    try:
        return x.ndim
    except AttributeError:
        return len(shape(x))


@compose
def 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 : int
    """
    try:
        return x.size
    except AttributeError:
        return math.prod(shape(x))


def conj(x):
    """Array conjugate."""
    return do("conj", x)


def transpose(x, *args):
    """Array transpose."""
    return do("transpose", x, *args)


def dag(x):
    """Array Hermitian transpose."""
    try:
        return x.H
    except AttributeError:
        backend = _infer_class_backend_cached(x.__class__)
        return do("conj", do("transpose", x, like=backend), like=backend)


def real(x):
    """Array real part."""
    return do("real", x)


def imag(x):
    """Array imaginary part."""
    return do("imag", x)


def reshape(x, shape):
    """Array reshaped."""
    try:
        return x.reshape(shape)
    except AttributeError:
        return do("reshape", x, shape)


def to_backend_dtype(dtype_name, like):
    """Turn string specifier ``dtype_name`` into dtype of backend ``like``."""
    if not isinstance(like, str):
        like = _infer_class_backend_cached(like.__class__)

    try:
        return get_lib_fn(like, dtype_name)
    except ImportError:
        return dtype_name


@compose
def get_dtype_name(x):
    """Find string specifier ``dtype_name`` of array ``x``."""
    dtype = x.dtype
    try:
        return dtype.name
    except AttributeError:
        return str(dtype)


_COMPLEX_DTYPES = {"complex64", "complex128"}
_DOUBLE_DTYPES = {"float64", "complex128"}
_DTYPE_MAP = {
    (False, False): "float32",
    (False, True): "float64",
    (True, False): "complex64",
    (True, True): "complex128",
}


def get_common_dtype(*arrays):
    """Compute the minimal dtype sufficient for ``arrays``."""
    dtypes = set(map(get_dtype_name, arrays))
    has_complex = not _COMPLEX_DTYPES.isdisjoint(dtypes)
    has_double = not _DOUBLE_DTYPES.isdisjoint(dtypes)
    return _DTYPE_MAP[has_complex, has_double]


def astype(x, dtype_name, **kwargs):
    """Cast array as type ``dtype_name`` - tries ``x.astype`` first."""
    dtype = to_backend_dtype(dtype_name, like=x)
    try:
        return x.astype(dtype, **kwargs)
    except AttributeError:
        return do("astype", x, dtype, **kwargs)


def to_numpy(x):
    """Get a numpy version of array ``x``."""
    return do("to_numpy", x)


# -------------------------- some common wrappers --------------------------- #


def svd_not_full_matrices_wrapper(fn):
    @functools.wraps(fn)
    def default_not_full_matrices(*args, **kwargs):
        kwargs.setdefault("full_matrices", False)
        return fn(*args, **kwargs)

    return default_not_full_matrices


def svd_sUV_to_UsVH_wrapper(fn):
    @functools.wraps(fn)
    def numpy_like(*args, **kwargs):
        s, U, V = fn(*args, **kwargs)
        xp = get_namespace(V)
        return U, s, xp.conj(xp.swapaxes(V, -2, -1))

    return numpy_like


def svd_UsV_to_UsVH_wrapper(fn):
    @functools.wraps(fn)
    def numpy_like(*args, **kwargs):
        U, s, V = fn(*args, **kwargs)
        xp = get_namespace(V)
        return U, s, xp.conj(xp.swapaxes(V, -2, -1))

    return numpy_like


def svd_manual_full_matrices_kwarg(fn):
    @functools.wraps(fn)
    def numpy_like(*args, full_matrices=False, **kwargs):
        U, s, VH = fn(*args, **kwargs)

        if not full_matrices:
            U, VH = U[:, : s.size], VH[: s.size, :]

        return U, s, VH

    return numpy_like


def qr_allow_fat(fn):
    @functools.wraps(fn)
    def numpy_like(a, **kwargs):
        xp = get_namespace(a)

        *_, m, n = xp.shape(a)

        if m >= n:
            # square or thin
            return fn(a, **kwargs)

        Q, R_sq = fn(a[..., :, :m])
        Qdag = xp.conj(xp.swapaxes(Q, -2, -1))
        R_r = Qdag @ a[..., :, m:]
        R = xp.concatenate((R_sq, R_r), axis=-1)

        return Q, R

    return numpy_like


def tril_to_band_part(fn):
    @functools.wraps(fn)
    def numpy_like(x, k=0):
        if k < 0:
            raise ValueError(
                "'k' must be positive to recreate 'numpy.tril' "
                "behaviour with 'tensorflow.matrix_band_part'."
            )

        return fn(x, -1, k)

    return numpy_like


def triu_to_band_part(fn):
    @functools.wraps(fn)
    def numpy_like(x, k=0):
        if k > 0:
            raise ValueError(
                "'k' must be negative to recreate 'numpy.triu' "
                "behaviour with 'tensorflow.matrix_band_part'."
            )

        return fn(x, -k, -1)

    return numpy_like


def cholesky_lower(fn):
    @functools.wraps(fn)
    def cholesky_numpy_like(a):
        return fn(a, lower=True)

    return cholesky_numpy_like


def binary_allow_1d_rhs_wrap(fn):
    @functools.wraps(fn)
    def allow_1d_rhs(a, b):
        need_to_convert = ndim(a) != ndim(b)
        if need_to_convert:
            b = reshape(b, (*shape(b), 1))
        x = fn(a, b)
        if need_to_convert:
            x = reshape(x, shape(x)[:-1])
        return x

    return allow_1d_rhs


def scale_random_uniform_manually(fn):
    @functools.wraps(fn)
    def numpy_like(low=0.0, high=1.0, size=None, dtype=None, **kwargs):
        if size is None:
            size = ()

        x = fn(size, **kwargs)

        if (low != 0.0) or (high != 1.0):
            x = (high - low) * x + low

        if (dtype is not None) and get_dtype_name(x) != dtype:
            x = astype(x, dtype)
        return x

    return numpy_like


def scale_random_normal_manually(fn):
    @functools.wraps(fn)
    def numpy_like(loc=0.0, scale=1.0, size=None, dtype=None, **kwargs):
        if size is None:
            size = ()

        x = fn(size, **kwargs)

        if (loc != 0.0) or (scale != 1.0):
            x = scale * x + loc

        if (dtype is not None) and get_dtype_name(x) != dtype:
            x = astype(x, dtype)
        return x

    return numpy_like


def 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.
    """

    @functools.wraps(fn)
    def with_dtype(*args, dtype=None, **kwargs):
        A = fn(*args, **kwargs)
        if (dtype is not None) and (dtype != get_dtype_name(A)):
            A = astype(A, dtype)
        return A

    return with_dtype


def 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.
    """

    @functools.wraps(fn)
    def translated_function(*args, **kwargs):
        new_kwargs = {}
        translation = translator.copy()

        # convert args, pairing them off with kwargs
        for arg_value in args:
            new_arg_name = translation.popitem(last=False)[1][0]
            new_kwargs[new_arg_name] = arg_value

        # convert kwargs - but only those in the translation
        for key, value in kwargs.items():
            try:
                new_kwargs[translation.pop(key)[0]] = value
            except KeyError:
                new_kwargs[key] = value

        # set remaining default kwargs
        for opt in translation.values():
            if len(opt) == 2:
                # backend_name, default_value
                new_kwargs[opt[0]] = opt[1]
            # else, no default value -> don't inject

        return fn(**new_kwargs)

    return translated_function


def make_translator(t):
    return functools.partial(translate_wrapper, translator=OrderedDict(t))


def complex_add_re_im(re, im):
    return re + 1j * im


def allclose(x, y, rtol=1e-05, atol=1e-08):
    return do("all", do("abs", x - y) <= atol + rtol * do("abs", y))


def _handle_size_to_shape(size=None):
    if size is None:
        return ()
    try:
        return tuple(size)
    except TypeError:
        return (size,)


# ----------------------------- Custom dispatchers -------------------------- #


def 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.
    """

    @functools.wraps(fn)
    def wrapped(args, kwargs):
        return fn(*args, **kwargs)

    return wrapped


def 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)``.
    """
    if raw_signature:
        dispatcher = wrap_args_kwargs_from_raw(dispatcher)

    _DISPATCHERS[fun] = dispatcher


def default_dispatcher(args, kwargs):
    """Try to infer backend from first argument passed to function."""
    try:
        return _infer_class_backend_cached(args[0].__class__)
    except IndexError:
        raise TypeError("No args to infer backend from.")


# lookup of custom dispatcher methods, for cases when backend cannot be
#     inferred accurately from first argument.
_DISPATCHERS = defaultdict(lambda: default_dispatcher)


def join_array_dispatcher(args, kwargs):
    """Dispatcher for functions where first argument is a sequence."""
    try:
        return _infer_class_backend_cached(args[0][0].__class__)
    except (TypeError, ValueError):
        # user passed an empty sequence, or something non-iterable
        # try to infer backend from first argument as fallback
        return _infer_class_backend_cached(args[0].__class__)


# List of functions listed in numpy API as array joining operations
register_dispatch("concatenate", join_array_dispatcher, raw_signature=False)
register_dispatch("stack", join_array_dispatcher, raw_signature=False)
register_dispatch("block", join_array_dispatcher, raw_signature=False)
register_dispatch("vstack", join_array_dispatcher, raw_signature=False)
register_dispatch("hstack", join_array_dispatcher, raw_signature=False)
register_dispatch("dstack", join_array_dispatcher, raw_signature=False)
register_dispatch("column_stack", join_array_dispatcher, raw_signature=False)
register_dispatch("row_stack", join_array_dispatcher, raw_signature=False)


def 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.
    """
    return infer_backend_multi(*args)


register_dispatch("einsum", einsum_dispatcher, raw_signature=False)


def 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.
    """
    return infer_backend_multi(*args[:2])


register_dispatch("tensordot", binary_dispatcher, raw_signature=False)
register_dispatch("matmul", binary_dispatcher, raw_signature=False)
register_dispatch("multiply", binary_dispatcher, raw_signature=False)
register_dispatch("divide", binary_dispatcher, raw_signature=False)
register_dispatch("true_divide", binary_dispatcher, raw_signature=False)
register_dispatch("add", binary_dispatcher, raw_signature=False)
register_dispatch("subtract", binary_dispatcher, raw_signature=False)

# TODO: register other binary functions?

# --------------- object to act as drop-in replace for numpy ---------------- #


class InjectDtypeDevice:
    """Wrapper that injects defaultdtype and device arguments into a function"""

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

    def __init__(self, fn, device=None, dtype=None):
        self._fn = fn
        self._device = device
        self._dtype = dtype

    def __call__(self, *args, **kwargs):
        if self._device is not None and "device" not in kwargs:
            kwargs["device"] = self._device
        if self._dtype is not None and "dtype" not in kwargs:
            kwargs["dtype"] = self._dtype
        return self._fn(*args, **kwargs)

    def __repr__(self):
        return (
            f"InjectDtypeDevice(fn={self._fn}, "
            f"device={self._device}, "
            f"dtype={self._dtype})"
        )


_NAME_SPACE_SUBMODULES = {
    "random",
    "linalg",
    "scipy",
    "scipy.linalg",
    "scipy.sparse.linalg",
}


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

    Parameters
    ----------
    like : array_like, str, or None
        The backend to use, or an object to infer the backend from. If None,
        the default behavior is to use `autoray.do` and auto dispatch backend
        at function call time. If given, the functions are cached at first
        call.
    device : str, optional
        The device to use for the backend. If None, it will be inferred from
        the `like` paramater is that is array-like or set to None.
    dtype : str, optional
        The dtype to use for the backend. If None, it will be inferred from
        the `like` parameter if that is array-like or set to None.
    submodule : str, optional
        This is used internally when nesting attribute lookups, e.g.
        `xp.random.normal`, `xp.linalg.eigh`.
    """

    def __init__(
        self,
        like=None,
        device=None,
        dtype=None,
        submodule=None,
    ):
        if like is None:
            # use autoray.do and auto dispatch
            self._backend = None
        elif isinstance(like, str):
            # commit to a specific given backend
            self._backend = like
        else:
            # commit to a specific backend inferred from array-like
            self._backend = _infer_class_backend_cached(like.__class__)

        if device is None:
            if like is not None:
                if hasattr(like, "device"):
                    device = like.device
        self._device = device

        if dtype is None:
            if like is not None:
                if hasattr(like, "dtype"):
                    dtype = like.dtype
        self._dtype = dtype
        self._submodule = submodule

    def _get_submodule(self, name):
        new = object.__new__(type(self))
        new._backend = self._backend
        new._device = self._device
        new._dtype = self._dtype
        new._submodule = name
        return new

    def _get_fn(self, name):
        if name.startswith("__") and name.endswith("__"):
            # raise correct error for dunder methods, so
            # that debuggers and inspectors work correctly
            raise AttributeError(
                f"'{self.__class__.__name__}' object has no attribute '{name}'"
            )

        if self._submodule is not None:
            # prepend the submodule name
            name = f"{self._submodule}.{name}"

        if name in _NAME_SPACE_SUBMODULES:
            # note that other submodules can be accessed, these are just the
            # ones with functions that we possibly want to translate
            return self._get_submodule(name)

        if self._backend is None:
            # use autoray.do and auto dispatch
            return functools.partial(do, name)

        fn = get_lib_fn(self._backend, name)

        # possibly wrap for dtype and device injection
        if name in _CREATION_ROUTINES:
            key = (self._backend, name)
            try:
                # check for backend specific defaults
                inject_dtype, inject_device = _CREATION_INJECT[key]
            except KeyError:
                # populate with non backend specific defaults
                inject_dtype, inject_device = _CREATION_ROUTINES[name]
                _CREATION_INJECT[key] = (inject_dtype, inject_device)

            if not inject_dtype:
                # this is not a function that accepts dtype
                dtype_to_inject = None
            else:
                dtype_to_inject = self._dtype

            if not inject_device:
                # this is not a function that accepts device
                device_to_inject = None
            else:
                device_to_inject = self._device

            if (dtype_to_inject is not None) or (device_to_inject is not None):
                # only wrap if we actually inject something
                fn = InjectDtypeDevice(fn, device_to_inject, dtype_to_inject)

        return fn

    def __getattr__(self, name):
        x = self._get_fn(name)
        self.__dict__[name] = x
        return x

    def __repr__(self):
        return (
            f"{self.__class__.__name__}("
            f"backend={self._backend}, "
            f"device={self._device}, "
            f"dtype={self._dtype}, "
            f"submodule={self._submodule}"
            ")"
        )


_NAMESPACE_CACHE = {}


def 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
    -------
    AutoNamespace
        An automatic namespace object.
    """
    if like is not None:
        if not isinstance(like, str):
            # array
            cls = like.__class__
            if device is None:
                try:
                    device = like.device
                except AttributeError:
                    device = None
            if dtype is None:
                try:
                    dtype = like.dtype
                except AttributeError:
                    dtype = None
        else:
            # manually specified backend string
            cls = like
    else:
        # namespace functions will dispatch at call time
        cls = None

    key = (cls, device, dtype, submodule)

    try:
        xp = _NAMESPACE_CACHE[key]
    except KeyError:
        xp = _NAMESPACE_CACHE[key] = AutoNamespace(
            like=like,
            device=device,
            dtype=dtype,
            submodule=submodule,
        )

    return xp


# note this is the backend agnostic numpy mimic (`import autoray.numpy as np`)
# not the namespace for backend numpy (`xp = get_namespace(like="numpy")`)
numpy = get_namespace()


# --------------------------------------------------------------------------- #
#                             specific functions                              #
# --------------------------------------------------------------------------- #

# ------------------------------ standard-lib ------------------------------- #

register_module_alias("decimal", "math")
register_module_alias("builtins", "numpy")


_builtin_dtype_lookup = {
    int: "int64",
    float: "float64",
    complex: "complex128",
}


@get_dtype_name.register("builtins")
def builtins_get_dtype_name(x):
    return _builtin_dtype_lookup[x.__class__]


register_function("builtins", "complex", complex)

# ---------------------------------- numpy ---------------------------------- #


@register_function("numpy", "to_numpy")
@register_function("builtins", "to_numpy")
def numpy_to_numpy(x):
    return do("asarray", x, like="numpy")


register_module_alias("numpy.scipy", "scipy")

register_submodule_alias("numpy", "linalg.expm", "scipy.linalg")
register_submodule_alias("numpy", "linalg.lu", "scipy.linalg")

register_custom_wrapper("numpy", "linalg.svd", svd_not_full_matrices_wrapper)
register_custom_wrapper("numpy", "random.normal", with_dtype_wrapper)
register_custom_wrapper("numpy", "random.uniform", with_dtype_wrapper)

register_function("numpy", "complex", complex_add_re_im)

# ---------------------------------- cupy ----------------------------------- #


@register_function("cupy", "to_numpy")
def cupy_to_numpy(x):  # pragma: no cover
    return x.get()


register_module_alias("cupy.scipy", "cupyx.scipy")

register_custom_wrapper("cupy", "linalg.svd", svd_not_full_matrices_wrapper)

register_function("cupy", "complex", complex_add_re_im)

# ----------------------------------- jax ----------------------------------- #


@register_function("jax", "to_numpy")
def jax_to_numpy(x):
    return do("asarray", x, like="numpy")


class JaxDefaultRNG:
    """Stateful but deterministic random number generator for JAX following
    numpy's Generator API, compatible with `jax.jit`.
    """

    def __init__(self, seed, **kwargs):
        import jax

        self.jax = jax
        self.key = jax.random.key(seed, **kwargs)

    def binomial(self, n, p, size=None, **kwargs):
        self.key, subkey = self.jax.random.split(self.key)
        shape = _handle_size_to_shape(size)
        return self.jax.random.binomial(
            subkey, n=n, p=p, shape=shape, **kwargs
        )

    def choice(self, a, size=None, replace=True, p=None, axis=0, **kwargs):
        self.key, subkey = self.jax.random.split(self.key)
        shape = _handle_size_to_shape(size)
        return self.jax.random.choice(
            subkey,
            a,
            shape=shape,
            replace=replace,
            p=p,
            axis=axis,
            **kwargs,
        )

    def exponential(self, scale=1.0, size=None, **kwargs):
        self.key, subkey = self.jax.random.split(self.key)
        shape = _handle_size_to_shape(size)
        x = self.jax.random.exponential(subkey, shape=shape, **kwargs)
        if scale != 1.0:
            x *= scale
        return x

    def gumbel(self, loc=0.0, scale=1.0, size=None, **kwargs):
        self.key, subkey = self.jax.random.split(self.key)
        shape = _handle_size_to_shape(size)
        x = self.jax.random.gumbel(subkey, shape=shape, **kwargs)
        if scale != 1.0:
            x *= scale
        if loc != 0.0:
            x += loc
        return x

    def integers(self, low, high=None, size=None, **kwargs):
        self.key, subkey = self.jax.random.split(self.key)
        shape = _handle_size_to_shape(size)
        if high is None:
            high = low
            low = 0
        return self.jax.random.randint(
            subkey,
            shape=shape,
            minval=low,
            maxval=high,
            **kwargs,
        )

    def normal(self, loc=0.0, scale=1.0, size=None, **kwargs):
        self.key, subkey = self.jax.random.split(self.key)
        shape = _handle_size_to_shape(size)
        x = self.jax.random.normal(subkey, shape=shape, **kwargs)
        if scale != 1.0:
            x *= scale
        if loc != 0.0:
            x += loc
        return x

    def permutation(self, x, **kwargs):
        self.key, subkey = self.jax.random.split(self.key)
        return self.jax.random.permutation(subkey, x, **kwargs)

    def poisson(self, lam=1.0, size=None, **kwargs):
        self.key, subkey = self.jax.random.split(self.key)
        shape = _handle_size_to_shape(size)
        x = self.jax.random.poisson(subkey, lam, shape=shape, **kwargs)
        return x

    def random(self, size=None, **kwargs):
        return self.uniform(size=size, **kwargs)

    def uniform(self, low=0.0, high=1.0, size=None, **kwargs):
        self.key, subkey = self.jax.random.split(self.key)
        shape = _handle_size_to_shape(size)
        return self.jax.random.uniform(
            subkey, shape=shape, minval=low, maxval=high, **kwargs
        )


@register_function("jax", "random.default_rng")
def jax_default_rng(seed, **kwargs):
    if isinstance(seed, JaxDefaultRNG):
        return seed
    return JaxDefaultRNG(seed, **kwargs)


_JAX_RANDOM_KEY = None


@register_function("jax", "random.seed")
def jax_random_seed(seed=None):
    from jax.random import PRNGKey

    global _JAX_RANDOM_KEY
    if seed is None:
        from random import SystemRandom

        seed = SystemRandom().randint(-(2**63), 2**63 - 1)  # inclusive high
    _JAX_RANDOM_KEY = PRNGKey(seed)


def jax_random_get_key():
    from jax.random import split

    global _JAX_RANDOM_KEY
    if _JAX_RANDOM_KEY is None:
        jax_random_seed()
    _JAX_RANDOM_KEY, subkey = split(_JAX_RANDOM_KEY)
    return subkey


@register_function("jax", "random.uniform")
def jax_random_uniform(low=0.0, high=1.0, size=None, **kwargs):
    from jax.random import uniform

    if size is None:
        size = ()
    return uniform(
        jax_random_get_key(), shape=size, minval=low, maxval=high, **kwargs
    )


@register_function("jax", "random.normal")
def jax_random_normal(loc=0.0, scale=1.0, size=None, **kwargs):
    from jax.random import normal

    if size is None:
        size = ()
    x = normal(jax_random_get_key(), shape=size, **kwargs)
    if scale != 1.0:
        x *= scale
    if loc != 0.0:
        x += loc
    return x


register_backend_alias("jaxlib", "jax")

register_backend(JaxDefaultRNG, "jax")

register_module_alias("jax.scipy", "jax.scipy")
register_module_alias("jax", "jax.numpy")

register_submodule_alias("jax", "complex", "jax.lax")
register_submodule_alias("jax", "linalg.expm", "jax.scipy.linalg")
register_submodule_alias("jax", "linalg.householder_product", "jax.lax.linalg")

# n.b. jax supports fat QR but *not* when computing gradients
#     https://github.com/jax-ml/jax/issues/23533
register_custom_wrapper("jax", "linalg.qr", qr_allow_fat)
register_custom_wrapper("jax", "linalg.svd", svd_not_full_matrices_wrapper)

# --------------------------------- aesara ---------------------------------- #


register_module_alias("aesara", "aesara.tensor")


@shape.register("aesara")
@register_function("aesara", "shape")
def aesara_shape(x):
    return x.type.shape


# -------------------------------- autograd --------------------------------- #


register_module_alias("autograd", "autograd.numpy")

register_custom_wrapper(
    "autograd", "linalg.svd", svd_not_full_matrices_wrapper
)

register_function("autograd", "complex", complex_add_re_im)


@register_function("autograd", "take")
def autograd_take(x, indices, axis=None):
    # NOTE: autograd take doesn't support grad
    if axis is None:
        return x.ravel()[indices]
    else:
        selector = tuple(
            slice(None) if i != axis else indices for i in range(x.ndim)
        )
        return x[selector]


# ---------------------------------- dask ----------------------------------- #


def dask_to_numpy(x):
    return x.compute()


@register_custom_wrapper("dask", "eye")
def dask_eye_wrapper(eye_fn):
    # Make M work as positional argument
    @functools.wraps(eye_fn)
    def numpy_like(N, M=None, **kwargs):
        return eye_fn(N, M=M, **kwargs)

    return numpy_like


register_module_alias("dask", "dask.array")

register_func_alias("dask", "abs", "absolute")
register_func_alias("dask", "identity", "eye")

register_custom_wrapper("dask", "linalg.cholesky", cholesky_lower)
register_custom_wrapper("dask", "linalg.svd", svd_manual_full_matrices_kwarg)
register_custom_wrapper("dask", "random.normal", with_dtype_wrapper)
register_custom_wrapper("dask", "random.uniform", with_dtype_wrapper)

register_function("dask", "complex", complex_add_re_im)
register_function("dask", "to_numpy", dask_to_numpy)

# ---------------------------------- mars ----------------------------------- #


def mars_to_numpy(x):
    return x.to_numpy()


register_module_alias("mars", "mars.tensor")

register_custom_wrapper("mars", "linalg.cholesky", cholesky_lower)

register_function("mars", "complex", complex_add_re_im)
register_function("mars", "to_numpy", mars_to_numpy)


# ----------------------------------- ctf ----------------------------------- #


def ctf_array(x):
    return do("astensor", x, like="ctf")


def ctf_to_numpy(x):
    return x.to_nparray()


def ctf_count_nonzero(x):
    return (x != 0).astype(int).sum()


@get_dtype_name.register("ctf")
def ctf_get_dtype_name(x):
    return x.dtype.__name__


register_submodule_alias("ctf", "complex128", "numpy")
register_submodule_alias("ctf", "complex64", "numpy")
register_submodule_alias("ctf", "float32", "numpy")
register_submodule_alias("ctf", "float64", "numpy")
register_submodule_alias("ctf", "linalg.eigh", "ctf")
register_submodule_alias("ctf", "linalg.norm", "ctf")
register_submodule_alias("ctf", "linalg.qr", "ctf")
register_submodule_alias("ctf", "linalg.svd", "ctf")

register_func_alias("ctf", "random.uniform", "random")

register_custom_wrapper("ctf", "random.uniform", scale_random_uniform_manually)

register_function("ctf", "allclose", allclose)
register_function("ctf", "array", ctf_array)
register_function("ctf", "complex", complex_add_re_im)
register_function("ctf", "count_nonzero", ctf_count_nonzero)
register_function("ctf", "to_numpy", ctf_to_numpy)

# ------------------------------- sparse------------------------------------- #


# sparse uses numpys __array_func__ interface
for f in (
    "arccos",
    "arccosh",
    "arcsin",
    "arcsinh",
    "arctan",
    "arctanh",
    "cos",
    "cosh",
    "exp",
    "log",
    "log10",
    "log2",
    "sign",
    "sin",
    "sinh",
    "sqrt",
    "tan",
    "tanh",
    "tensordot",
    # NB put tensordot here, as sparse.tensordot can produce dense (numpy)
    # arrays but errors when both inputs are dense - we want nested calls to
    # tensordot to handle this
):
    register_submodule_alias("sparse", f, "numpy")

register_func_alias("sparse", "identity", "eye")


@register_function("sparse", "array")
def sparse_array(x):
    return do("COO.from_numpy", x, like="sparse")


@register_function("sparse", "to_numpy")
def sparse_to_numpy(x):
    return x.todense()


@register_function("sparse", "transpose")
def sparse_transpose(x, axes=None):
    return x.transpose(axes)


@register_function("sparse", "reshape")
def sparse_reshape(x, shape):
    return x.reshape(shape)


@register_function("sparse", "sum")
def sparse_sum(x, axis=None, keepdims=False, dtype=None, out=None):
    return x.sum(axis=axis, keepdims=keepdims, dtype=dtype, out=out)


@register_function("sparse", "prod")
def sparse_prod(x, axis=None, keepdims=False, dtype=None, out=None):
    return x.prod(axis=axis, keepdims=keepdims, dtype=dtype, out=out)


@register_function("sparse", "conj")
def sparse_conj(x):
    return x.conj()


@register_function("sparse", "real")
def sparse_real(x):
    return x.real


@register_function("sparse", "imag")
def sparse_imag(x):
    return x.imag


@register_function("sparse", "count_nonzero")
def sparse_count_nonzero(x):
    return x.nnz


@register_function("sparse", "complex")
def sparse_complex(re, im):
    return complex_add_re_im(re, im)


@register_function("sparse", "random.uniform")
def sparse_random_uniform(low=0.0, high=1.0, size=None, dtype=None, **kwargs):
    def rvs(nnz):
        return do(
            "random.uniform", low, high, (nnz,), dtype=dtype, like="numpy"
        )

    return do("random", size, data_rvs=rvs, **kwargs, like="sparse")


@register_function("sparse", "random.normal")
def sparse_random_normal(loc=0.0, scale=1.0, size=None, dtype=None, **kwargs):
    def rvs(nnz):
        return do(
            "random.normal", loc, scale, (nnz,), dtype=dtype, like="numpy"
        )

    return do("random", size, data_rvs=rvs, **kwargs, like="sparse")


# ------------------------------- tensorflow -------------------------------- #


@functools.cache
def get_tensorflow():
    import tensorflow as tf

    return tf


@register_custom_wrapper("tensorflow", "pad")
def tensorflow_pad_wrap(tf_pad):
    def numpy_like(array, pad_width, mode="constant", constant_values=0):
        if mode != "constant":
            raise NotImplementedError

        try:
            if len(pad_width) == 1:
                pad_width = pad_width * ndim(array)
        except TypeError:
            pad_width = ((pad_width, pad_width),) * ndim(array)

        return tf_pad(
            array, pad_width, mode="CONSTANT", constant_values=constant_values
        )

    return numpy_like


@register_custom_wrapper("tensorflow", "linalg.norm")
def tensorflow_wrap_norm(tf_norm):
    def wrapped_norm(x, ord=None, axis=None, keepdims=False, **kwargs):
        if ord is None:
            ord = "euclidean"
        return tf_norm(x, ord=ord, axis=axis, keepdims=keepdims, **kwargs)

    return wrapped_norm


class TensorflowDefaultRNG:
    """Stateful random number generator for TensorFlow following numpy's
    Generator API, compatible with `tf.function`.
    """

    def __init__(self, seed=None, **kwargs):
        import tensorflow as tf  # type: ignore

        self.tf = tf

        if seed is None:
            self.generator = tf.random.Generator.from_non_deterministic_state(
                **kwargs
            )
        else:
            self.generator = tf.random.Generator.from_seed(seed, **kwargs)

    def uniform(self, low=0.0, high=1.0, size=None, **kwargs):
        shape = _handle_size_to_shape(size)
        x = self.generator.uniform(
            shape=shape, minval=low, maxval=high, **kwargs
        )
        return x

    def random(self, size=None, **kwargs):
        return self.uniform(size=size, **kwargs)

    def normal(self, loc=0.0, scale=1.0, size=None, **kwargs):
        shape = _handle_size_to_shape(size)
        x = self.generator.normal(
            shape=shape, mean=loc, stddev=scale, **kwargs
        )
        return x

    def integers(self, low, high=None, size=None, **kwargs):
        shape = _handle_size_to_shape(size)
        if high is None:
            high = low
            low = 0
        return self.generator.uniform(
            shape=shape,
            minval=low,
            maxval=high,
            dtype=self.tf.int32,
            **kwargs,
        )

    # def binomial(self, n, p, size=None, **kwargs):
    #     raise NotImplementedError

    # def choice(self, a, size=None, replace=True, p=None, axis=0, **kwargs):
    #     raise NotImplementedError

    # def permutation(self, x, **kwargs):
    #     raise NotImplementedError


@register_function("tensorflow", "random.default_rng")
def tensorflow_default_rng(seed, **kwargs):
    if isinstance(seed, TensorflowDefaultRNG):
        return seed
    return TensorflowDefaultRNG(seed, **kwargs)


register_backend(TensorflowDefaultRNG, "tensorflow")

register_module_alias("tensorflow.linalg", "tensorflow.linalg")
register_module_alias("tensorflow.random", "tensorflow.random")
register_module_alias("tensorflow", "tensorflow.experimental.numpy")

# these aren't in experimental.numpy
register_submodule_alias("tensorflow", "astype", "tensorflow")
register_submodule_alias("tensorflow", "cast", "tensorflow")
register_submodule_alias("tensorflow", "complex", "tensorflow")
register_submodule_alias("tensorflow", "pad", "tensorflow")

register_func_alias("tensorflow", "astype", "cast")

register_custom_wrapper("tensorflow", "linalg.solve", binary_allow_1d_rhs_wrap)
register_custom_wrapper("tensorflow", "linalg.svd", svd_sUV_to_UsVH_wrapper)
register_custom_wrapper(
    "tensorflow",
    "random.normal",
    make_translator(
        [
            ("loc", ("mean", 0.0)),
            ("scale", ("stddev", 1.0)),
            ("size", ("shape", ())),
        ]
    ),
)
register_custom_wrapper(
    "tensorflow",
    "random.uniform",
    make_translator(
        [
            ("low", ("minval", 0.0)),
            ("high", ("maxval", 1.0)),
            ("size", ("shape", ())),
        ]
    ),
)


@register_function("tensorflow", "to_numpy")
def tensorflow_to_numpy(x):
    return x.numpy()


@register_function("tensorflow", "indices")
def tensorflow_indices(dimensions):
    _meshgrid = get_lib_fn("tensorflow", "meshgrid")
    _arange = get_lib_fn("tensorflow", "arange")
    return _meshgrid(*map(_arange, dimensions), indexing="ij")


@register_function("tensorflow", "swapaxes")
def tensorflow_swapaxes(a, axis1, axis2):
    tf = get_tensorflow()
    perm = list(range(a.ndim))
    perm[axis1], perm[axis2] = perm[axis2], perm[axis1]
    return tf.transpose(a, perm)


@register_function("tensorflow", "linalg.cholesky")
def tensorflow_cholesky(x, upper=False, **kwargs):
    tf = get_tensorflow()
    left = tf.linalg.cholesky(x, **kwargs)
    if upper:
        return tf.math.conj(tensorflow_swapaxes(left, -2, -1))
    return left


@register_function("tensorflow", "scipy.linalg.solve_triangular")
def tensorflow_solve_triangular(a, b, lower=False, **kwargs):
    tf = get_tensorflow()
    left = tf.linalg.triangular_solve(a, b, lower=lower, **kwargs)
    return left


# ---------------------------------- torch ---------------------------------- #


@functools.cache
def get_torch():
    import torch

    return torch


@shape.register("torch")
def torch_shape(x):
    # torch returns a Size object, we want tuple[int]
    return tuple(map(int, x.shape))


@size.register("torch")
def torch_size(x):
    return x.numel()


@functools.lru_cache(None)
def _torch_get_dtype_name(dtype):
    return str(dtype).split(".")[-1]


@get_dtype_name.register("torch")
def torch_get_dtype_name(x):
    return _torch_get_dtype_name(x.dtype)


@register_custom_wrapper("torch[alt]", "linalg.solve")
def torch_linalg_solve_wrap(fn):
    @binary_allow_1d_rhs_wrap
    def numpy_like(a, b):
        return fn(b, a)[0]

    return numpy_like


@register_custom_wrapper("torch", "tensordot")
def torch_tensordot_wrap(fn):
    @functools.wraps(fn)
    def numpy_like(a, b, axes=2):
        return fn(a, b, dims=axes)

    return numpy_like


@register_custom_wrapper("torch[alt]", "split")
def torch_split_wrap(fn):
    # for torch >=1.8 we can use tensor_split instead, but in current stable
    # release this function has not been added
    @functools.wraps(fn)
    def numpy_like(ary, indices_or_sections, axis=0, **kwargs):
        if isinstance(indices_or_sections, numbers.Integral):
            split_size = shape(ary)[axis] // indices_or_sections
            return fn(ary, split_size, dim=axis, **kwargs)
        else:
            # torch.split doesn't support empty splits
            if len(indices_or_sections) == 0:
                return (ary,)

            diff = do(
                "diff",
                indices_or_sections,
                prepend=0,
                append=shape(ary)[axis],
                like="numpy",
            )
            diff = list(diff)
            return fn(ary, diff, dim=axis)

    return numpy_like


@register_custom_wrapper("torch", "ones")
@register_custom_wrapper("torch", "zeros")
def torch_zeros_ones_wrap(fn):
    @functools.wraps(fn)
    def numpy_like(shape, dtype=None, **kwargs):
        if dtype is not None:
            dtype = to_backend_dtype(dtype, like="torch")
        return fn(shape, dtype=dtype, **kwargs)

    return numpy_like


@register_custom_wrapper("torch", "eye")
def torch_eye_wrap(fn):
    @functools.wraps(fn)
    def numpy_like(N, M=None, dtype=None, **kwargs):
        if dtype is not None:
            dtype = to_backend_dtype(dtype, like="torch")
        if M is not None:
            return fn(N, m=M, dtype=dtype, **kwargs)
        else:
            return fn(N, dtype=dtype, **kwargs)

    return numpy_like


@register_custom_wrapper("torch", "sort")
def torch_sort_wrap(fn):
    @functools.wraps(fn)
    def numpy_like(a, axis=-1):
        return fn(a, dim=axis)[0]

    return numpy_like


@register_custom_wrapper("torch", "flip")
def torch_flip_wrap(torch_flip):
    def numpy_like(x, axis=None):
        if axis is None:
            dims = tuple(range(x.ndimension()))
        elif isinstance(axis, int):
            dims = (axis,)
        else:
            # already tuple/list
            dims = axis
        return torch_flip(x, dims)

    return numpy_like


class TorchDefaultRNG:
    def __init__(self, seed=None, device=None):
        self._torch = get_torch()
        self._generator = self._torch.Generator(device=device)
        if seed is not None:
            self._generator.manual_seed(seed)

    # def binomial(self, n, p, size=None, **kwargs):
    #     raise NotImplementedError()

    # def choice(self, a, size=None, replace=True, p=None, axis=0, **kwargs):
    #     raise NotImplementedError()

    # def exponential(self, scale=1.0, size=None, **kwargs):
    #     raise NotImplementedError()

    # def gumbel(self, loc=0.0, scale=1.0, size=None, **kwargs):
    #     raise NotImplementedError()

    def integers(self, low, high=None, size=None, **kwargs):
        if high is None:
            high = low
            low = 0
        size = _handle_size_to_shape(size)
        return self._torch.randint(
            low, high, size, generator=self._generator, **kwargs
        )

    # def poisson(self, lam=1.0, size=None, **kwargs):
    #     raise NotImplementedError()

    def normal(self, loc=0.0, scale=1.0, size=None, **kwargs):
        size = _handle_size_to_shape(size)
        x = self._torch.randn(size, generator=self._generator, **kwargs)
        if scale != 1.0:
            x = x * scale
        if loc != 0.0:
            x = x + loc
        return x

    def random(self, size=None, **kwargs):
        size = _handle_size_to_shape(size)
        return self._torch.rand(size, generator=self._generator, **kwargs)

    def permutation(self, x, **kwargs):
        if isinstance(x, int):
            return self._torch.randperm(x, generator=self._generator, **kwargs)

        axis = kwargs.get("axis", 0)
        n = x.shape[axis]
        perm_indices = self._torch.randperm(
            n, generator=self._generator, device=x.device
        )
        return self._torch.index_select(x, axis, perm_indices)

    def uniform(self, low=0.0, high=1.0, size=None, **kwargs):
        size = _handle_size_to_shape(size)
        x = self._torch.rand(size, generator=self._generator, **kwargs)
        if low != 0.0 or high != 1.0:
            x = x * (high - low) + low
        return x

    def choice(self, a, size=None, replace=True, p=None):
        scalar = reshape = None
        if size is None:
            scalar = True
            size = 1
        elif not isinstance(size, numbers.Integral):
            # assume tuple
            reshape = size
            size = functools.reduce(operator.mul, size)

        if isinstance(a, numbers.Integral):
            a = self._torch.arange(a, device=self._generator.device)

        # with torch have to first generate indices

        if (p is not None) or (not replace):
            if p is None:
                # uniform distribution, multinomial renormalizes internally
                p = self._torch.ones(a.shape[0], device=self._generator.device)

            idx = self._torch.multinomial(
                p,
                num_samples=size,
                replacement=replace,
                generator=self._generator,
            )
        else:
            # quicker for uniform *with* replacement
            idx = self._torch.randint(
                low=0,
                high=a.shape[0],
                size=(size,),
                generator=self._generator,
            )

        # then take the samples!
        random_choices = a[idx]

        if scalar:
            random_choices = random_choices.squeeze(0)

        if reshape is not None:
            random_choices = random_choices.reshape(reshape)

        return random_choices


@register_function("torch", "random.default_rng")
def torch_default_rng(seed, **kwargs):
    if isinstance(seed, TorchDefaultRNG):
        return seed
    return TorchDefaultRNG(seed, **kwargs)


register_backend(TorchDefaultRNG, "torch")

register_submodule_alias("torch", "linalg.expm", "torch")
register_submodule_alias("torch", "scipy.linalg.expm", "torch")
register_submodule_alias("torch", "random.normal", "torch")
register_submodule_alias("torch", "random.uniform", "torch")

register_func_alias("torch", "array", "tensor")
register_func_alias("torch", "asarray", "as_tensor")
register_func_alias("torch", "clip", "clamp")
register_func_alias("torch", "concatenate", "cat")
register_func_alias("torch", "conjugate", "conj")
register_func_alias("torch", "equal", "eq")
register_func_alias("torch", "expand_dims", "unsqueeze")
register_func_alias("torch", "identity", "eye")
register_func_alias("torch", "linalg.expm", "matrix_exp")
register_func_alias("torch", "max", "amax")
register_func_alias("torch", "min", "amin")
register_func_alias("torch", "power", "pow")
register_func_alias("torch", "random.normal", "randn")
register_func_alias("torch", "random.uniform", "rand")
register_func_alias("torch", "scipy.linalg.expm", "matrix_exp")
register_func_alias("torch", "split", "tensor_split")
register_func_alias("torch", "take_along_axis", "take_along_dim")

register_custom_wrapper(
    "torch",
    "clip",
    make_translator(
        [("a", ("input",)), ("a_min", ("min",)), ("a_max", ("max",))]
    ),
)
register_custom_wrapper(
    "torch",
    "concatenate",
    make_translator([("arrays", ("tensors",)), ("axis", ("dim", 0))]),
)
register_custom_wrapper(
    "torch",
    "count_nonzero",
    make_translator([("a", ("input",)), ("axis", ("dim", None))]),
)
register_custom_wrapper(
    "torch",
    "diagonal",
    make_translator(
        [("a", ("input",)), ("axis1", ("dim1", 0)), ("axis2", ("dim2", 1))]
    ),
)
register_custom_wrapper(
    "torch", "empty", make_translator([("shape", ("size",))])
)
register_custom_wrapper(
    "torch",
    "expand_dims",
    make_translator([("a", ("input",)), ("axis", ("dim",))]),
)
register_custom_wrapper("torch", "linalg.svd", svd_not_full_matrices_wrapper)
register_custom_wrapper("torch", "random.normal", scale_random_normal_manually)
register_custom_wrapper(
    "torch", "random.uniform", scale_random_uniform_manually
)

register_custom_wrapper(
    "torch",
    "stack",
    make_translator([("arrays", ("tensors",)), ("axis", ("dim", 0))]),
)
register_custom_wrapper(
    "torch",
    "tril",
    make_translator([("m", ("input",)), ("k", ("diagonal", 0))]),
)
register_custom_wrapper(
    "torch",
    "triu",
    make_translator([("m", ("input",)), ("k", ("diagonal", 0))]),
)
register_custom_wrapper(
    "torch",
    "take_along_axis",
    make_translator(
        [("arr", ("input",)), ("indices", ("indices",)), ("axis", ("dim", -1))]
    ),
)
register_custom_wrapper(
    "torch",
    "linalg.norm",
    make_translator(
        [
            ("x", ("input",)),
            ("ord", ("ord", None)),
            ("axis", ("dim", None)),
            ("keepdims", ("keepdim", False)),
        ]
    ),
)

_torch_reduce_translation = [
    ("a", ("input",)),
    ("axis", ("dim",)),
    ("keepdims", ("keepdim",)),
]
for f in ("sum", "max", "min", "prod", "mean", "median", "std", "var"):
    # TODO: search "keepdim" in torch docs to find more
    register_custom_wrapper(
        "torch", f, make_translator(_torch_reduce_translation)
    )


for f in _CREATION_ROUTINES:
    register_creation_routine("torch", f, inject_device=True)
register_creation_routine(
    "torch", "array", inject_dtype=False, inject_device=True
)
register_creation_routine(
    "torch", "asarray", inject_dtype=False, inject_device=True
)
register_creation_routine(
    "torch", "random.default_rng", inject_dtype=False, inject_device=True
)

# for older versions of torch, can provide some alternative implementations
register_module_alias("torch[alt]", "torch")

register_submodule_alias("torch[alt]", "linalg.norm", "torch")
register_submodule_alias("torch[alt]", "linalg.qr", "torch")
register_submodule_alias("torch[alt]", "linalg.solve", "torch")
register_submodule_alias("torch[alt]", "linalg.svd", "torch")

register_custom_wrapper("torch[alt]", "linalg.qr", qr_allow_fat)
register_custom_wrapper("torch[alt]", "linalg.svd", svd_UsV_to_UsVH_wrapper)


@register_function("torch", "to_numpy")
def torch_to_numpy(x):
    return x.detach().cpu().numpy()


@register_function("torch", "copy")
def torch_copy(x):
    return x.detach().clone()


@register_function("torch", "transpose")
def torch_transpose(x, axes=None):
    if axes is None:
        axes = reversed(range(0, x.dim()))
    return x.permute(*axes)


@register_function("torch", "astype")
def torch_astype(x, dtype):
    return x.to(dtype=to_backend_dtype(dtype, like=x))


@register_function("torch", "complex")
def torch_complex(re, im):
    return complex_add_re_im(re, im)


@register_function("torch", "real")
def torch_real(x):
    # torch doesn't support calling real on real arrays
    try:
        if x.is_complex():
            return x.real
    except AttributeError:
        pass
    return x


@register_function("torch", "imag")
def torch_imag(x):
    # torch doesn't support calling imag on real arrays
    try:
        if x.is_complex():
            return x.imag
    except AttributeError:
        pass
    return get_torch().zeros_like(x)


@register_function("torch[alt]", "linalg.eigh")
def torch_linalg_eigh(x):
    torch = get_torch()
    return tuple(torch.symeig(x, eigenvectors=True))


@register_function("torch[alt]", "linalg.eigvalsh")
def torch_linalg_eigvalsh(x):
    torch = get_torch()
    return tuple(torch.symeig(x, eigenvectors=False))


@register_function("torch", "scipy.linalg.solve_triangular")
def torch_scipy_linalg_solve_triangular(
    a,
    b,
    lower=False,
    unit_diagonal=False,
    **kwargs,
):
    torch = get_torch()

    return torch.linalg.solve_triangular(
        a, b, upper=not lower, unitriangular=unit_diagonal, **kwargs
    )


@register_function("torch", "pad")
def torch_pad(array, pad_width, mode="constant", constant_values=0):
    torch = get_torch()

    if mode != "constant":
        raise NotImplementedError

    try:
        # numpy takes pads like ((0, 0), (1, 1), ... (n-1, n-1))
        # torch takes pads like (n-1, n-1, n-2, n-2, n-3, n-3, ...)
        pad = tuple(itertools.chain.from_iterable(pad_width))[::-1]

        # a single tuple was specified ((a, b),) - use for all axes
        if len(pad) == 2:
            pad = pad * array.ndimension()

    except TypeError:
        # assume int
        pad = (pad_width,) * 2 * array.ndimension()

    return torch.nn.functional.pad(
        array,
        pad=pad,
        mode=mode,
        value=constant_values,
    )


@register_function("torch", "indices")
def torch_indices(dimensions):
    torch = get_torch()
    return torch.meshgrid(*map(torch.arange, dimensions), indexing="ij")


@register_function("torch", "take")
def torch_take(a, indices, axis=None):
    torch = get_torch()

    if isinstance(indices, (tuple, list)) and len(indices) == 1:
        # need to convert list/tuple dimension to tensor, and not squeeze
        indices = torch.as_tensor(indices[0], device=a.device)
        unsqueeze = True
        squeeze = False
    else:
        indices = torch.as_tensor(indices, device=a.device)
        # XXX: if scalar can't yet use torch.select as it can't vmap
        squeeze = unsqueeze = indices.ndim == 0

    if unsqueeze:
        # promote scalar indices to 1D
        indices = torch.unsqueeze(indices, dim=0)

    # perform the take!
    a = torch.index_select(a, dim=axis, index=indices)

    if squeeze:
        # remove scalar dimension
        return a.squeeze(dim=axis)

    return a


@register_function("torch", "trace")
def torch_trace(x, axis1=0, axis2=1):
    if x.dim() == 2:
        # prefer built in
        return x.trace()
    return x.diagonal(dim1=axis1, dim2=axis2).sum(-1)


# ---------------------------------- mxnet ---------------------------------- #


def mxnet_to_numpy(x):
    return x.asnumpy()


register_module_alias("mxnet", "mxnet.numpy")

register_function("mxnet", "to_numpy", mxnet_to_numpy)


# --------------------------------- paddle ---------------------------------- #

_paddle_dtype_name_conversion = {
    "BOOL": "bool",
    "INT8": "int8",
    "INT16": "int16",
    "INT32": "int32",
    "INT64": "int64",
    "FP16": "float16",
    "FP32": "float32",
    "FP64": "float64",
    "COMPLEX64": "complex64",
    "COMPLEX128": "complex128",
}


@get_dtype_name.register("paddle")
def paddle_get_dtype_name(x):
    return _paddle_dtype_name_conversion[x.dtype.name]


@shape.register("paddle")
def paddle_shape(x):
    # convert from list
    return tuple(x.shape)


def paddle_wrap_reduction(fn):
    def numpy_like(*args, **kwargs):
        keepdims = kwargs.pop("keepdims", None)
        if keepdims is not None:
            kwargs["keepdim"] = keepdims
        return fn(*args, **kwargs)

    return numpy_like


def paddle_split_wrap(fn):
    # paddle doesn't seem to have `tensor_split always`

    @functools.wraps(fn)
    def numpy_like(ary, indices_or_sections, axis=0, **kwargs):
        if isinstance(indices_or_sections, numbers.Integral):
            return fn(ary, indices_or_sections, axis=axis, **kwargs)
        else:
            diff = do(
                "diff",
                indices_or_sections,
                prepend=0,
                append=shape(ary)[axis],
                like="numpy",
            )
            diff = list(diff)
            return fn(ary, diff, axis=axis)

    return numpy_like


register_module_alias("paddle[alt]", "paddle")

register_submodule_alias("paddle", "random.normal", "paddle")
register_submodule_alias("paddle", "random.uniform", "paddle")

register_func_alias("paddle", "asarray", "to_tensor")
register_func_alias("paddle", "concatenate", "concat")
register_func_alias("paddle", "identity", "eye")
register_func_alias("paddle", "power", "pow")
register_func_alias("paddle", "random.normal", "randn")
register_func_alias("paddle", "random.uniform", "rand")
register_func_alias("paddle", "split", "tensor_split")

register_custom_wrapper(
    "paddle", "random.normal", scale_random_normal_manually
)
register_custom_wrapper(
    "paddle", "random.uniform", scale_random_uniform_manually
)
register_custom_wrapper("paddle[alt]", "split", paddle_split_wrap)
register_custom_wrapper(
    "paddle", "tril", make_translator([("m", ("x",)), ("k", ("diagonal", 0))])
)
register_custom_wrapper(
    "paddle", "triu", make_translator([("m", ("x",)), ("k", ("diagonal", 0))])
)
for f in ("sum", "max", "min", "prod", "mean", "std", "var"):
    register_custom_wrapper("paddle", f, paddle_wrap_reduction)


@register_function("paddle", "imag")
def paddle_imag(x):
    # paddle doesn't support calling imag on real arrays
    try:
        if x.is_complex():
            return x.imag()
    except AttributeError:
        pass
    return do("zeros_like", x)


@register_function("paddle", "indices")
def paddle_indices(dimensions):
    _meshgrid = get_lib_fn("paddle", "meshgrid")
    _arange = get_lib_fn("paddle", "arange")
    return _meshgrid(*map(_arange, dimensions), indexing="ij")


@register_function("paddle", "pad")
def paddle_pad(array, pad_width, mode="constant", constant_values=0):
    if mode != "constant":
        raise NotImplementedError

    try:
        # numpy takes pads like ((0, 0), (1, 1), ... (n-1, n-1))
        # paddle takes pads like (0, 0, 1, 1, 2, 2, ...)
        pad = tuple(itertools.chain.from_iterable(pad_width))

        # a single tuple was specified ((a, b),) - use for all axes
        if len(pad) == 2:
            pad = pad * array.ndim

    except TypeError:
        # assume int
        pad = (pad_width,) * 2 * array.ndim

    return do(
        "nn.functional.pad",
        array,
        pad=pad,
        mode=mode,
        value=constant_values,
        like="paddle",
    )


@register_function("paddle", "ravel")
def paddle_ravel(x):
    return x.reshape((-1,))


@register_function("paddle", "real")
def paddle_real(x):
    # paddle doesn't support calling real on real arrays
    try:
        if x.is_complex():
            return x.real()
    except AttributeError:
        pass
    return x


@register_function("paddle", "to_numpy")
def paddle_to_numpy(x):
    return x.numpy()


@register_function("paddle", "transpose")
def paddle_transpose(a, axes=None):
    if axes is None:
        axes = tuple(range(a.ndim - 1, -1, -1))
    return a.transpose(perm=axes)


# -------------------------------- pytensor --------------------------------- #


@shape.register("pytensor")
def pytensor_shape(x):
    return x.type.shape


def pytensor_wrap_qr_with_shapes(fn):
    import pytensor.tensor as pt

    @functools.wraps(fn)
    def qr_shaped(x, **kwargs):
        *b, m, n = x.type.shape
        k = min(m, n)
        q, r = fn(x, mode="economic", **kwargs)
        q = pt.specify_shape(q, (*b, m, k))
        r = pt.specify_shape(r, (*b, k, n))
        return q, r

    return qr_shaped


def pytensor_wrap_svd_with_shapes(fn):
    import pytensor.tensor as pt

    @functools.wraps(fn)
    def svd_shaped(x, full_matrices=False, **kwargs):
        *b, m, n = x.type.shape

        k = min(m, n)
        if full_matrices:
            u_shape = (*b, m, m)
            vh_shape = (*b, n, n)
        else:
            u_shape = (*b, m, k)
            vh_shape = (*b, k, n)

        u, s, vh = fn(x, full_matrices=full_matrices, **kwargs)
        u = pt.specify_shape(u, u_shape)
        s = pt.specify_shape(s, (*b, k))
        vh = pt.specify_shape(vh, vh_shape)
        return u, s, vh

    return svd_shaped


register_module_alias("pytensor", "pytensor.tensor")

register_custom_wrapper("pytensor", "linalg.qr", pytensor_wrap_qr_with_shapes)
register_custom_wrapper(
    "pytensor", "linalg.svd", pytensor_wrap_svd_with_shapes
)
