jax_sn.tools.memory_analysis module

Tools for analyzing memory usage and timing of JAX functions, particularly for linear operators parameterized by certain variables.

pytree_nbytes(pytree)

Computes the total number of bytes used by all array leaves in a pytree.

Parameters:

pytreeAny

A pytree containing arrays.

Returns:

: int

The total number of bytes used by all array leaves in the pytree.

make_memory_parametrised_grad(method_or_f)

custom_timeit(func, number=100, ignore_n: int = 0)

class OperatorAnalysis

Bases: ABC

Class to simplify analysis of memory usage and timing of linear operators. Provides static methods to analyze memory and timing for operator methods (e.g. mv, inv).

static analyze_memory(operator: AbstractLinearOperator, method: Union[str, Callable])

Analyzes memory usage for a operator method (mv or inv). First it initializes the operator with the provided arguments and keyword arguments.

It uses _make_op to create a function that applies the specified method of the operator some input (just ones, for memory it does not matter).

This function is then lowered and compiled and its memory usage is analyzed using _filter_jit_f_to_memory.

Parameters:

• operator (AbstractLinearOperator) – The linear operator to analyze.

• method (Union[str, Callable]) – The method of the operator to use, e.g. “mv” or “inv”.

Returns:

mem_mv

static analyze_timing(operator: AbstractLinearOperator, method: Union[str, Callable], n_samples: int = 100, ignore_n: int = 1)

Analyzes timing for a operator method.

It uses _make_op to create a function that applies the specified method of the operator some input (just ones, for timing it does not matter).

The function is then timed using custom_timeit, and the result is returned.

Parameters:

• operator (AbstractLinearOperator) – The linear operator to analyze.

• method (Union[str, Callable]) – The method of the operator to use, e.g. “mv” or “inv”.

• n_samples (int) – The number of samples to use for timing, by default 100.

• ignore_n (int) – The number of initial runs to ignore for timing (e.g. for compilation caching), by default 1.

Returns:

duration_mv

analyze_grad_memory(cls: type, operator: AbstractLinearOperator, method: Union[str, Callable])

Analyzes memory usage for the gradient of a parameterised operator method (mv or inv). First it initializes the operator with the provided arguments and keyword arguments.

It uses _make_theta_op to create a function that applies the specified method of the operator some input (just ones, for memory it does not matter) and sums the output. This gives a scalar output, allowing the use of jax.value_and_grad to compute the gradient with respect to the parameter theta.

This function is then lowered and compiled and its memory usage is analyzed using _filter_jit_f_to_memory. This is finally printed.

Parameters:

• cls (type) – The parametrisation case class, providing theta_base and reconstruct_operator methods.

• operator (AbstractLinearOperator) – The linear operator to analyze.

• method (Union[str, Callable]) – The method of the operator to use, e.g. “mv” or “inv”.

Returns:

mem_mv, mem_grad_mv mem_mv : Dict[str, float]

Memory usage statistics for the operator method [in MB].

mem_grad_mvDict[str, float]

Memory usage statistics for the gradient of the operator method [in MB].

analyze_grad_timing(cls: type, operator: AbstractLinearOperator, method: Union[str, Callable], n_samples: int = 100, ignore_n: int = 1)

Analyzes timing for the gradient of a parameterised operator method.

It uses _make_theta_op to create a function that applies the specified method of the operator some input (just ones, for timing it does not matter) and sums the output. This gives a scalar output, allowing the use of jax.value_and_grad to compute the gradient with respect to the parameter theta.

The function and its gradient are then timed using custom_timeit, and the results are returned.

Parameters:

• cls (type) – The parametrisation case class, providing theta_base and reconstruct_operator methods.

• operator (AbstractLinearOperator) – The linear operator to analyze.

• method (Union[str, Callable]) – The method of the operator to use, e.g. “mv” or “inv”.

• n_samples (int) – The number of samples to use for timing, by default 100.

• ignore_n (int) – The number of initial runs to ignore for timing (e.g. for compilation caching), by default 1.

Returns:

duration_mv, duration_grad_mv duration_mv : float

The average time taken for the operator method.

duration_grad_mvfloat

The average time taken for the gradient of the operator method.

assert_pytree_allclose(a, b, *, rtol=1e-12, atol=1e-12)