jax.tree_util.tree_broadcast#

jax.tree_util.tree_broadcast(prefix_tree, full_tree, is_leaf=None)[源代码]#

jax.tree.broadcast() 的别名。

参数:
  • prefix_tree (Any)

  • full_tree (Any)

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

返回类型:

任意类型