jax.tree_util.tree_map_with_path

jax.tree_util.tree_map_with_path(f, tree, *rest, is_leaf=None, is_leaf_takes_path=False)

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

参数:

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

• tree (Any)

• rest (Any)

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

• is_leaf_takes_path (bool)

返回类型:

任意类型