jax.tree_util.tree_structure#

jax.tree_util.tree_structure(tree, is_leaf=None)[source]#

Gets the treedef for a pytree.

Parameters:

tree (Any) – the pytree for which to get the leaves
is_leaf (None | Callable[[Any], bool]) – 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.

Returns:

a PyTreeDef representing the structure of the tree.

Return type:

pytreedef

Example

>>> import jax
>>> jax.tree.structure([1, (2, 3), [4, 5]])
PyTreeDef([*, (*, *), [*, *]])

jax.tree_util.tree_structure