StatefulFunction

class braincore.transform.StatefulFunction(fun, static_argnums=(), axis_env=None, abstracted_axes=None, state_returns=('read', 'write'), cache_type=None)

A wrapper class for a function that collects the states that are read and written by the function. The states are collected by the function and returned as a StateDictManager instance. The StateDictManager instance can be used to manage the states in the JAX program. The class provides a function called states that returns the states that are read and written by the function. The class provides a function called to_state_manager that returns a StateDictManager instance that contains the states that are read and written by the function. The class provides a function called __call__ that wraps the function and returns the states that are read and written by the function and the output of the function.

Parameters:

• fun (Callable) – The function whose jaxpr is to be computed. Its positional arguments and return value should be arrays, scalars, or standard Python containers (tuple/list/dict) thereof.

• static_argnums (Union[int, Iterable[int]]) – See the jax.jit() docstring.

• axis_env (Optional[Sequence[tuple[Hashable, int]]]) – Optional, a sequence of pairs where the first element is an axis name and the second element is a positive integer representing the size of the mapped axis with that name. This parameter is useful when lowering functions that involve parallel communication collectives, and it specifies the axis name/size environment that would be set up by applications of jax.pmap().

• abstracted_axes (Optional[Any]) – Optional, a pytree with the same structure as the input arguments to fun. The leaves of the pytree can be either None or a dict with axis names as keys and integers as values. If the leaf is None, then the corresponding axis is not abstracted. If the leaf is a dict, then the corresponding axis is abstracted, and the dict specifies the axis name and size. The abstracted axes are used to infer the input type of the function. If None, then all axes are abstracted.

• state_returns (Union[str, Tuple[str, ...]]) – Optional, a string or a tuple of strings. The default is ('read', 'write'). The strings specify the categories of states to be returned by the wrapped function. The categories are 'read' and 'write'. If the category is 'read', then the wrapped function returns the states that are read by the function. If the category is 'write', then the wrapped function returns the states that are written by the function. If the category is 'read' and 'write', then the wrapped function returns both the read and write states.

clear_cache()

Clear the compilation cache.

Return type:

None

compile_and_get_states_by_static_args(*args, **kwargs)

Get the states that are read and written by the function.

Parameters:

• *args – The arguments to the function.

• **kwargs – The keyword arguments to the function.

Return type:

Tuple[State, ...]

Returns:

The states that are read and written by the function.

get_arg_cache_key(*args, **kwargs)

Get the static arguments from the arguments.

Parameters:

*args – The arguments to the function.

Return type:

Tuple

Returns:

The static arguments.

get_jaxpr(cache_key=())

Read the JAX Jaxpr representation of the function.

Parameters:

cache_key (Hashable) – The hashable key.

Return type:

ClosedJaxpr

Returns:

The JAX Jaxpr representation of the function.

get_out_shapes(cache_key=())

Read the output shapes of the function.

Parameters:

cache_key (Hashable) – The hashable key.

Return type:

Any

Returns:

The output shapes of the function.

get_out_treedef(cache_key=())

Read the output tree of the function.

Parameters:

cache_key (Hashable) – The hashable key.

Return type:

Any

Returns:

The output tree of the function.

get_read_states(cache_key=())

Read the states that are read by the function.

Parameters:

cache_key (Hashable) – The hashable key.

Return type:

Tuple[State, ...]

Returns:

The states that are read by the function.

get_states(cache_key=())

Read the states that are read and written by the function.

Parameters:

cache_key (Hashable) – The hashable key.

Return type:

Tuple[State, ...]

Returns:

The states that are read and written by the function.

get_write_states(cache_key=())

Read the states that are written by the function.

Parameters:

cache_key (Hashable) – The hashable key.

Return type:

Tuple[State, ...]

Returns:

The states that are written by the function.

jaxpr_call(state_vals, *args, **kwargs)

Call the function at the JAX Jaxpr level.

Parameters:

• state_vals – The state values.

• *args – The arguments to the function.

• **kwargs – The keyword arguments to the function.

Return type:

Any

Returns:

State values and the function output.

jaxpr_call_auto(*args, **kwargs)

Call the function at the JAX Jaxpr level with automatic state management.

Parameters:

• *args – The arguments to the function.

• **kwargs – The keyword arguments to the function.

Return type:

Any

Returns:

The output of the function.

make_jaxpr(*args, **kwargs)

Creates a function that produces its jaxpr given example args.

A ClosedJaxpr representation of fun on those arguments. If the argument return_shape is True, then the returned function instead returns a pair where the first element is the ClosedJaxpr representation of fun and the second element is a pytree representing the structure, shape, dtypes, and named shapes of the output of fun.

Parameters:

• *args – The arguments to the function.

• **kwargs – The keyword arguments to the function.