jax.tree.all

目录

jax.tree.all#

jax.tree.all(tree, *, is_leaf=None)[source]#

在树的叶子节点上调用 all()。

参数:
  • tree (Any) – 要评估的 pytree

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

返回值:

布尔值 True 或 False

返回类型:

result

示例

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