autoray.lazy.core
=================

.. py:module:: autoray.lazy.core

.. autoapi-nested-parse::

   Core lazy array functionality.



Attributes
----------

.. autoapisummary::

   autoray.lazy.core._EMPTY_DICT
   autoray.lazy.core.get_depth
   autoray.lazy.core.get_data
   autoray.lazy.core.get_dtype
   autoray.lazy.core.PLACEHOLDER
   autoray.lazy.core._materialize_dispatch
   autoray.lazy.core._stringify_dispatch
   autoray.lazy.core._SHARING_STACK
   autoray.lazy.core._DTYPES_REAL_EQUIV
   autoray.lazy.core._DTYPES_COMPLEX_EQUIV
   autoray.lazy.core.permute_dims
   autoray.lazy.core.add
   autoray.lazy.core.eq
   autoray.lazy.core.equal
   autoray.lazy.core.floordivide
   autoray.lazy.core.ge
   autoray.lazy.core.gt
   autoray.lazy.core.le
   autoray.lazy.core.lt
   autoray.lazy.core.mod
   autoray.lazy.core.multiply
   autoray.lazy.core.ne
   autoray.lazy.core.pow_
   autoray.lazy.core.sub
   autoray.lazy.core.truedivide
   autoray.lazy.core.sin
   autoray.lazy.core.cos
   autoray.lazy.core.tan
   autoray.lazy.core.arcsin
   autoray.lazy.core.arccos
   autoray.lazy.core.arctan
   autoray.lazy.core.sinh
   autoray.lazy.core.cosh
   autoray.lazy.core.tanh
   autoray.lazy.core.arcsinh
   autoray.lazy.core.arccosh
   autoray.lazy.core.arctanh
   autoray.lazy.core.sqrt
   autoray.lazy.core.exp
   autoray.lazy.core.log
   autoray.lazy.core.log2
   autoray.lazy.core.log10
   autoray.lazy.core.conj
   autoray.lazy.core.sign
   autoray.lazy.core.abs_
   autoray.lazy.core.angle
   autoray.lazy.core.real
   autoray.lazy.core.imag
   autoray.lazy.core.sum_
   autoray.lazy.core.prod
   autoray.lazy.core.min_
   autoray.lazy.core.max_


Classes
-------

.. autoapisummary::

   autoray.lazy.core.Function
   autoray.lazy.core.Placeholder
   autoray.lazy.core.LazyArray


Functions
---------

