liesel.goose.pytree.split_and_transpose#

liesel.goose.pytree.split_and_transpose(pytree, axis=0)[source]#

Splits the leaves in a pytree along one axis and transposes the tree such that it’s a list of pytrees.

It assumes that all leaves have the same dimensionality along the chosen axis.