:py:mod:`pybaum` ================ .. py:module:: pybaum Submodules ---------- .. toctree:: :titlesonly: :maxdepth: 1 _version/index.rst config/index.rst equality/index.rst registry/index.rst registry_entries/index.rst tree_util/index.rst typecheck/index.rst Package Contents ---------------- Functions ~~~~~~~~~ .. autoapisummary:: pybaum.get_registry pybaum.leaf_names pybaum.tree_equal pybaum.tree_flatten pybaum.tree_just_flatten pybaum.tree_just_yield pybaum.tree_map pybaum.tree_multimap pybaum.tree_unflatten pybaum.tree_update pybaum.tree_yield .. py:function:: get_registry(types=None, include_defaults=True) Create a pytree registry. :param types: A list strings with the names of types that should be included in the registry, i.e. considered containers and not leaves by the functions that work with pytrees. Currently we support: - "tuple" - "dict" - "list" - :class:`collections.namedtuple` or :class:`typing.NamedTuple` - :obj:`None` - :class:`collections.OrderedDict` - "numpy.ndarray" - "jax.numpy.ndarray" - "pandas.Series" - "pandas.DataFrame" :type types: list :param include_defaults: Whether the default pytree containers "tuple", "dict" "list", "None", "namedtuple" and "OrderedDict" should be included even if not specified in `types`. :type include_defaults: bool :returns: A pytree registry. :rtype: dict .. py:function:: leaf_names(tree, is_leaf=None, registry=None, separator='_') Construct names for leaves in a pytree. :param tree: a pytree to flatten. :param is_leaf: An optionally specified function that will be called at each flattening step. It should return a boolean, which indicates whether the flattening should traverse the current object, or if it should be stopped immediately, with the whole subtree being treated as a leaf. :type is_leaf: callable or None :param registry: A pytree container registry that determines which types are considered container objects that should be flattened. `is_leaf` can override this in the sense that types that are in the registry are still considered a leaf but it cannot declare something a container that is not in the registry. None means that the default registry is used, i.e. that dicts, tuples and lists are considered containers. "extended" means that in addition numpy arrays and params DataFrames are considered containers. Passing a dictionary where the keys are types and the values are dicts with the entries "flatten", "unflatten" and "names" allows to completely override the default registries. :type registry: dict or None :param separator: String that separates the building blocks of the leaf name. :type separator: str :returns: List of strings with names for pytree leaves. :rtype: list .. py:function:: tree_equal(tree, other, is_leaf=None, registry=None, equality_checkers=None) Determine if two pytrees are equal. Two pytrees are considered equal if their leaves are equal and the names of their leaves are equal. While this definition of equality might not always make sense it makes sense in most cases and can be implemented relatively easily. :param tree: A pytree. :param other: Another pytree. :param is_leaf: An optionally specified function that will be called at each flattening step. It should return a boolean, which indicates whether the flattening should traverse the current object, or if it should be stopped immediately, with the whole subtree being treated as a leaf. :type is_leaf: callable or None :param registry: A pytree container registry that determines which types are considered container objects that should be flattened. `is_leaf` can override this in the sense that types that are in the registry are still considered a leaf but it cannot declare something a container that is not in the registry. None means that the default registry is used, i.e. that dicts, tuples and lists are considered containers. "extended" means that in addition numpy arrays and params DataFrames are considered containers. Passing a dictionary where the keys are types and the values are dicts with the entries "flatten", "unflatten" and "names" allows to completely override the default registries. :type registry: dict or None :param equality_checkers: A dictionary where keys are types and values are functions which assess equality for the type of object. :type equality_checkers: dict, None :returns: bool .. py:function:: tree_flatten(tree, is_leaf=None, registry=None) Flatten a pytree and create a treedef. :param tree: a pytree to flatten. :param is_leaf: An optionally specified function that will be called at each flattening step. It should return a boolean, which indicates whether the flattening should traverse the current object, or if it should be stopped immediately, with the whole subtree being treated as a leaf. :type is_leaf: callable or None :param registry: A pytree container registry that determines which types are considered container objects that should be flattened. ``is_leaf`` can override this in the sense that types that are in the registry are still considered a leaf but it cannot declare something a container that is not in the registry. None means that the default registry is used, i.e. that dicts, tuples and lists are considered containers. "extended" means that in addition numpy arrays and params DataFrames are considered containers. Passing a dictionary where the keys are types and the values are dicts with the entries "flatten", "unflatten" and "names" allows to completely override the default registries. :type registry: dict or None :returns: A pair where the first element is a list of leaf values and the second element is a treedef representing the structure of the flattened tree. .. py:function:: tree_just_flatten(tree, is_leaf=None, registry=None) Flatten a pytree without creating a treedef. :param tree: a pytree to flatten. :param is_leaf: An optionally specified function that will be called at each flattening step. It should return a boolean, which indicates whether the flattening should traverse the current object, or if it should be stopped immediately, with the whole subtree being treated as a leaf. :type is_leaf: callable or None :param registry: A pytree container registry that determines which types are considered container objects that should be flattened. ``is_leaf`` can override this in the sense that types that are in the registry are still considered a leaf but it cannot declare something a container that is not in the registry. None means that the default registry is used, i.e. that dicts, tuples and lists are considered containers. "extended" means that in addition numpy arrays and params DataFrames are considered containers. Passing a dictionary where the keys are types and the values are dicts with the entries "flatten", "unflatten" and "names" allows to completely override the default registries. :type registry: dict or None :returns: A pair where the first element is a list of leaf values and the second element is a treedef representing the structure of the flattened tree. .. py:function:: tree_just_yield(tree, is_leaf=None, registry=None) Yield leafs from a pytree without creating a treedef. :param tree: a pytree. :param is_leaf: An optionally specified function that will be called at each yield step. It should return a boolean, which indicates whether the generator should traverse the current object, or if it should be stopped immediately, with the whole subtree being treated as a leaf. :type is_leaf: callable or None :param registry: A pytree container registry that determines which types are considered container objects that should be yielded. ``is_leaf`` can override this in the sense that types that are in the registry are still considered a leaf but it cannot declare something a container that is not in the registry. None means that the default registry is used, i.e. that dicts, tuples and lists are considered containers. "extended" means that in addition numpy arrays and params DataFrames are considered containers. Passing a dictionary where the keys are types and the values are dicts with the entries "flatten", "unflatten" and "names" allows to completely override the default registries. :type registry: dict or None :returns: A generator of leaf values. .. py:function:: tree_map(func, tree, is_leaf=None, registry=None) Apply func to all leaves in tree. :param func: Function applied to each leaf in the tree. :type func: callable :param tree: A pytree. :param is_leaf: An optionally specified function that will be called at each flattening step. It should return a boolean, which indicates whether the flattening should traverse the current object, or if it should be stopped immediately, with the whole subtree being treated as a leaf. :type is_leaf: callable or None :param registry: A pytree container registry that determines which types are considered container objects that should be flattened. `is_leaf` can override this in the sense that types that are in the registry are still considered a leaf but it cannot declare something a container that is not in the registry. None means that the default registry is used, i.e. that dicts, tuples and lists are considered containers. "extended" means that in addition numpy arrays and params DataFrames are considered containers. Passing a dictionary where the keys are types and the values are dicts with the entries "flatten", "unflatten" and "names" allows to completely override the default registries. :type registry: dict or None :returns: modified copy of tree. .. py:function:: tree_multimap(func, *trees, is_leaf=None, registry=None) Apply func to leaves of multiple pytrees. :param func: Function applied to each leaf corresponding leaves of multiple py trees. :type func: callable :param trees: An arbitrary number of pytrees. All trees need to have the same structure. :param is_leaf: An optionally specified function that will be called at each flattening step. It should return a boolean, which indicates whether the flattening should traverse the current object, or if it should be stopped immediately, with the whole subtree being treated as a leaf. :type is_leaf: callable or None :param registry: A pytree container registry that determines which types are considered container objects that should be flattened. `is_leaf` can override this in the sense that types that are in the registry are still considered a leaf but it cannot declare something a container that is not in the registry. None means that the default registry is used, i.e. that dicts, tuples and lists are considered containers. "extended" means that in addition numpy arrays and params DataFrames are considered containers. Passing a dictionary where the keys are types and the values are dicts with the entries "flatten", "unflatten" and "names" allows to completely override the default registries. :type registry: dict or None :returns: tree with the same structure as the elements in trees. .. py:function:: tree_unflatten(treedef, leaves, is_leaf=None, registry=None) Reconstruct a pytree from the treedef and a list of leaves. The inverse of :func:`tree_flatten`. :param treedef: the treedef to with information needed for reconstruction. :param leaves: the list of leaves to use for reconstruction. The list must match the leaves of the treedef. :type leaves: list :param is_leaf: An optionally specified function that will be called at each flattening step. It should return a boolean, which indicates whether the flattening should traverse the current object, or if it should be stopped immediately, with the whole subtree being treated as a leaf. :type is_leaf: callable or None :param registry: A pytree container registry that determines which types are considered container objects that should be flattened. `is_leaf` can override this in the sense that types that are in the registry are still considered a leaf but it cannot declare something a container that is not in the registry. None means that the default registry is used, i.e. that dicts, tuples and lists are considered containers. "extended" means that in addition numpy arrays and params DataFrames are considered containers. Passing a dictionary where the keys are types and the values are dicts with the entries "flatten", "unflatten" and "names" allows to completely override the default registries. :type registry: dict or None :returns: The reconstructed pytree, containing the ``leaves`` placed in the structure described by ``treedef``. .. py:function:: tree_update(tree, other, is_leaf=None, registry=None) Update leaves in a pytree with leaves from another pytree. The second pytree must be compatible with the first one but can be smaller. For example, lists can be shorter, dictionaries can contain subsets of entries, etc. :param tree: A pytree. :param other: Another pytree. :param is_leaf: An optionally specified function that will be called at each flattening step. It should return a boolean, which indicates whether the flattening should traverse the current object, or if it should be stopped immediately, with the whole subtree being treated as a leaf. :type is_leaf: callable or None :param registry: A pytree container registry that determines which types are considered container objects that should be flattened. `is_leaf` can override this in the sense that types that are in the registry are still considered a leaf but it cannot declare something a container that is not in the registry. None means that the default registry is used, i.e. that dicts, tuples and lists are considered containers. "extended" means that in addition numpy arrays and params DataFrames are considered containers. Passing a dictionary where the keys are types and the values are dicts with the entries "flatten", "unflatten" and "names" allows to completely override the default registries. :type registry: dict or None :returns: Updated pytree. .. py:function:: tree_yield(tree, is_leaf=None, registry=None) Yield leafs from a pytree and create the tree definition. :param tree: a pytree. :param is_leaf: An optionally specified function that will be called at each yield step. It should return a boolean, which indicates whether the generator should traverse the current object, or if it should be stopped immediately, with the whole subtree being treated as a leaf. :type is_leaf: callable or None :param registry: A pytree container registry that determines which types are considered container objects that should be yielded. ``is_leaf`` can override this in the sense that types that are in the registry are still considered a leaf but it cannot declare something a container that is not in the registry. None means that the default registry is used, i.e. that dicts, tuples and lists are considered containers. "extended" means that in addition numpy arrays and params DataFrames are considered containers. Passing a dictionary where the keys are types and the values are dicts with the entries "flatten", "unflatten" and "names" allows to completely override the default registries. :type registry: dict or None :returns: A pair where the first element is a generator of leaf values and the second element is a treedef representing the structure of the flattened tree.