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.allclose 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.moveaxis autoray.lazy.ones autoray.lazy.reshape autoray.lazy.shared_intermediates autoray.lazy.sort autoray.lazy.split autoray.lazy.stack autoray.lazy.swapaxes 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 :value: '' .. 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:: as_string_pretty() Render this LazyArray as a nice string, for display rather than evaluation purposes. .. py:method:: show(filler=' ', max_lines=None, max_depth=None, show_args=True) Show the computational graph as a nested directory structure. :param filler: The string to use to fill in the empty space on the left. :type filler: str, optional :param max_lines: The maximum number of lines to show. By default all lines are shown. :type max_lines: int, optional :param max_depth: The maximum depth to show. By default all depths are shown. :type max_depth: int, optional :param show_args: Whether to show the arguments to each function in the graph. :type show_args: bool, optional .. 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) .. 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:function:: allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False) .. 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, **kwargs) .. 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:function:: moveaxis(a, source, destination) .. 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:: swapaxes(a, axis1, axis2) .. 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