jax.tree.all

jax.tree.all(tree, *, is_leaf=None)

对树的所有叶子调用 all()。

参数:

• tree (Any) – 要评估的 pytree

• is_leaf (Callable[[Any], bool] | None) – 一个可选指定的函数，将在每次扁平化步骤中被调用。它应返回一个布尔值，指示是否应遍历当前对象，或者是否应立即停止，将整个子树视为叶子。

返回:

布尔值 真 或 假

返回类型:

result

示例

>>> import jax
>>> jax.tree.all([True, {'a': True, 'b': (True, True)}])
True
>>> jax.tree.all([False, (True, False)])
False

参见

• jax.tree.reduce()

• jax.tree.leaves()