.. autoapisummary::

   autoray.lazy.core.is_lazy_array
   autoray.lazy.core.to_queue
   autoray.lazy.core.descend
   autoray.lazy.core.ascend
   autoray.lazy.core.compute
   autoray.lazy.core.compute_constants
   autoray.lazy.core.get_source
   autoray.lazy.core.ensure_lazy
   autoray.lazy.core.find_lazy
   autoray.lazy.core.materialize_larray
   autoray.lazy.core.materialize_tuple
   autoray.lazy.core.materialize_list
   autoray.lazy.core.materialize_dict
   autoray.lazy.core.materialize_identity
   autoray.lazy.core.maybe_materialize
   autoray.lazy.core.stringify_larray
   autoray.lazy.core.stringify_tuple
   autoray.lazy.core.stringify_list
   autoray.lazy.core.stringify_dict
   autoray.lazy.core.stringify_identity
   autoray.lazy.core.stringify
   autoray.lazy.core.currently_sharing
   autoray.lazy.core.get_sharing_cache
   autoray.lazy.core._add_sharing_cache
   autoray.lazy.core._remove_sharing_cache
   autoray.lazy.core.shared_intermediates
   autoray.lazy.core.maybe_id
   autoray.lazy.core.hash_args_kwargs
   autoray.lazy.core.lazy_cache
   autoray.lazy.core.dtype_real_equiv
   autoray.lazy.core.dtype_complex_equiv
   autoray.lazy.core._find_common_dtype
   autoray.lazy.core.find_common_dtype
   autoray.lazy.core._find_common_backend_cached
   autoray.lazy.core.find_common_backend
   autoray.lazy.core.find_broadcast_shape
   autoray.lazy.core.Variable
   autoray.lazy.core.array
   autoray.lazy.core.transpose
   autoray.lazy.core._reshape_tuple
   autoray.lazy.core.find_full_reshape
   autoray.lazy.core.reshape
   autoray.lazy.core.expand_dims
   autoray.lazy.core.broadcast_to
   autoray.lazy.core.getitem_hasher
   autoray.lazy.core.get_sliced_size
   autoray.lazy.core.getitem
   autoray.lazy.core.tensordot
   autoray.lazy.core._basic_einsum_parse_input
   autoray.lazy.core._get_parse_einsum_input
   autoray.lazy.core.einsum
   autoray.lazy.core.trace
   autoray.lazy.core.diag
   autoray.lazy.core.diagonal
   autoray.lazy.core.matmul
   autoray.lazy.core.kron
   autoray.lazy.core.clip
   autoray.lazy.core.flip
   autoray.lazy.core.sort
   autoray.lazy.core.argsort
   autoray.lazy.core.stack
   autoray.lazy.core.concatenate
   autoray.lazy.core.split
   autoray.lazy.core.where
   autoray.lazy.core._get_py_shape
   autoray.lazy.core.take
   autoray.lazy.core.take_along_axis
   autoray.lazy.core.make_binary_func
   autoray.lazy.core.complex_
   autoray.lazy.core.make_unary_func
   autoray.lazy.core.make_reduction_func
   autoray.lazy.core.empty
   autoray.lazy.core.zeros
   autoray.lazy.core.ones
   autoray.lazy.core.eye
   autoray.lazy.core.identity
   autoray.lazy.core.lazy_get_dtype_name
   autoray.lazy.core.lazy_astype


Module Contents
---------------

.. py:data:: _EMPTY_DICT

.. py:data:: get_depth

.. py:data:: get_data

.. py:data:: get_dtype

.. py:function:: is_lazy_array(x)

   Check if ``x`` is a lazy array.


.. py:function:: to_queue(lz)

   Parse a pytree of lazy arrays into a queue of nodes, sorted by depth.
   This is useful for traversing the computational graph of multiple outputs
   in topological order.


.. py:function:: descend(lz)

   Generate each unique computational node. Use ``ascend`` if you need to
   visit children before parents.

   :param lz: The output node(s) of the computational graph to descend.
   :type lz: pytree of LazyArray

   :Yields: *LazyArray*


.. py:function:: ascend(lz)

   Generate each unique computational node, from leaves to root. I.e. a
   topological ordering of the computational graph. Moreover, the nodes
   are visited 'deepest first'.

   :param lz: The output node(s) of the computational graph to ascend to.
   :type lz: pytree of LazyArray

   :Yields: *LazyArray*


.. py:function:: compute(lz)

   Compute the value of one or more lazy arrays. All nodes that are
   computed clear any references to their function, arguments and
   dependencies, and store the result in their ``_data`` attribute.

   :param lz: The output node(s) of the computational graph to compute.
   :type lz: pytree of LazyArray

   :returns: The computed value(s) of the lazy array(s).
   :rtype: array or tuple of array


.. py:function:: compute_constants(lz, variables)

   Fold constant arrays - everything not dependent on ``variables`` -
   into the graph.

   :param lz: The output node(s) of the computational graph.
   :type lz: pytree of LazyArray
   :param variables: Nodes that should be treated as variable. I.e. any descendants will
                     not be folded into the graph.
   :type variables: pytree of LazyArray


.. py:function:: get_source(lz, params=None)

   Write the source code of an unravelled version of the computational
   graph, injecting required runtime objects into ``params``.

   :param lz: The output node(s) of the computational graph to write the source code
              for. Their corresponding label is ``f"x{id(node)}"`` in the
              source code.
   :type lz: LazyArray or sequence of LazyArray

   :returns: The source code of the computational graph, suitable for ``exec``.
   :rtype: str


