autoray.experimental.complexity_tracing¶

Functionality for tracing through an autoray.lazy computation and estimating the cost and scaling.

In the following there are cost_* functions that estimate the total cost of a given operation, including sub-leading factors. There are also cost_scaling_* functions that only consider the leading factor of the cost, so that we can prime number decompose it and extract the scaling.

Attributes¶

`cost_scaling_tensordot`
`cost_scaling_matmul`
`cost_scaling_einsum`
`_LINEAR_COSTS`
`_NOTHING_COSTS`
`COSTS`
`COST_SCALINGS`

Functions¶

`_get_batch_matrix_sizes`(x)	Get the batch size and ordered matrix dimensions of `x`.
`cost_tensordot`(x)
`cost_qr`(x)
`cost_svd`(x)
`cost_eigh`(x)
`cost_scaling_linalg`(x)	Here we only care about the leading factor of the cost, which we need to
`cost_matmul`(x)
`cost_einsum`(x)
`cost_linear`(x)
`cost_nothing`(x)
`cost_node`(x[, allow_missed])
`compute_cost`(z[, print_missed, allow_missed])	Estimate the total cost of one or more lazy output nodes.
`prime_factors`(→ list[int])
`is_prime`(→ bool)
`closest_prime`(→ int)
`frequencies`(it)
`_check_factor_map`(factor_map)
`compute_cost_scalings`(z, factor_map[, print_missed, ...])	Estimate cost scalings for one or more lazy output nodes.

Module Contents¶

autoray.experimental.complexity_tracing._get_batch_matrix_sizes(x)¶: Get the batch size and ordered matrix dimensions of x.

autoray.experimental.complexity_tracing.cost_tensordot(x)¶

autoray.experimental.complexity_tracing.cost_scaling_tensordot¶

autoray.experimental.complexity_tracing.cost_qr(x)¶

autoray.experimental.complexity_tracing.cost_svd(x)¶

autoray.experimental.complexity_tracing.cost_eigh(x)¶

autoray.experimental.complexity_tracing.cost_scaling_linalg(x)¶: Here we only care about the leading factor of the cost, which we need to preserve so that we can prime number decompose it.

autoray.experimental.complexity_tracing.cost_matmul(x)¶

autoray.experimental.complexity_tracing.cost_scaling_matmul¶

autoray.experimental.complexity_tracing.cost_einsum(x)¶

autoray.experimental.complexity_tracing.cost_scaling_einsum¶

autoray.experimental.complexity_tracing.cost_linear(x)¶

autoray.experimental.complexity_tracing.cost_nothing(x)¶

autoray.experimental.complexity_tracing._LINEAR_COSTS¶

autoray.experimental.complexity_tracing._NOTHING_COSTS¶

autoray.experimental.complexity_tracing.COSTS¶

autoray.experimental.complexity_tracing.cost_node(x, allow_missed=True)¶

autoray.experimental.complexity_tracing.compute_cost(z, print_missed=True, allow_missed=True)¶

Estimate the total cost of one or more lazy output nodes.

Shared dependencies of multiple output nodes are counted once.

Parameters:

z (pytree of LazyArray) – The output node or nodes to trace.
print_missed (bool, optional) – Whether to warn about operations without a registered cost.
allow_missed (bool, optional) – Whether to omit operations without a registered cost. If False, raise a ValueError listing them instead.

autoray.experimental.complexity_tracing.COST_SCALINGS¶

autoray.experimental.complexity_tracing.prime_factors(n) → list[int]¶

autoray.experimental.complexity_tracing.is_prime(n: int) → bool¶

autoray.experimental.complexity_tracing.closest_prime(nt: int) → int¶

autoray.experimental.complexity_tracing.frequencies(it)¶

autoray.experimental.complexity_tracing._check_factor_map(factor_map)¶

autoray.experimental.complexity_tracing.compute_cost_scalings(z, factor_map, print_missed=True, allow_missed=True)¶

Estimate cost scalings for one or more lazy output nodes.

Parameters:

z (pytree of LazyArray) – The output node or nodes to trace. Shared dependencies are counted once.
factor_map (dict[str, int]) – Mapping from dimension labels to distinct prime numbers.
print_missed (bool, optional) – Whether to warn about unregistered operations and prime factors.
allow_missed (bool, optional) – Whether to omit operations without a registered scaling. If False, raise a ValueError listing them instead.