maps

jaxns.internals.maps

Module Contents

logger

replace_index(operand, update, start_index)

Replaces an index or slice with an update. If update is too big to respect start_index then start_index is shifted, which will give non-intuitive results.

get_index(operand, start_index, length)

prepare_func_args(f)

Takes a callable(a,b,…,z=Z) and prepares it into callable(**kwargs), such that only a,b,…,z are taken from **kwargs and the rest ignored.

This allows f(**kwarg) to work even if f() is missing some keys from kwargs.

Parameters:

f – callable(a,b,...,z=Z)

Returns:

callable(**kwargs) where **kwargs are the filtered for args of the original function.

F

FV

chunked_pmap(f, chunk_size=None, unroll=1)

A version of pmap which chunks the input into smaller pieces to avoid memory issues.

Parameters:

• f (Callable[Ellipsis, FV]) – callable

• chunk_size (Optional[int]) – the size of the chunks. Default is len(devices())

• unroll (int) – the number of times to unroll the computation

Returns:

a chunked version of f

Return type:

Callable[Ellipsis, FV]

prepad(a, chunksize)

Parameters:

chunksize (int) –

T

remove_chunk_dim(py_tree)

Remove the chunk dimension from a pytree

Parameters:

py_tree (T) – pytree to remove chunk dimension from

Returns:

pytree with chunk dimension removed

Return type:

T

add_chunk_dim(py_tree, chunk_size)

Add a chunk dimension to a pytree

Parameters:

• py_tree (T) – pytree to add chunk dimension to

• chunk_size (int) – size of chunk dimension

Returns:

pytree with chunk dimension added

Return type:

T

chunked_vmap(f, chunk_size=None, unroll=1)

A version of vmap which chunks the input into smaller pieces to avoid memory issues.

Parameters:

• f – the function to be mapped

• chunk_size (Optional[int]) – the size of the chunks. Default is len(devices())

• unroll (int) – the number of times to unroll the computation

Returns: