jax.tree_util.tree_broadcast

jax.tree_util.tree_broadcast(prefix_tree, full_tree, is_leaf=None)

Alias of jax.tree.broadcast().

Parameters:

• prefix_tree (Any)

• full_tree (Any)

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

Return type:

Any