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

.. py:module:: autoray.lazy


Submodules
----------

.. toctree::
   :maxdepth: 1

   /autoapi/autoray/lazy/core/index
   /autoapi/autoray/lazy/draw/index
   /autoapi/autoray/lazy/linalg/index


Attributes
----------

.. autoapisummary::

   autoray.lazy.add
   autoray.lazy.angle
   autoray.lazy.arccos
   autoray.lazy.arccosh
   autoray.lazy.arcsin
   autoray.lazy.arcsinh
   autoray.lazy.arctan
   autoray.lazy.arctanh
   autoray.lazy.conj
   autoray.lazy.cos
   autoray.lazy.cosh
   autoray.lazy.equal
   autoray.lazy.exp
   autoray.lazy.floordivide
   autoray.lazy.imag
   autoray.lazy.log
   autoray.lazy.log2
   autoray.lazy.log10
   autoray.lazy.multiply
   autoray.lazy.permute_dims
   autoray.lazy.prod
   autoray.lazy.real
   autoray.lazy.sign
   autoray.lazy.sin
   autoray.lazy.sinh
   autoray.lazy.sqrt
   autoray.lazy.tan
   autoray.lazy.tanh
   autoray.lazy.truedivide
   autoray.lazy.abs
   autoray.lazy.max
   autoray.lazy.min
   autoray.lazy.sum


Classes
-------

.. autoapisummary::

   autoray.lazy.Function
   autoray.lazy.LazyArray


Functions
---------

.. autoapisummary::

   autoray.lazy.Variable
   autoray.lazy.argsort
   autoray.lazy.array
   autoray.lazy.ascend
   autoray.lazy.broadcast_to
   autoray.lazy.clip
   autoray.lazy.compute
   autoray.lazy.concatenate
   autoray.lazy.descend
   autoray.lazy.diag
   autoray.lazy.diagonal
   autoray.lazy.einsum
   autoray.lazy.empty
   autoray.lazy.expand_dims
   autoray.lazy.eye
   autoray.lazy.flip
   autoray.lazy.get_source
   autoray.lazy.identity
   autoray.lazy.kron
   autoray.lazy.matmul
   autoray.lazy.ones
   autoray.lazy.reshape
   autoray.lazy.shared_intermediates
   autoray.lazy.sort
   autoray.lazy.split
   autoray.lazy.stack
   autoray.lazy.take
   autoray.lazy.take_along_axis
   autoray.lazy.tensordot
   autoray.lazy.trace
   autoray.lazy.transpose
   autoray.lazy.where
   autoray.lazy.zeros
   autoray.lazy.complex


Package Contents
----------------

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

.. py:data:: angle

.. py:data:: arccos

.. py:data:: arccosh

.. py:data:: arcsin

.. py:data:: arcsinh

.. py:data:: arctan

.. py:data:: arctanh

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

.. py:function:: array(x)

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


.. 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:: broadcast_to(array, shape)

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

.. 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:: concatenate(arrays, axis=0)

.. py:data:: conj

.. py:data:: cos

.. py:data:: cosh

.. 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:: diag(a, k=0)

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

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

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

   Lazy creation of an empty array with a given shape.


.. py:data:: equal

.. py:data:: exp

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

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

   Lazy creation of the identity matrix of size N.


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

.. py:data:: floordivide

.. 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:function:: identity(n, *, backend='numpy', **kwargs)

   Lazy creation of the identity matrix of size n.


.. py:data:: imag

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

.. py:data:: log

.. py:data:: log2

.. py:data:: log10

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

.. py:data:: multiply

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

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


.. py:data:: permute_dims

.. py:data:: prod

.. py:data:: real

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

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

.. py:data:: sin

.. py:data:: sinh

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

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

.. py:data:: sqrt

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

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

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

.. py:data:: tan

.. py:data:: tanh

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

.. py:function:: trace(a)

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

.. py:data:: truedivide

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

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

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


.. py:data:: abs

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

.. py:data:: max

.. py:data:: min

.. py:data:: sum

