jax.tree_util.treedef_tuple

jax.tree_util.treedef_tuple(treedefs)

Makes a tuple treedef from an iterable of child treedefs.

Parameters:

treedefs (Iterable[PyTreeDef]) – iterable of PyTree structures

Returns:

a single treedef representing a tuple of the structures

Return type:

PyTreeDef

Examples

>>> import jax
>>> x = [1, 2, 3]
>>> y = {'a': 4, 'b': 5}
>>> x_tree = jax.tree.structure(x)
>>> y_tree = jax.tree.structure(y)
>>> xy_tree = jax.tree_util.treedef_tuple([x_tree, y_tree])
>>> xy_tree == jax.tree.structure((x, y))
True

See also

• jax.tree_util.treedef_children()