nifty8.re.tree_math.forest_math module

assert_arithmetics(obj, *args, **kwargs)

get_map(map) → Callable

has_arithmetics(obj, additional_attributes=())

map_forest(f: Callable, in_axes: int | Tuple = 0, out_axes: int | Tuple = 0, tree_transpose_output: bool = True, map: str | Callable = 'vmap', **kwargs) → Callable

map_forest_mean(method, map='vmap', *args, **kwargs) → Callable

mean(forest)

mean_and_std(forest, correct_bias=True)

random_like(key: ~collections.abc.Iterable, primals, rng: ~typing.Callable = <function normal>)

stack(arrays, axis=0)

tree_shape(tree: T) → T

unite(x, y, op=<built-in function add>)

Unites two Vector-like objects.

If a key is contained in both objects, then the fields at that key are combined.

unstack(stack, axis=0)