autoray.autoray =============== .. py:module:: autoray.autoray .. autoapi-nested-parse:: AUTORAY - backend agnostic array operations. Copyright 2019-2023 Johnnie Gray Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. Attributes ---------- .. autoapisummary:: autoray.autoray._global_backend autoray.autoray._inferrer_global autoray.autoray._infer_auto autoray.autoray._global_backends_threadaware autoray.autoray._inferrers_threadaware autoray.autoray._importing_thrid autoray.autoray._backend_lock autoray.autoray._CUSTOM_BACKENDS autoray.autoray.multi_class_priorities autoray.autoray._CREATION_ROUTINES autoray.autoray._CREATION_INJECT autoray.autoray._BACKEND_ALIASES autoray.autoray._MODULE_ALIASES autoray.autoray._SUBMODULE_ALIASES autoray.autoray._FUNC_ALIASES autoray.autoray._CUSTOM_WRAPPERS autoray.autoray._FUNCS autoray.autoray._COMPOSED_FUNCTION_GENERATORS autoray.autoray.TREE_MAP_REGISTRY autoray.autoray.TREE_APPLY_REGISTRY autoray.autoray.TREE_ITER_REGISTRY autoray.autoray.IS_CONTAINER_CACHE autoray.autoray.TREE_MAPPER_CACHE autoray.autoray.TREE_ITER_CACHE autoray.autoray.TREE_APPLIER_CACHE autoray.autoray.LEAF autoray.autoray._COMPLEX_DTYPES autoray.autoray._DOUBLE_DTYPES autoray.autoray._DTYPE_MAP autoray.autoray._DISPATCHERS autoray.autoray.numpy autoray.autoray._builtin_dtype_lookup autoray.autoray._JAX_RANDOM_KEY autoray.autoray._torch_reduce_translation autoray.autoray._paddle_dtype_name_conversion Classes ------- .. autoapisummary:: autoray.autoray.Leaf autoray.autoray.Composed autoray.autoray.NumpyMimic Functions --------- .. autoapisummary:: autoray.autoray.do autoray.autoray._default_infer_from_sig autoray.autoray._default_infer_from_sig_threadaware autoray.autoray._always_the_same autoray.autoray.get_backend autoray.autoray.set_backend autoray.autoray.backend_like autoray.autoray.register_backend autoray.autoray._infer_class_backend_cached autoray.autoray.infer_backend autoray.autoray._infer_class_backend_multi_cached autoray.autoray.infer_backend_multi autoray.autoray.register_creation_routine autoray.autoray._maybe_inject_dtype_device autoray.autoray._choose_backend autoray.autoray.choose_backend autoray.autoray.import_lib_fn autoray.autoray.get_lib_fn autoray.autoray.register_function autoray.autoray.tree_register_container autoray.autoray.is_not_container autoray.autoray.is_array autoray.autoray.identity autoray.autoray.tree_map autoray.autoray.empty autoray.autoray.tree_iter autoray.autoray.nothing autoray.autoray.tree_apply autoray.autoray.is_leaf_placeholder autoray.autoray.tree_flatten autoray.autoray.tree_unflatten autoray.autoray.tree_map_tuple autoray.autoray.tree_iter_tuple autoray.autoray.tree_apply_tuple autoray.autoray.tree_map_list autoray.autoray.tree_iter_list autoray.autoray.tree_apply_list autoray.autoray.tree_map_dict autoray.autoray.tree_iter_dict autoray.autoray.tree_apply_dict autoray.autoray.compose autoray.autoray.shape autoray.autoray.ndim autoray.autoray.size autoray.autoray.conj autoray.autoray.transpose autoray.autoray.dag autoray.autoray.real autoray.autoray.imag autoray.autoray.reshape autoray.autoray.to_backend_dtype autoray.autoray.get_dtype_name autoray.autoray.get_common_dtype autoray.autoray.astype autoray.autoray.to_numpy autoray.autoray.svd_not_full_matrices_wrapper autoray.autoray.svd_sUV_to_UsVH_wrapper autoray.autoray.svd_UsV_to_UsVH_wrapper autoray.autoray.svd_manual_full_matrices_kwarg autoray.autoray.qr_allow_fat autoray.autoray.tril_to_band_part autoray.autoray.triu_to_band_part autoray.autoray.cholesky_lower autoray.autoray.binary_allow_1d_rhs_wrap autoray.autoray.scale_random_uniform_manually autoray.autoray.scale_random_normal_manually autoray.autoray.with_dtype_wrapper autoray.autoray.translate_wrapper autoray.autoray.make_translator autoray.autoray.complex_add_re_im autoray.autoray.allclose autoray.autoray.wrap_args_kwargs_from_raw autoray.autoray.register_dispatch autoray.autoray.default_dispatcher autoray.autoray.join_array_dispatcher autoray.autoray.einsum_dispatcher autoray.autoray.binary_dispatcher autoray.autoray._get_mimic_function_or_attribute autoray.autoray.builtins_get_dtype_name autoray.autoray.numpy_to_numpy autoray.autoray.cupy_to_numpy autoray.autoray.jax_random_seed autoray.autoray.jax_random_get_key autoray.autoray.jax_random_uniform autoray.autoray.jax_random_normal autoray.autoray.jax_to_numpy autoray.autoray.aesara_shape autoray.autoray.dask_to_numpy autoray.autoray.dask_eye_wrapper autoray.autoray.mars_to_numpy autoray.autoray.ctf_array autoray.autoray.ctf_to_numpy autoray.autoray.ctf_count_nonzero autoray.autoray.ctf_get_dtype_name autoray.autoray.sparse_array autoray.autoray.sparse_to_numpy autoray.autoray.sparse_transpose autoray.autoray.sparse_reshape autoray.autoray.sparse_sum autoray.autoray.sparse_prod autoray.autoray.sparse_conj autoray.autoray.sparse_real autoray.autoray.sparse_imag autoray.autoray.sparse_count_nonzero autoray.autoray.sparse_random_uniform autoray.autoray.sparse_random_normal autoray.autoray.tensorflow_to_numpy autoray.autoray.tensorflow_indices autoray.autoray.tensorflow_pad_wrap autoray.autoray.torch_shape autoray.autoray.torch_size autoray.autoray.torch_to_numpy autoray.autoray.torch_copy autoray.autoray.torch_transpose autoray.autoray.torch_count_nonzero autoray.autoray.torch_astype autoray.autoray._torch_get_dtype_name autoray.autoray.torch_get_dtype_name autoray.autoray.torch_real autoray.autoray.torch_imag autoray.autoray.torch_linalg_solve_wrap autoray.autoray.torch_linalg_eigh autoray.autoray.torch_linalg_eigvalsh autoray.autoray.torch_tensordot_wrap autoray.autoray.torch_pad autoray.autoray.torch_split_wrap autoray.autoray.torch_zeros_ones_wrap autoray.autoray.torch_eye_wrap autoray.autoray.torch_sort_wrap autoray.autoray.torch_indices autoray.autoray.torch_flip_wrap autoray.autoray.mxnet_to_numpy autoray.autoray.paddle_get_dtype_name autoray.autoray.paddle_shape autoray.autoray.paddle_to_numpy autoray.autoray.paddle_transpose autoray.autoray.paddle_real autoray.autoray.paddle_imag autoray.autoray.paddle_indices autoray.autoray.paddle_ravel autoray.autoray.paddle_pad autoray.autoray.paddle_wrap_reduction autoray.autoray.paddle_split_wrap Module Contents --------------- .. py:function:: do(fn, *args, like=None, **kwargs) Do function named ``fn`` on ``(*args, **kwargs)``, peforming single dispatch to retrieve ``fn`` based on whichever library defines the class of the ``args[0]``, or the ``like`` keyword argument if specified. .. rubric:: Examples Works on numpy arrays: >>> import numpy as np >>> x_np = np.random.uniform(size=[5]) >>> y_np = do('sqrt', x_np) >>> y_np array([0.32464973, 0.90379787, 0.85037325, 0.88729814, 0.46768083]) >>> type(y_np) numpy.ndarray Works on cupy arrays: >>> import cupy as cp >>> x_cp = cp.random.uniform(size=[5]) >>> y_cp = do('sqrt', x_cp) >>> y_cp array([0.44541656, 0.88713113, 0.92626237, 0.64080557, 0.69620767]) >>> type(y_cp) cupy.core.core.ndarray Works on tensorflow arrays: >>> import tensorflow as tf >>> x_tf = tf.random.uniform(shape=[5]) >>> y_tf = do('sqrt', x_tf) >>> y_tf >>> 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) .. py:function:: _default_infer_from_sig(fn, *args, **kwargs) This is the default backend dispatcher, used if no global backend has been set. Hot swapping this function out as below avoids having to check manually for a global backend or worse, a thread aware global backend, on every call to ``do``. .. py:data:: _global_backend :value: None .. py:data:: _inferrer_global .. py:data:: _infer_auto .. py:data:: _global_backends_threadaware .. py:data:: _inferrers_threadaware .. py:data:: _importing_thrid .. py:data:: _backend_lock .. py:function:: _default_infer_from_sig_threadaware(fn, args, kwargs) .. py:function:: _always_the_same(fn, args, kwargs, backend) .. py:function:: get_backend(get_globally='auto') Return the universally set backend, if any. :param get_globally: Which backend to return: - True: return the globally set backend, if any. - False: return the backend set for the current thread, if any. - "auto": return the globally set backend, if this thread is the thread that imported autoray. Otherwise return the backend set for the current thread, if any. :type get_globally: {"auto", False, True}, optional :returns: **backend** -- The name of the backend, or None if no backend is set. :rtype: str or None .. py:function:: set_backend(like, set_globally='auto') Set a default global backend. The argument ``like`` can be an explicit backend name or an ``array``. :param like: The backend to set. If an array, the backend of the array's class will be set. :type like: str or array :param set_globally: Whether to set the backend globally or for the current thread: - True: set the backend globally. - False: set the backend for the current thread. - "auto": set the backend globally if this thread is the thread that imported autoray. Otherwise set the backend for the current thread. Only one thread should ever call this function with ``set_globally=True``, (by default this is importing thread). :type set_globally: {"auto", False, True}, optional .. py:function:: backend_like(like, set_globally='auto') Context manager for setting a default backend. The argument ``like`` can be an explicit backend name or an ``array`` to infer it from. :param like: The backend to set. If an array, the backend of the array's class will be set. :type like: str or array :param set_globally: Whether to set the backend globally or for the current thread: - True: set the backend globally. - False: set the backend for the current thread. - "auto": set the backend globally if this thread is the thread that imported autoray. Otherwise set the backend for the current thread. Only one thread should ever call this function with ``set_globally=True``, (by default this is importing thread). :type set_globally: {"auto", False, True}, optional .. py:data:: _CUSTOM_BACKENDS .. py:function:: register_backend(cls, name) Register the name (and by default the module or submodule) of a custom array class. :param cls: The array class itself. :type cls: type :param name: The name of the backend that should be used for this class. By default this wil be assumed to be the location of the relevant functions for this class, but this can be overridden. :type name: str .. py:function:: _infer_class_backend_cached(cls) .. py:function:: infer_backend(array) Get the name of the library that defined the class of ``array`` - unless ``array`` is directly a subclass of ``numpy.ndarray``, in which case assume ``numpy`` is the desired backend. .. py:data:: multi_class_priorities .. py:function:: _infer_class_backend_multi_cached(classes) .. py:function:: infer_backend_multi(*arrays) Infer which backend should be used for a function that takes multiple arguments. This assigns a priority to each backend, and returns the backend with the highest priority. By default, the priority is: - ``builtins``: -2 - ``numpy``: -1 - other backends: 0 - ``autoray.lazy``: 1 I.e. when mixing with ``numpy``, other array libraries are preferred, when mixing with ``autoray.lazy``, ``autoray.lazy`` is preferred. This has quite low overhead due to caching. .. py:data:: _CREATION_ROUTINES .. py:data:: _CREATION_INJECT .. py:function:: register_creation_routine(backend, fn, inject_dtype=True, inject_device=False) Register a function that creates a new array, with `dtype` and possibly `device` kwargs, that should be inferred from the like argument. This is not necessary for array creation routines that don't accept either. :param backend: The backend to register the function for. :type backend: str :param fn: The name of the function to register. :type fn: str :param inject_dtype: Whether to inject a `dtype` argument based on the `like` argument. :type inject_dtype: bool, optional :param inject_device: Whether to inject a `device` argument based on the `like` argument. :type inject_device: bool, optional .. py:function:: _maybe_inject_dtype_device(backend, fn, args, kwargs, like) .. py:function:: _choose_backend(fn, args, kwargs, like=None) Private function to choose a backend based on function name and signature, which passes args and kwargs by reference for performance and also to allow injection of dtype and device arguments for array creation routines. .. py:function:: choose_backend(fn, *args, like=None, **kwargs) Choose a backend based on function name, arguments, and the ``like`` keyword argument. The default, if ``like`` is not specified, is to infer the backend from the function call, the default of which is simply to use the first argument, if no custom dispatcher is found. Otherwise the backend is chosen based on the ``like`` argument - which can be an explicit backend name or an arbitrary object. .. py:data:: _BACKEND_ALIASES .. py:data:: _MODULE_ALIASES .. py:data:: _SUBMODULE_ALIASES .. py:data:: _FUNC_ALIASES .. py:data:: _CUSTOM_WRAPPERS .. py:data:: _FUNCS .. py:data:: _COMPOSED_FUNCTION_GENERATORS .. py:function:: import_lib_fn(backend, fn) .. py:function:: get_lib_fn(backend, fn) Cached retrieval of correct function for backend, all the logic for finding the correct funtion only runs the first time. :param backend: The module defining the array class to dispatch on. :type backend: str :param fn: The function to retrieve. :type fn: str :rtype: callable .. py:function:: register_function(backend, name, fn, wrap=False) Directly provide your own function. :param backend: The name of the backend to register the function for. :type backend: str :param name: Name of the function, e.g. `'sum'` or `'linalg.svd'`. :type name: str :param fn: The function to register. :type fn: callable :param wrap: Whether to wrap the old function like ``fn(old_fn)`` rather than directly supply the entire new function. :type wrap: bool, optional .. py:data:: TREE_MAP_REGISTRY .. py:data:: TREE_APPLY_REGISTRY .. py:data:: TREE_ITER_REGISTRY .. py:function:: tree_register_container(cls, mapper, iterator, applier) Register a new container type for use with ``tree_map`` and ``tree_apply``. :param cls: The container type to register. :type cls: type :param mapper: A function that takes ``f``, ``tree`` and ``is_leaf`` and returns a new tree of type ``cls`` with ``f`` applied to all leaves. :type mapper: callable :param applier: A function that takes ``f``, ``tree`` and ``is_leaf`` and applies ``f`` to all leaves in ``tree``. :type applier: callable .. py:data:: IS_CONTAINER_CACHE .. py:function:: is_not_container(x) The default function to determine if an object is a leaf. This simply checks if the object is an instance of any of the registered container types. .. py:function:: is_array(x) An alternative leaf tester for addressing only arrays within trees. .. py:function:: identity(f, tree, is_leaf) .. py:data:: TREE_MAPPER_CACHE .. py:function:: tree_map(f, tree, is_leaf=is_not_container) Map ``f`` over all leaves in ``tree``, returning a new pytree. :param f: A function to apply to all leaves in ``tree``. :type f: callable :param tree: A nested sequence of tuples, lists, dicts and other objects. :type tree: pytree :param is_leaf: A function to determine if an object is a leaf, ``f`` is only applied to objects for which ``is_leaf(x)`` returns ``True``. :type is_leaf: callable :rtype: pytree .. py:function:: empty(tree, is_leaf) .. py:data:: TREE_ITER_CACHE .. py:function:: tree_iter(tree, is_leaf=is_not_container) Iterate over all leaves in ``tree``. :param f: A function to apply to all leaves in ``tree``. :type f: callable :param tree: A nested sequence of tuples, lists, dicts and other objects. :type tree: pytree :param is_leaf: A function to determine if an object is a leaf, ``f`` is only applied to objects for which ``is_leaf(x)`` returns ``True``. :type is_leaf: callable .. py:function:: nothing(f, tree, is_leaf) .. py:data:: TREE_APPLIER_CACHE .. py:function:: tree_apply(f, tree, is_leaf=is_not_container) Apply ``f`` to all leaves in ``tree``, no new pytree is built. :param f: A function to apply to all leaves in ``tree``. :type f: callable :param tree: A nested sequence of tuples, lists, dicts and other objects. :type tree: pytree :param is_leaf: A function to determine if an object is a leaf, ``f`` is only applied to objects for which ``is_leaf(x)`` returns ``True``. :type is_leaf: callable .. py:class:: Leaf A singleton object to use as a placeholder in a pytree, for unflattening. .. py:attribute:: __slots__ :value: () .. py:method:: __repr__() .. py:data:: LEAF .. py:function:: is_leaf_placeholder(x) .. py:function:: tree_flatten(tree, is_leaf=is_not_container, get_ref=False) Flatten ``tree`` into a list of leaves. :param tree: A nested sequence of tuples, lists, dicts and other objects. :type tree: pytree :param is_leaf: A function to determine if an object is a leaf, only objects for which ``is_leaf(x)`` returns ``True`` are returned in the flattened list. :type is_leaf: callable :param get_ref: If ``True``, a reference tree is also returned which can be used to reconstruct the original tree from a flattened list. :type get_ref: bool :returns: * **objs** (*list*) -- The flattened list of leaf objects. * **(ref_tree)** (*pytree*) -- If ``get_ref`` is ``True``, a reference tree, with leaves of ``Leaf``, is returned which can be used to reconstruct the original tree. .. py:function:: tree_unflatten(objs, tree, is_leaf=is_leaf_placeholder) Unflatten ``objs`` into a pytree of the same structure as ``tree``. :param objs: A sequence of objects to be unflattened into a pytree. :type objs: sequence :param tree: A nested sequence of tuples, lists, dicts and other objects, the objs will be inserted into a new pytree of the same structure. :type tree: pytree :param is_leaf: A function to determine if an object is a leaf, only objects for which ``is_leaf(x)`` returns ``True`` will have the next item from ``objs`` inserted. By default checks for the ``Leaf`` object inserted by ``tree_flatten(..., get_ref=True)``. :type is_leaf: callable :rtype: pytree .. py:function:: tree_map_tuple(f, tree, is_leaf) .. py:function:: tree_iter_tuple(tree, is_leaf) .. py:function:: tree_apply_tuple(f, tree, is_leaf) .. py:function:: tree_map_list(f, tree, is_leaf) .. py:function:: tree_iter_list(tree, is_leaf) .. py:function:: tree_apply_list(f, tree, is_leaf) .. py:function:: tree_map_dict(f, tree, is_leaf) .. py:function:: tree_iter_dict(tree, is_leaf) .. py:function:: tree_apply_dict(f, tree, is_leaf) .. py:class:: Composed(fn, name=None) Compose an ``autoray.do`` using function. See the main wrapper ``compose``. .. py:attribute:: _default_fn .. py:attribute:: _name :value: None .. py:attribute:: _supply_backend .. py:method:: register(backend, fn=None) Register a different implementation for ``backend``. .. py:method:: make_function(backend) Make a new function for the specific ``backend``. .. py:method:: __call__(*args, like=None, **kwargs) .. py:method:: __repr__() .. py:function:: compose(fn, *, name=None) Take a function consisting of multiple ``autoray.do`` calls and compose it into a new, single, named function, registered with ``autoray.do``. This creates a default implementation of this function for each new backend encountered without explicitly having to write each out, but also allows for specific implementations to be overridden for specific backends. If the function takes a ``backend`` argument, it will be supplied with the backend name, to save having to re-choose the backend. Specific implementations can be provided by calling the ``register`` method of the composed function, or it can itself be used like a decorator:: @compose def foo(x): ... @foo.register("numpy") @numba.njit def foo_numba(x): ... :param fn: The funtion to compose, and its default implementation. :type fn: callable :param name: The name of the composed function. If not provided, the name of the function will be used. :type name: str, optional .. py:function:: shape(x) Get the shape of an array as a tuple of int. This should be preferred to calling `x.shape` directly, as it: 1. Allows customization (e.g. for torch and aesara which return different types for shape - use `@shape.register(backend)` to customize the behavior from this default implementation). 2. Can be used on nested lists and tuples, without calling numpy. :param x: The array to get the shape of. It can be an arbitrary nested list or tuple of arrays and scalars, but is assumed not to be ragged. :type x: array_like :returns: **shape** -- The size of each dimension of the array. :rtype: tuple of int .. py:function:: ndim(x) Get the number of dimensions of an array. This should be preferred to calling `x.ndim`, since not all backends implement that, and it can also be called on nested lists and tuples. :param x: The array to get the number of dimensions of. It can be an arbitrary nested list or tuple of arrays and scalars. :type x: array_like :returns: **ndim** :rtype: int .. py:function:: size(x) Get the size, or number of elements, of an array. This should be preferred to calling `x.size`, since not all backends implement that, and it can also be called on nested lists and tuples. :param x: The array to get the size of. It can be an arbitrary nested list or tuple of arrays and scalars. :type x: array_like :returns: **size** :rtype: int .. py:function:: conj(x) Array conjugate. .. py:function:: transpose(x, *args) Array transpose. .. py:function:: dag(x) Array Hermitian transpose. .. py:function:: real(x) Array real part. .. py:function:: imag(x) Array imaginary part. .. py:function:: reshape(x, shape) Array reshaped. .. py:function:: to_backend_dtype(dtype_name, like) Turn string specifier ``dtype_name`` into dtype of backend ``like``. .. py:function:: get_dtype_name(x) Find string specifier ``dtype_name`` of array ``x``. .. py:data:: _COMPLEX_DTYPES .. py:data:: _DOUBLE_DTYPES .. py:data:: _DTYPE_MAP .. py:function:: get_common_dtype(*arrays) Compute the minimal dtype sufficient for ``arrays``. .. py:function:: astype(x, dtype_name, **kwargs) Cast array as type ``dtype_name`` - tries ``x.astype`` first. .. py:function:: to_numpy(x) Get a numpy version of array ``x``. .. py:function:: svd_not_full_matrices_wrapper(fn) .. py:function:: svd_sUV_to_UsVH_wrapper(fn) .. py:function:: svd_UsV_to_UsVH_wrapper(fn) .. py:function:: svd_manual_full_matrices_kwarg(fn) .. py:function:: qr_allow_fat(fn) .. py:function:: tril_to_band_part(fn) .. py:function:: triu_to_band_part(fn) .. py:function:: cholesky_lower(fn) .. py:function:: binary_allow_1d_rhs_wrap(fn) .. py:function:: scale_random_uniform_manually(fn) .. py:function:: scale_random_normal_manually(fn) .. py:function:: with_dtype_wrapper(fn) Add ability to handle `dtype` keyword. If not None, `dtype` should be specified as a string, otherwise conversion will happen regardless. .. py:function:: translate_wrapper(fn, translator) Wrap a function to match the api of another according to a translation. The ``translator`` entries in the form of an ordered dict should have entries like: (desired_kwarg: (backend_kwarg, default_value)) with the order defining the args of the function. .. py:function:: make_translator(t) .. py:function:: complex_add_re_im(re, im) .. py:function:: allclose(x, y, rtol=1e-05, atol=1e-08) .. py:function:: wrap_args_kwargs_from_raw(fn) Take a function with signature ``(*args, **kwargs)`` and wrap it to accept a single tuple of args and a dict of kwargs. .. py:function:: register_dispatch(fun, dispatcher, raw_signature=True) Register a new dispatcher, a function that takes the arguments and keyword arguments of a function and returns the backend to use, when the backend is not explicitly given. This is useful in case the backend to be used by a function cannot be inferred from the first argument. :param fun: The name of the function to register the dispatcher for. :type fun: str :param dispatcher: The dispatcher function to use. This should take the arguments and keyword arguments of the function and return the backend to use. :type dispatcher: callable :param raw_signature: The ``dispatcher`` has signature ``(*args, **kwargs)`` if ``True``, otherwise it has signature ``(args, kwargs)``. :type raw_signature: bool, optional .. py:function:: default_dispatcher(args, kwargs) Try to infer backend from first argument passed to function. .. py:data:: _DISPATCHERS .. py:function:: join_array_dispatcher(args, kwargs) Dispatcher for functions where first argument is a sequence. .. py:function:: einsum_dispatcher(args, kwargs) Dispatcher for handling einsum. einsum can be called with a str equation as the first argument, or with 'interleaved' inputs. This dispatcher handles both cases and also takes into account all arrays. .. py:function:: binary_dispatcher(args, kwargs) There are cases when we want to take into account both backends of two arguments, e.g. a lazy variable and a constant array. .. py:function:: _get_mimic_function_or_attribute(self, fn) .. py:class:: NumpyMimic(submodule=None) A class to mimic the syntax of using `numpy` directly. .. py:attribute:: submodule :value: None .. py:method:: __getattribute__(attr) .. py:method:: __repr__() :staticmethod: .. py:data:: numpy .. py:data:: _builtin_dtype_lookup .. py:function:: builtins_get_dtype_name(x) .. py:function:: numpy_to_numpy(x) .. py:function:: cupy_to_numpy(x) .. py:data:: _JAX_RANDOM_KEY :value: None .. py:function:: jax_random_seed(seed=None) .. py:function:: jax_random_get_key() .. py:function:: jax_random_uniform(low=0.0, high=1.0, size=None, **kwargs) .. py:function:: jax_random_normal(loc=0.0, scale=1.0, size=None, **kwargs) .. py:function:: jax_to_numpy(x) .. py:function:: aesara_shape(x) .. py:function:: dask_to_numpy(x) .. py:function:: dask_eye_wrapper(fn) .. py:function:: mars_to_numpy(x) .. py:function:: ctf_array(x) .. py:function:: ctf_to_numpy(x) .. py:function:: ctf_count_nonzero(x) .. py:function:: ctf_get_dtype_name(x) .. py:function:: sparse_array(x) .. py:function:: sparse_to_numpy(x) .. py:function:: sparse_transpose(x, axes=None) .. py:function:: sparse_reshape(x, shape) .. py:function:: sparse_sum(x, axis=None, keepdims=False, dtype=None, out=None) .. py:function:: sparse_prod(x, axis=None, keepdims=False, dtype=None, out=None) .. py:function:: sparse_conj(x) .. py:function:: sparse_real(x) .. py:function:: sparse_imag(x) .. py:function:: sparse_count_nonzero(x) .. py:function:: sparse_random_uniform(low=0.0, high=1.0, size=None, dtype=None, **kwargs) .. py:function:: sparse_random_normal(loc=0.0, scale=1.0, size=None, dtype=None, **kwargs) .. py:function:: tensorflow_to_numpy(x) .. py:function:: tensorflow_indices(dimensions) .. py:function:: tensorflow_pad_wrap(tf_pad) .. py:function:: torch_shape(x) .. py:function:: torch_size(x) .. py:function:: torch_to_numpy(x) .. py:function:: torch_copy(x) .. py:function:: torch_transpose(x, axes=None) .. py:function:: torch_count_nonzero(x) .. py:function:: torch_astype(x, dtype) .. py:function:: _torch_get_dtype_name(dtype) .. py:function:: torch_get_dtype_name(x) .. py:function:: torch_real(x) .. py:function:: torch_imag(x) .. py:function:: torch_linalg_solve_wrap(fn) .. py:function:: torch_linalg_eigh(x) .. py:function:: torch_linalg_eigvalsh(x) .. py:function:: torch_tensordot_wrap(fn) .. py:function:: torch_pad(array, pad_width, mode='constant', constant_values=0) .. py:function:: torch_split_wrap(fn) .. py:function:: torch_zeros_ones_wrap(fn) .. py:function:: torch_eye_wrap(fn) .. py:function:: torch_sort_wrap(fn) .. py:function:: torch_indices(dimensions) .. py:function:: torch_flip_wrap(torch_flip) .. py:data:: _torch_reduce_translation :value: [('a', ('input',)), ('axis', ('dim',)), ('keepdims', ('keepdim',))] .. py:function:: mxnet_to_numpy(x) .. py:data:: _paddle_dtype_name_conversion .. py:function:: paddle_get_dtype_name(x) .. py:function:: paddle_shape(x) .. py:function:: paddle_to_numpy(x) .. py:function:: paddle_transpose(a, axes=None) .. py:function:: paddle_real(x) .. py:function:: paddle_imag(x) .. py:function:: paddle_indices(dimensions) .. py:function:: paddle_ravel(x) .. py:function:: paddle_pad(array, pad_width, mode='constant', constant_values=0) .. py:function:: paddle_wrap_reduction(fn) .. py:function:: paddle_split_wrap(fn)