jax.tree_util.tree_structure

jax.tree_util.tree_structure(tree, is_leaf=None)

Gets the treedef for a pytree.

Parameters:

• tree (Any) –

• is_leaf (None | Callable[[Any], bool]) –

Return type:

PyTreeDef