jax

b1604459 - Clarify the intended purpose of tree_util.

Clarify the intended purpose of tree_util. Most importantly, this removes the initial paragraph which was easy to misinterpret to imply that this module was not JAX-specific.