autoray ======= .. py:module:: autoray Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/autoray/_version/index /autoapi/autoray/autoray/index /autoapi/autoray/compiler/index /autoapi/autoray/experimental/index /autoapi/autoray/grad/index /autoapi/autoray/lazy/index Attributes ---------- .. autoapisummary:: autoray.numpy Classes ------- .. autoapisummary:: autoray.DoFunc Functions --------- .. autoapisummary:: autoray.astype autoray.backend_like autoray.compose autoray.conj autoray.dag autoray.do autoray.get_backend autoray.get_common_dtype autoray.get_dtype_name autoray.get_lib_fn autoray.get_namespace autoray.imag autoray.infer_backend autoray.infer_backend_device_dtype autoray.infer_backend_multi autoray.is_array autoray.is_scalar autoray.ndim autoray.real autoray.register_backend autoray.register_function autoray.reshape autoray.set_backend autoray.shape autoray.size autoray.to_backend_dtype autoray.to_numpy autoray.transpose autoray.tree_apply autoray.tree_flatten autoray.tree_iter autoray.tree_map autoray.tree_unflatten autoray.autojit autoray.stop_gradient Package Contents ---------------- .. py:class:: DoFunc(fn: str) Get an automatic dispatch (i.e. backend selection deferred to call time) callable for function named ``fn``. Slightly faster equivalent to ``functools.partial(do, fn)``. .. rubric:: Examples `DoFunc` objects have a fixed operation but can still be called on any type of array: >>> sqrt = DoFunc('sqrt') >>> sqrt >>> import numpy as np >>> sqrt(np.random.uniform(size=[5])) array([0.32464973, 0.90379787, 0.85037325, 0.88729814, 0.46768083]) >>> import cupy as cp >>> sqrt(cp.random.uniform(size=[5])) array([0.44541656, 0.88713113, 0.92626237, 0.64080557, 0.69620767]) >>> import tensorflow as tf >>> sqrt(tf.random.uniform(shape=[5])) .. py:attribute:: __slots__ :value: ('fn',) .. py:attribute:: fn .. py:method:: __call__(*args, like=None, **kwargs) .. py:method:: __repr__() .. py:function:: astype(x, dtype_name, **kwargs) Cast array as type ``dtype_name`` - tries ``x.astype`` first. .. py:function:: backend_like(like, set_globally='auto') Context manager for setting a default backend. The argument ``like`` can be an explicit backend name or an ``array`` to infer it from. :param like: The backend to set. If an array, the backend of the array's class will be set. :type like: str or array :param set_globally: Whether to set the backend globally or for the current thread: - True: set the backend globally. - False: set the backend for the current thread. - "auto": set the backend globally if this thread is the thread that imported autoray. Otherwise set the backend for the current thread. Only one thread should ever call this function with ``set_globally=True``, (by default this is importing thread). :type set_globally: {"auto", False, True}, optional .. py:function:: 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:: conj(x) Array conjugate. .. py:function:: dag(x) Array Hermitian transpose. .. py:function:: do(fn: str, *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. :param fn: Name of the function to do, e.g. 'sum' or 'linalg.svd'. :type fn: str :param args: Positional arguments to pass to the function. :param like: 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. :type like: str or array, optional :param kwargs: Keyword arguments to pass to the function. .. 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:: 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:: get_common_dtype(*arrays) Compute the minimal dtype sufficient for ``arrays``. .. py:function:: get_dtype_name(x) Find string specifier ``dtype_name`` of array ``x``. .. 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:: 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. :param like: An array-like object to dispatch on, an explicit backend name, or None. :type like: array-like, str or None, optional :param device: The device to use for array creation, or None to infer from `like`. :type device: str or None, optional :param dtype: The data type to use for array creation, or None to infer from `like`. :type dtype: str or None, optional :returns: An automatic namespace object. :rtype: AutoNamespace .. py:function:: imag(x) Array imaginary part. .. py:function:: infer_backend(array) Get the name of the library that defined the class of ``array`` - unless ``array`` is directly a subclass of ``numpy.ndarray``, in which case assume ``numpy`` is the desired backend. .. py:function:: infer_backend_device_dtype(like, device=None, dtype=None) Infer the backend, device and dtype from `like`, with optional overrides for device and dtype. The dispatcher is cached on `like.__class__` to avoid repeated lookups of the same attributes for the same type of array. :param like: The array to infer the backend, device and dtype from. If str, an explicit backend name. If None, the backend None is simply returned. :type like: array-like or str or None :param device: If given, an explicit device to use. If None, and `like` is an array with a device attribute, that is used. :type device: str or device_like, optional :param dtype: If given, an explicit dtype to use. If None, and `like` is an array with a dtype attribute, that is used. :type dtype: str or dtype_like, optional :returns: * **backend** (*str or None*) -- The inferred backend name, or None if `like` is None. * **device** (*str or device_like or None*) -- The inferred device, or None if not given and not found on `like`. * **dtype** (*str or dtype_like or None*) -- The inferred dtype, or None if not given and not found on `like`. .. py:function:: infer_backend_multi(*arrays) Infer which backend should be used for a function that takes multiple arguments. This assigns a priority to each backend, and returns the backend with the highest priority. By default, the priority is: - ``builtins``: -2 - ``numpy``: -1 - other backends: 0 - ``autoray.lazy``: 1 I.e. when mixing with ``numpy``, other array libraries are preferred, when mixing with ``autoray.lazy``, ``autoray.lazy`` is preferred. This has quite low overhead due to caching. .. py:function:: 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. :param x: Object to check. :type x: object :rtype: bool .. seealso:: :py:obj:`is_scalar` .. py:function:: 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. :param x: Object to check. :type x: object :rtype: bool .. seealso:: :py:obj:`is_array` .. 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:data:: numpy .. py:function:: real(x) Array real part. .. py:function:: register_backend(cls, name) Register the name (and by default the module or submodule) of a custom array class. :param cls: The array class itself. :type cls: type :param name: The name of the backend that should be used for this class. By default this wil be assumed to be the location of the relevant functions for this class, but this can be overridden. :type name: str .. py:function:: register_function(backend, name, fn=None, wrap=False) Directly provide a custom implementation. :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. If not supplied, this function can be used as a decorator with ``backend`` and ``name`` only. :type fn: callable, optional :param wrap: 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`. :type wrap: bool, optional .. py:function:: reshape(x, shape) Array reshaped. .. 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:: 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:: 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:: to_backend_dtype(dtype_name, like) Turn string specifier ``dtype_name`` into dtype of backend ``like``. .. py:function:: to_numpy(x) Get a numpy version of array ``x``. .. py:function:: transpose(x, *args) Array transpose. .. py:function:: tree_apply(f, tree, is_leaf=is_not_container) Apply ``f`` to all leaves in ``tree``, no new pytree is built. :param f: A function to apply to all leaves in ``tree``. :type f: callable :param tree: A nested sequence of tuples, lists, dicts and other objects. :type tree: pytree :param is_leaf: A function to determine if an object is a leaf, ``f`` is only applied to objects for which ``is_leaf(x)`` returns ``True``. :type is_leaf: callable .. py:function:: tree_flatten(tree, is_leaf=is_not_container, get_ref=False) Flatten ``tree`` into a list of leaves. :param tree: A nested sequence of tuples, lists, dicts and other objects. :type tree: pytree :param is_leaf: A function to determine if an object is a leaf, only objects for which ``is_leaf(x)`` returns ``True`` are returned in the flattened list. :type is_leaf: callable :param get_ref: If ``True``, a reference tree is also returned which can be used to reconstruct the original tree from a flattened list. :type get_ref: bool :returns: * **objs** (*list*) -- The flattened list of leaf objects. * **(ref_tree)** (*pytree*) -- If ``get_ref`` is ``True``, a reference tree, with leaves of ``Leaf``, is returned which can be used to reconstruct the original tree. .. py:function:: tree_iter(tree, is_leaf=is_not_container) Iterate over all leaves in ``tree``. :param f: A function to apply to all leaves in ``tree``. :type f: callable :param tree: A nested sequence of tuples, lists, dicts and other objects. :type tree: pytree :param is_leaf: A function to determine if an object is a leaf, ``f`` is only applied to objects for which ``is_leaf(x)`` returns ``True``. :type is_leaf: callable .. py:function:: tree_map(f, tree, is_leaf=is_not_container) Map ``f`` over all leaves in ``tree``, returning a new pytree. :param f: A function to apply to all leaves in ``tree``. :type f: callable :param tree: A nested sequence of tuples, lists, dicts and other objects. :type tree: pytree :param is_leaf: A function to determine if an object is a leaf, ``f`` is only applied to objects for which ``is_leaf(x)`` returns ``True``. :type is_leaf: callable :rtype: pytree .. py:function:: tree_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:: autojit(fn=None, *, backend=None, compiler_opts=None) Just-in-time compile an ``autoray`` function, automatically choosing the backend based on the input arrays, or via keyword argument. The backend used to do the compilation can be set in three ways: 1. Automatically based on the arrays the function is called with, i.e. ``cfn(*torch_arrays)`` will use ``torch.jit.trace``. 2. In this wrapper, ``@autojit(backend='jax')``, to provide a specific default instead. 3. When you call the function ``cfn(*arrays, backend='torch')`` to override on a per-call basis. If the arrays supplied are of a different backend type to the compiler, then the returned array will also be converted back, i.e. ``cfn(*numpy_arrays, backend='tensorflow')`` will return a ``numpy`` array. The ``'python'`` backend simply extracts and unravels all the ``do`` calls into a code object using ``compile`` which is then run with ``exec``. This makes use of shared intermediates and constant folding, strips away any python scaffoliding, and is compatible with any library, but the resulting function is not 'low-level' in the same way as the other backends. :param fn: The autoray function to compile. :type fn: callable :param backend: If set, use this as the default backend. The options are: - ``'python'``: extract and unravel all the ``do`` calls into a code object using ``compile`` which is then run with ``exec``. - ``'jax'``: use `jax.jit` to compile the function. - ``'tensorflow'``: use `tf.function` to compile the function. - ``'torch'``: use `torch.jit.trace` to compile the function. - ``'pytensor'``: use `pytensor.function` to compile the function. If not set, the backend will be inferred from the input arrays, or it can be set at call time with the ``backend`` keyword argument. :type backend: str, optional :param compiler_opts: Dict of dicts when you can supply options for each compiler backend separately, e.g.: ``@autojit(compiler_opts={'tensorflow': {'jit_compile': True}})``. :type compiler_opts: dict[dict], optional :returns: **cfn** -- The function with auto compilation. :rtype: callable .. py:function:: stop_gradient(x) Stop gradient flow through array ``x``. In autodiff backends (JAX, PyTorch, TensorFlow, etc.), this detaches ``x`` from the computational graph so that no gradients are propagated through it. For non-autodiff backends (NumPy, CuPy, etc.), this is a no-op and returns ``x`` unchanged. :param x: The array to stop gradient flow through. :type x: array :returns: An array with the same value as ``x`` but detached from the autodiff computational graph. :rtype: array .. rubric:: Examples With JAX:: >>> import jax, jax.numpy as jnp, autoray as ar >>> x = jnp.array([1.0, 2.0, 3.0]) >>> ar.stop_gradient(x) # equivalent to jax.lax.stop_gradient(x) With PyTorch:: >>> import torch, autoray as ar >>> x = torch.tensor([1.0, 2.0], requires_grad=True) >>> y = ar.stop_gradient(x) # equivalent to x.detach() >>> y.requires_grad False With NumPy (no-op):: >>> import numpy as np, autoray as ar >>> x = np.array([1.0, 2.0]) >>> ar.stop_gradient(x) is x True