jax.tree_util.tree_map#

jax.tree_util.tree_map(f, tree, *rest, is_leaf=None)[原始碼]#

別名為 jax.tree.map()。

參數:

  • f (Callable[..., Any])

  • tree (Any)

  • rest (Any)

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

回傳型別:

Any