autoray.lazy

Submodules

• autoray.lazy.core
• autoray.lazy.draw
• autoray.lazy.linalg

Package Contents

Classes

Function	Get a compiled (by python compile), function that performs the
LazyArray	A lazy array representing a node in a computational graph.

Functions

ascend(lz)	Generate each unique computational node, from leaves to root. I.e. a
descend(lz)	Generate each unique computational node. Use ascend if you need to
compute(lz)	Compute the value of one or more lazy arrays. All nodes that are
get_source(lz[, params])	Write the source code of an unravelled version of the computational
Variable(shape[, backend])	Create a LazyArray from a shape only, representing a leaf node
shared_intermediates([cache])	Context in which intermediate results are shared.
array(x)	Create a LazyArray from an input array, representing a leaf node
transpose(a[, axes])
reshape(a, newshape)
tensordot(a, b[, axes])
einsum(*operands)
trace(a)
diag(a[, k])
matmul(x1, x2)
kron(x1, x2)
clip(a, a_min, a_max)
flip(a[, axis])
sort(a[, axis])
argsort(a[, axis])
stack(arrays[, axis])
concatenate(arrays[, axis])
split(ary, indices_or_sections[, axis])
where(condition, x, y)
take(x, indices)
complex(re, im)

Attributes

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

autoray.lazy.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’.

Parameters:

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

Yields:

LazyArray

autoray.lazy.descend(lz)

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

Parameters:

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

Yields:

LazyArray

autoray.lazy.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.

Parameters:

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

Returns:

The computed value(s) of the lazy array(s).

Return type:

array or tuple of array

autoray.lazy.get_source(lz, params=None)

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

Parameters:

lz (LazyArray or sequence of LazyArray) – 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.

Returns:

The source code of the computational graph, suitable for exec.

Return type:

str

class autoray.lazy.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.

Parameters:

• inputs (pytree of LazyArray) – The input node(s) of the computational graph.

• outputs (pytree of LazyArray) – The output node(s) of the computational graph.

• fold_constants (bool, optional) – If True, fold constant arrays (those with no dependence on inputs) into the graph ahead of compile.

See also

get_source, compute

__slots__ = ('_in_names', '_out_names', '_out_tree', '_source', '_code', '_locals')

__call__(*args)

__getstate__()

Helper for pickle.

__setstate__(state)

print_source()

Print the source code of the compiled function.

__repr__()

Return repr(self).

class autoray.lazy.LazyArray(backend, fn, args, kwargs, shape, deps=None)

A lazy array representing a node in a computational graph.

Parameters:

• backend (str) – 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.

• fn (callable) – 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).

• args (tuple) – The positional arguments to pass to fn, which might be LazyArray instances.

• kwargs (dict) – The keyword arguments to pass to fn, which might be LazyArray instances.

• shape (tuple) – The shape of the array that fn(*args, **kwargs) will return, or the shape of the array that data has.

• deps (tuple, optional) – 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.

property fn

The function to use to compute this array.

property fn_name

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

property args

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

property kwargs

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

property shape

property ndim

property size

property backend

property deps

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

property depth

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

property T

property H

property real

property imag

__slots__ = ('_backend', '_fn', '_args', '_kwargs', '_shape', '_data', '_deps', '_depth')

descend

ascend

compute_constants

get_source

plot

plot_graph

plot_circuit

plot_history_size_footprint

plot_history_functions

plot_history_functions_scatter

plot_history_functions_lines

plot_history_functions_image

plot_history_stats

plot_history_stats_counts

plot_history_stats_sizein

__array_ufunc__

classmethod from_data(data)

Create a new LazyArray directly from a concrete array.

classmethod from_shape(shape, backend='numpy')

Create a new LazyArray with a given shape.

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.

_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.

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.

as_string(params)

Create a string which evaluates to the lazy array creation.

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.

Parameters:

• variables (sequence of LazyArray) – Input nodes whose data can change between calls.

• fold_constants (bool, optional) – Compute all intermediates which do not depend on variables prior to compilation.

Returns:

fn – Function with signature fn(arrays).

Return type:

callable

show(filler=' ', max_lines=None, max_depth=None)

Show the computational graph as a nested directory structure.

history_num_nodes()

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

history_max_size()

Get the largest single tensor size appearing in this computation.

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.

Parameters:

include_inputs (bool, optional) – 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.

history_peak_size(include_inputs=True)

Get the peak combined intermediate size of this computation.

Parameters:

include_inputs (bool, optional) – 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.

history_total_size()

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

history_stats(fn)

Compute aggregate statistics about the computational graph.

Parameters:

fn (callable or str) – 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.

Returns:

stats – Dictionary mapping function names to the aggregate statistics.

Return type:

dict

history_fn_frequencies()

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

to_nx_digraph(variables=None)

Convert this LazyArray into a networkx.DiGraph.

__len__()

__iter__()

__getitem__(key)

__mul__(other)

__rmul__(other)

__add__(other)

__radd__(other)

__sub__(other)

__rsub__(other)

__floordiv__(other)

__rfloordiv__(other)

__truediv__(other)

__rtruediv__(other)

__pow__(other)

__rpow__(other)

__matmul__(other)

__rmatmul__(other)

__abs__()

__neg__()

__ne__(other)

Return self!=value.

__gt__(other)

Return self>value.

__lt__(other)

Return self<value.

__ge__(other)

Return self>=value.

__le__(other)

Return self<=value.

reshape(shape)

astype(dtype_name)

__repr__()

Return repr(self).

autoray.lazy.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.

autoray.lazy.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).

Parameters:

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

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.

Return type:

dict

autoray.lazy.array(x)

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

autoray.lazy.transpose(a, axes=None)

autoray.lazy.reshape(a, newshape)

autoray.lazy.tensordot(a, b, axes=2)

autoray.lazy.einsum(*operands)

autoray.lazy.trace(a)

autoray.lazy.diag(a, k=0)

autoray.lazy.matmul(x1, x2)

autoray.lazy.kron(x1, x2)

autoray.lazy.clip(a, a_min, a_max)

autoray.lazy.flip(a, axis=None)

autoray.lazy.sort(a, axis=-1)

autoray.lazy.argsort(a, axis=-1)

autoray.lazy.stack(arrays, axis=0)

autoray.lazy.concatenate(arrays, axis=0)

autoray.lazy.split(ary, indices_or_sections, axis=0)

autoray.lazy.where(condition, x, y)

autoray.lazy.take(x, indices)

autoray.lazy.multiply

autoray.lazy.add

autoray.lazy.floordivide

autoray.lazy.truedivide

autoray.lazy.sin

autoray.lazy.cos

autoray.lazy.tan

autoray.lazy.arcsin

autoray.lazy.arccos

autoray.lazy.arctan

autoray.lazy.sinh

autoray.lazy.cosh

autoray.lazy.tanh

autoray.lazy.arcsinh

autoray.lazy.arccosh

autoray.lazy.arctanh

autoray.lazy.sqrt

autoray.lazy.exp

autoray.lazy.log

autoray.lazy.log2

autoray.lazy.log10

autoray.lazy.conj

autoray.lazy.sign

autoray.lazy.angle

autoray.lazy.real

autoray.lazy.imag

autoray.lazy.prod

autoray.lazy.abs

autoray.lazy.sum

autoray.lazy.min

autoray.lazy.max

autoray.lazy.complex(re, im)