.. py:class:: Function(inputs, outputs, fold_constants=True)

   Get a compiled (by python ``compile``), function that performs the
   computational graph corresponding to ``inputs`` -> ``outputs``. The
   signature of the function is ``func(input_arrays) -> output_arrays``. As an
   intermediate step, the computational graph is traced to a flattened source
   code string.

   :param inputs: The input node(s) of the computational graph.
   :type inputs: pytree of LazyArray
   :param outputs: The output node(s) of the computational graph.
   :type outputs: pytree of LazyArray
   :param fold_constants: If True, fold constant arrays (those with no dependence on ``inputs``)
                          into the graph ahead of compile.
   :type fold_constants: bool, optional

   .. seealso:: :py:obj:`get_source`, :py:obj:`compute`


   .. py:attribute:: __slots__
      :value: ('_in_names', '_out_names', '_out_tree', '_source', '_code', '_locals')



   .. py:attribute:: _locals


   .. py:attribute:: _source


   .. py:attribute:: _code


   .. py:attribute:: _in_names


   .. py:attribute:: _out_names


   .. py:method:: __call__(*args)


   .. py:method:: __getstate__()


   .. py:method:: __setstate__(state)


   .. py:method:: print_source()

      Print the source code of the compiled function.



   .. py:method:: __repr__()


.. py:class:: Placeholder

   A singleton object to use as a placeholder in a LazyArray.


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



   .. py:method:: __repr__()


.. py:data:: PLACEHOLDER

