jax.tree_util.all_leaves#

jax.tree_util.all_leaves(iterable, is_leaf=None)[源代码]#

测试给定的可迭代对象中的所有元素是否都是叶子节点。

此函数在高级用例中很有用，例如，如果一个库允许在叶子节点的扁平化可迭代对象上进行任意映射操作，它可能想检查结果是否仍然是叶子节点的扁平化可迭代对象。

参数:

  • iterable (Iterable[Any]) – 叶子节点的迭代器。

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

返回:

一个布尔值，指示输入中的所有元素是否都是叶子节点。

返回类型:

bool

示例

>>> import jax
>>> tree = {"a": [1, 2, 3]}
>>> assert all_leaves(jax.tree_util.tree_leaves(tree))
>>> assert not all_leaves([tree])