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.moveaxis autoray.lazy.core.swapaxes 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.allclose autoray.lazy.core.parse_creation_shape 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 :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:: 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:: 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:: 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:: moveaxis(a, source, destination) .. py:function:: swapaxes(a, axis1, axis2) .. 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, **kwargs) .. 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, var_name=None) .. 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, var_name=None) .. py:data:: sum_ .. py:data:: prod .. py:data:: min_ .. py:data:: max_ .. py:function:: allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False) .. py:function:: parse_creation_shape(shape) .. 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)