jax.tree_util.tree_broadcast#

jax.tree_util.tree_broadcast(prefix_tree, full_tree, is_leaf=None)[source]#

Alias of jax.tree.broadcast().

Parameters:

  • prefix_tree (Any)

  • full_tree (Any)

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

Return type:

Any