jax.tree_util.tree_map#

jax.tree_util.tree_map(f, tree, *rest, is_leaf=None)[源代码]#

是 jax.tree.map() 的别名。

参数:

  • f (可调用对象[..., 任意类型])

  • tree (Any)

  • rest (Any)

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

返回类型:

任意类型