improve tree mismatch error in tree_map

[mattjj]

import jax
import jax.numpy as jnp

jax.tree.map(lambda x: x, {'a': {'b': [1,2,3]}}, {'a': {'b': [1,2,3,4]}})

Before:

ValueError: List arity mismatch: 4 != 3; list: [1, 2, 3, 4].

After:

ValueError: in tree_map, at the key path ['a']['b'] the first tree has a list of length 3
but the second tree has a list of length 4, so the lengths do not match:
[1, 2, 3]
vs
[1, 2, 3, 4]

---

jax.tree.map(lambda x: x, [1,2,3] * 50, [1,2,3,4] * 50)

After:

ValueError: in tree_map, at the tree root the first tree has a list of length 150
but the second tree has a list of length 200, so the lengths do not match.

Intended improvements:

1. show key path
2. show both mismatching subtree values if they're short (rather than just the latter tree's)

TODO (maybe in follow-ups...):

• validate on dicts, custom pytree nodes, etc
• tests

[sshleifer]

Thank you so much way better!

[IvyZX]

This change looks fantastic! Is there anything blocking it to be landed?

[gnecula]

I am also happy to see more thinking here. I have found equality_errors very useful, but hard to use. I recommend searching for usage of tree_util.equality_errors and seeing how we can simplify usage.

[mattjj]

Thanks, @IvyZX ! TBH I don't remember if there were blockers or if I just failed to follow up. Let me see if I can land it now...

@gnecula I agree. This PR is just a small step, and as you say we could really use some deeper thinking about error utilities.

[IvyZX]

Thank you so much!