pybaum is available on PyPI and Anaconda.org. Install it with
$ pip install pybaum
# or
$ conda install -c conda-forge pybaumpybaum provides tools to work with pytrees which is a concept borrowed from JAX.
What is a pytree?
In pybaum, we use the term pytree to refer to a tree-like structure built out of container-like Python objects. Classes are considered container-like if they are in the pytree registry, which by default includes lists, tuples, and dicts. That is:
- Any object whose type is not in the pytree container registry is considered a leaf pytree.
- Any object whose type is in the pytree container registry, and which contains pytrees, is considered a pytree.
For each entry in the pytree container registry, a container-like type is registered with a pair of functions that specify how to convert an instance of the container type to a (children, metadata) pair and how to convert such a pair back to an instance of the container type. Using these functions, pybaum can canonicalize any tree of registered container objects into tuples.
Example pytrees:
[1, "a", object()] # 3 leaves
(1, (2, 3), ()) # 3 leaves
[1, {"k1": 2, "k2": (3, 4)}, 5] # 5 leavespybaum can be extended to consider other container types as pytrees.