jax.tree.broadcast#

jax.tree.broadcast(prefix_tree, full_tree, is_leaf=None)[source]#

Broadcasts a tree prefix into the full structure of a given tree.

Parameters:

  • prefix_tree (Any) – a pytree that is a tree prefix of full_tree.

  • full_tree (Any) – a pytree with the structure to broadcast the prefix leaves into.

  • is_leaf (Callable[[Any], bool] | None) – an optionally specified function that will be called at each flattening step. It should return a boolean, with true stopping the traversal and the whole subtree being treated as a leaf, and false indicating the flattening should traverse the current object.

Returns:

A pytree matching the structure of full_tree where the leaves of prefix_tree have been broadcasted into the leaves of each corresponding subtree.

Return type:

Any

Examples

>>> import jax
>>> prefix = (1, 2, 3)
>>> full = (0, {'a': 0, 'b': 0}, (0, 0))
>>> jax.tree.broadcast(prefix, full)
(1, {'a': 2, 'b': 2}, (3, 3))

See also