liesel.goose.pytree#
Pytree utilities.
Functions
|
Converts every leaf in a pytree to a non-weak |
|
Concatenates all leaves in the list of pytrees along the given axis. |
Decorator for registering dataclasses as pytrees. |
|
|
Performs the same slice operation on every leaf. |
|
Splits the leaves in a pytree along one axis and transposes the tree such that it's a list of pytrees. |
|
Splits all leaves in a pytree into multiple sub-arrays. |
|
Squeezes all leaves in a pytree. |
|
Stacks all leaves in the list of pytrees along the given axis. |