.. py:class:: LazyArray(backend, fn, args, kwargs, shape, deps=None)

   A lazy array representing a node in a computational graph.

   :param backend: The backend of the array were it to be computed. This can be ``None``
                   but this may cause problems when propagating information about which
                   functions to import to child nodes.
   :type backend: str
   :param fn: The function to call to compute the array, presumable imported from
              ``backend``. This can be ``None`` if the array already has data (e.g.
              is an input).
   :type fn: callable
   :param args: The positional arguments to pass to ``fn``, which might be
                ``LazyArray`` instances.
   :type args: tuple
   :param kwargs: The keyword arguments to pass to ``fn``, which might be
                  ``LazyArray`` instances.
   :type kwargs: dict
   :param shape: The shape of the array that ``fn(*args, **kwargs)`` will return, or
                 the shape of the array that ``data`` has.
   :type shape: tuple
   :param deps: The ``LazyArray`` instances that ``fn(*args, **kwargs)`` depends on.
                If not specified, these will be automatically found from ``args`` and
                ``kwargs``, specifying them manually is slightly more efficient.
   :type deps: tuple, optional


   .. py:attribute:: __slots__
      :value: ('_backend', '_fn', '_args', '_kwargs', '_shape', '_data', '_deps', '_depth')



   .. py:attribute:: _backend


   .. py:attribute:: _fn


   .. py:attribute:: _args


   .. py:attribute:: _shape


   .. py:attribute:: _data
      :value: None



   .. py:method:: from_data(data)
      :classmethod:


      Create a new ``LazyArray`` directly from a concrete array.



   .. py:method:: from_shape(shape, backend='numpy')
      :classmethod:


      Create a new ``LazyArray`` with a given shape.



   .. py:method:: to(fn, args=None, kwargs=None, backend=None, shape=None, deps=None)

      Create a new ``LazyArray``, by default propagating backend, shape,
      and deps from the the current LazyArray.



   .. py:method:: _materialize()

      Recursively compute all required args and kwargs for this node
      before computing itself and dereferencing dependencies. Note using this
      to materialize a large computation from scratch should be avoided due
      to the recursion limit, use ``x.compute()`` instead.



   .. py:attribute:: descend


   .. py:attribute:: ascend


   .. py:method:: compute()

      Compute the value of this lazy array, clearing any references to the
      function, arguments and dependencies, and storing the result in the
      ``_data`` attribute as well as returning it.

      Unlike ``self._materialize()`` this avoids deep recursion.



   .. py:attribute:: compute_constants


   .. py:method:: as_string(params)

      Create a string which evaluates to the lazy array creation.



   .. py:attribute:: get_source


   .. py:method:: get_function(variables, fold_constants=True)

      Get a compiled function that computes ``fn(arrays)``, with ``fn``
      describing the computational graph of this ``LazyArray`` and ``arrays``
      corresponding to the downstream ``LazyArray`` nodes ``variables``.

      :param variables: Input nodes whose data can change between calls.
      :type variables: sequence of LazyArray
      :param fold_constants: Compute all intermediates which do not depend on ``variables``
                             prior to compilation.
      :type fold_constants: bool, optional

      :returns: **fn** -- Function with signature ``fn(arrays)``.
      :rtype: callable



   .. py:method:: show(filler=' ', max_lines=None, max_depth=None)

      Show the computational graph as a nested directory structure.



   .. py:method:: history_num_nodes()

      Return the number of unique computational nodes in the history of
      this ``LazyArray``.



   .. py:method:: history_max_size()

      Get the largest single tensor size appearing in this computation.



   .. py:method:: history_size_footprint(include_inputs=True)

      Get the combined size of intermediates at each step of the
      computation. Note this assumes that intermediates are immediately
      garbage collected when they are no longer required.

      :param include_inputs: Whether to include the size of the inputs in the computation. If
                             ``True`` It is assumed they can be garbage collected once used but
                             are all present at the beginning of the computation.
      :type include_inputs: bool, optional



   .. py:method:: history_peak_size(include_inputs=True)

      Get the peak combined intermediate size of this computation.

      :param include_inputs: Whether to include the size of the inputs in the computation. If
                             ``True`` It is assumed they can be garbage collected once used but
                             are all present at the beginning of the computation.
      :type include_inputs: bool, optional



   .. py:method:: history_total_size()

      The the total size of all unique arrays in the computational graph,
      possibly relevant e.g. for back-propagation algorithms.



   .. py:method:: history_stats(fn)

      Compute aggregate statistics about the computational graph.

      :param fn: Function to apply to each node in the computational graph. If a
                 string, one of 'count', 'sizein', 'sizeout' can be used to count
                 the number of nodes, the total size of the inputs, or the total
                 size of each output respectively.
      :type fn: callable or str

      :returns: **stats** -- Dictionary mapping function names to the aggregate statistics.
      :rtype: dict



   .. py:method:: history_fn_frequencies()

      Get a dictionary mapping function names to the number of times they
      are used in the computational graph.



   .. py:method:: to_nx_digraph(variables=None)

      Convert this ``LazyArray`` into a ``networkx.DiGraph``.



   .. py:attribute:: plot


   .. py:attribute:: plot_graph


   .. py:attribute:: plot_circuit


   .. py:attribute:: plot_history_size_footprint


   .. py:attribute:: plot_history_functions


   .. py:attribute:: plot_history_functions_scatter


   .. py:attribute:: plot_history_functions_lines


   .. py:attribute:: plot_history_functions_image


   .. py:attribute:: plot_history_stats


   .. py:attribute:: plot_history_stats_counts


   .. py:attribute:: plot_history_stats_sizein


   .. py:property:: fn

      The function to use to compute this array.


   .. py:property:: fn_name

      The name of the function to use to compute this array.


   .. py:property:: args

      The positional arguments to the function to use to compute this
      array.


   .. py:property:: kwargs

      The keyword arguments to the function to use to compute this
      array.


   .. py:property:: shape


   .. py:method:: __len__()


   .. py:method:: __iter__()


   .. py:property:: ndim


   .. py:property:: size


   .. py:property:: backend


   .. py:property:: deps

      A tuple of the dependencies, other LazyArray instances, of this
      array.


   .. py:property:: depth

      The maximum distance to any input array in the computational graph.


   .. py:method:: __hash__()


   .. py:method:: __getitem__(key)


   .. py:attribute:: __array_ufunc__
      :value: None



   .. py:method:: __mul__(other)


   .. py:method:: __rmul__(other)


   .. py:method:: __add__(other)


   .. py:method:: __radd__(other)


   .. py:method:: __sub__(other)


   .. py:method:: __rsub__(other)


   .. py:method:: __floordiv__(other)


   .. py:method:: __rfloordiv__(other)


   .. py:method:: __truediv__(other)


   .. py:method:: __rtruediv__(other)


   .. py:method:: __mod__(other)


   .. py:method:: __pow__(other)


   .. py:method:: __rpow__(other)


   .. py:method:: __matmul__(other)


   .. py:method:: __rmatmul__(other)


   .. py:method:: __abs__()


   .. py:method:: __neg__()


   .. py:method:: __eq__(other)


   .. py:method:: __ne__(other)


   .. py:method:: __gt__(other)


   .. py:method:: __lt__(other)


   .. py:method:: __ge__(other)


   .. py:method:: __le__(other)


   .. py:property:: T


   .. py:property:: H


   .. py:method:: reshape(shape, **kwargs)


   .. py:method:: astype(dtype_name)


   .. py:property:: dtype


   .. py:property:: real


   .. py:property:: imag


   .. py:method:: __array_namespace__(api_version=None)

      Support array api namespace - https://data-apis.org/array-api/.



   .. py:method:: __repr__()


