jax.tree_util.tree_flatten# jax.tree_util.tree_flatten(tree, is_leaf=None)[源代码]# 别名 jax.tree.flatten()。 参数: tree (Any) is_leaf (Callable[[Any], bool] | None) 返回类型: tuple[list[Leaf], PyTreeDef]
