etree.stack

etils.etree.stack(

trees: Iterable[etils.enp.array_types.typing.Array | list[etils.etree.typing._T | list[Tree] | tuple[Tree, ...] | dict[str, Tree]] | tuple[etils.etree.typing._T | list[Tree] | tuple[Tree, ...] | dict[str, Tree], ...] | dict[str, etils.etree.typing._T | list[Tree] | tuple[Tree, ...] | dict[str, Tree]]],

) → etils.enp.array_types.typing.Array | list[etils.etree.typing._T | list[Tree] | tuple[Tree, ...] | dict[str, Tree]] | tuple[etils.etree.typing._T | list[Tree] | tuple[Tree, ...] | dict[str, Tree], ...] | dict[str, etils.etree.typing._T | list[Tree] | tuple[Tree, ...] | dict[str, Tree]]

Stack a tree of Iterable[Array].

Supports jax, tf, np.

Example:

etree.stack([
    {'a': np.array([1])},
    {'a': np.array([2])},
    {'a': np.array([3])},
]) == {
    'a': np.array([[1], [2], [3]])
}

Parameters:

trees – The list of tree to stack

Returns:

Tree of arrays.