.. py:function:: ensure_lazy(array)

.. py:function:: find_lazy(x)

   Recursively search for ``LazyArray`` instances in pytrees.


.. py:function:: materialize_larray(x)

.. py:function:: materialize_tuple(x)

.. py:function:: materialize_list(x)

.. py:function:: materialize_dict(x)

.. py:function:: materialize_identity(x)

.. py:data:: _materialize_dispatch

.. py:function:: maybe_materialize(x)

   Recursively evaluate LazyArray instances in tuples, lists and dicts.


.. py:function:: stringify_larray(x, params)

.. py:function:: stringify_tuple(x, params)

.. py:function:: stringify_list(x, params)

.. py:function:: stringify_dict(x, params)

.. py:function:: stringify_identity(x, params)

.. py:data:: _stringify_dispatch

.. py:function:: stringify(x, params)

   Recursively stringify LazyArray instances in tuples, lists and dicts.


.. py:data:: _SHARING_STACK

.. py:function:: currently_sharing()

   Check if we are currently sharing a cache -- thread specific.


.. py:function:: get_sharing_cache()

   Return the most recent sharing cache -- thread specific.


.. py:function:: _add_sharing_cache(cache)

.. py:function:: _remove_sharing_cache()

.. py:function:: shared_intermediates(cache=None)

   Context in which intermediate results are shared.

   Note that intermediate LazyArray instances (which can reference actual
   data) will not be garbage collected until
   1. this context exits, and
   2. the yielded cache is garbage collected (if it was captured).

   :param cache: If specified, a user-stored dict in which intermediate lazy arrays will
                 be stored. This can be used to interleave sharing contexts.
   :type cache: dict

   :returns: **cache** -- A dictionary in which sharing results are stored. If ignored,
             sharing results will be garbage collected when this context is
             exited. This dict can be passed to another context to resume
             sharing.
   :rtype: dict


.. py:function:: maybe_id(x)

.. py:function:: hash_args_kwargs(fn_name, *args, **kwargs)

.. py:function:: lazy_cache(fn_name, hasher=None)

   Decorator to mark a function as being lazy cacheable.

   :param fn_name: The name to use for the function in the cache.
   :type fn_name: str
   :param hasher: A function with signature ``hasher(fn_name, *args, **kwargs)`` that
                  returns a hashable key for the cache. If not specified, the default
                  is to use ``hash_args_kwargs``.
   :type hasher: callable


.. py:data:: _DTYPES_REAL_EQUIV

.. py:data:: _DTYPES_COMPLEX_EQUIV

.. py:function:: dtype_real_equiv(dtype_name)

.. py:function:: dtype_complex_equiv(dtype_name)

.. py:function:: _find_common_dtype(array_types, scalar_types)

.. py:function:: find_common_dtype(*xs)

