What is the purpose of jax.tree_util.tree_all?

[NeilGirdhar]

Why was jax.tree_util.tree_all added? It's defined to be

def tree_all(tree):
  return all(tree_leaves(tree))

Its only use in the JAX library is in jax.test_util where it seems to be superfluous?

def check_eq(xs, ys, err_msg=''):
  assert_close = partial(_assert_numpy_allclose, err_msg=err_msg)
  tree_all(tree_multimap(assert_close, xs, ys))

could be just

def check_eq(xs, ys, err_msg=''):
  assert_close = partial(_assert_numpy_allclose, err_msg=err_msg)
  tree_multimap(assert_close, xs, ys)

Unlike most of the other JAX functions (tree_map, tree_leaves, anything in jax.numpy, etc.), tree_all doesn't work with tracers. I thought I could simplify some code with it only to find that it gave rise to concretization errors. This is very confusing.

Perhaps a better definition would be:

def tree_all(tree: Any) -> bool | BooleanArray:
    return tree_reduce(jnp.logical_and, tree, True)

With this definition, a reasonable definition of tree_allclose would be

def tree_allclose(x: T, y: T) -> bool | BooleanArray:
     return tree_all(tree_multimap(jnp.allclose, x, y))

[jakevdp]

tree_all has been around since the initial commit of JAX on github: a30e858

Folks like @froystig & @mattjj may remember more context on this, but it appears it has always been basically a convenience routine used in testing. You're probably right that it might be cleaner to remove it from the public API at this point, since it's so rarely used and so easily implemented.

[mattjj]

Dougal wrote the pytree code, and a lot of the other code in the first commit :)

(Technically, the two versions of check_eq in the OP aren't exactly the same, in that one makes one top-level assertion while the other makes an assertion at each leaf. They'd raise errors in the same cases but not with the same stack traces.)

Originally tree_util.py wasn't meant to be a public API, but people started using it. I think tree_all was just a utility for tests, where using np rather than jnp made sense.

The alternative tree_all definition could work, but IMO we should probably just delete tree_all! The smaller the API, the better.

[NeilGirdhar]

@mattjj Regarding check_eq, I think you're mistaken? They both do the assertions in _assert_numpy_allclose, which doesn't return anything. tree_all is anding together a bunch of Nones and then throwing the result out, which does nothing.

It would have been different if the code had been something like

assert tree_all(tree_multimap(jnp.allclose, xs, ys))

I assume that's what might have been intended.

I l ike your motivation to keep the API small. JAX is joy to learn because the API is so compact and well thought-out.