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 (Optional[Callable[[Any], bool]]) –

Return type

PyTreeDef