.. py:function:: _find_common_backend_cached(names)

.. py:function:: find_common_backend(*xs)

.. py:function:: find_broadcast_shape(xshape, yshape)

.. py:function:: Variable(shape, backend=None)

   Create a ``LazyArray`` from a shape only, representing a leaf node
   in the computational graph. It can only act as a placeholder for data.


.. py:function:: array(x)

   Create a ``LazyArray`` from an input array, representing a leaf node
   in the computational graph.


.. py:function:: transpose(a, axes=None)

.. py:data:: permute_dims

.. py:function:: _reshape_tuple(a, newshape, **kwargs)

.. py:function:: find_full_reshape(newshape, size)

.. py:function:: reshape(a, newshape, **kwargs)

.. py:function:: expand_dims(a, axis)

.. py:function:: broadcast_to(array, shape)

.. py:function:: getitem_hasher(_, a, key)

.. py:function:: get_sliced_size(d, start, stop, step)

.. py:function:: getitem(a, key)

.. py:function:: tensordot(a, b, axes=2)

.. py:function:: _basic_einsum_parse_input(operands)

.. py:function:: _get_parse_einsum_input()

.. py:function:: einsum(*operands)

.. py:function:: trace(a)

.. py:function:: diag(a, k=0)

.. py:function:: diagonal(a, offset=0, axis1=0, axis2=1)

.. py:function:: matmul(x1, x2)

.. py:function:: kron(x1, x2)

.. py:function:: clip(a, a_min, a_max)

.. py:function:: flip(a, axis=None)

.. py:function:: sort(a, axis=-1)

.. py:function:: argsort(a, axis=-1, **kwargs)

.. py:function:: stack(arrays, axis=0)

.. py:function:: concatenate(arrays, axis=0)

.. py:function:: split(ary, indices_or_sections, axis=0)

.. py:function:: where(condition, x, y)

.. py:function:: _get_py_shape(x)

   Infer the shape of a possibly nested list/tuple object.


.. py:function:: take(x, indices)

.. py:function:: take_along_axis(arr, indices, axis=-1)

.. py:function:: make_binary_func(name, fn)

.. py:data:: add

.. py:data:: eq

.. py:data:: equal

.. py:data:: floordivide

.. py:data:: ge

.. py:data:: gt

.. py:data:: le

.. py:data:: lt

.. py:data:: mod

.. py:data:: multiply

.. py:data:: ne

.. py:data:: pow_

.. py:data:: sub

.. py:data:: truedivide

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

.. py:function:: make_unary_func(name)

.. py:data:: sin

.. py:data:: cos

.. py:data:: tan

.. py:data:: arcsin

.. py:data:: arccos

.. py:data:: arctan

.. py:data:: sinh

.. py:data:: cosh

.. py:data:: tanh

.. py:data:: arcsinh

.. py:data:: arccosh

.. py:data:: arctanh

.. py:data:: sqrt

.. py:data:: exp

.. py:data:: log

.. py:data:: log2

.. py:data:: log10

.. py:data:: conj

.. py:data:: sign

.. py:data:: abs_

.. py:data:: angle

.. py:data:: real

.. py:data:: imag

.. py:function:: make_reduction_func(name)

.. py:data:: sum_

.. py:data:: prod

.. py:data:: min_

.. py:data:: max_

.. py:function:: empty(shape, *, backend='numpy', **kwargs)

   Lazy creation of an empty array with a given shape.


.. py:function:: zeros(shape, *, backend='numpy', **kwargs)

   Lazy creation of an array filled with zeros with a given shape.


.. py:function:: ones(shape, *, backend='numpy', **kwargs)

   Lazy creation of an array filled with ones with a given shape.


.. py:function:: eye(N, *, backend='numpy', **kwargs)

   Lazy creation of the identity matrix of size N.


.. py:function:: identity(n, *, backend='numpy', **kwargs)

   Lazy creation of the identity matrix of size n.


.. py:function:: lazy_get_dtype_name(x)

.. py:function:: lazy_astype(x, dtype_name)

