Add jax.tree_util.register_simple

[jakevdp]

This PR proposes a new API, jax.tree_util.register_simple that will simplify pytree registration and promote best practices in the most common of cases. Example usage with dataclass, although this will work on any class that can be flattened and unflattened via simple attribute access:

import jax
import jax.numpy as jnp
from dataclasses import dataclass

@jax.tree_util.register_simple(
    dynamic_attributes=['x', 'y'],
    static_attributes=['val'])
@dataclass
class MyContainer:
    x: jax.Array
    y: jax.Array
    val: str

m = MyContainer(jnp.zeros(4), jnp.arange(4), val='name')

leaves, tree = jax.tree.flatten(m)
m2 = jax.tree.unflatten(tree, leaves)

print(m)
# MyContainer(x=Array([0., 0., 0., 0.], dtype=float32), y=Array([0, 1, 2, 3], dtype=int32), val='name')
print(m2)
# MyContainer(x=Array([0., 0., 0., 0.], dtype=float32), y=Array([0, 1, 2, 3], dtype=int32), val='name')

One benefit of this is that it's somewhat self-documenting: we mention in the documentation that children should contain dynamic data, and aux_data should contain static data, but in practice users often miss this detail (e.g. we've had several questions in recent weeks about problems arising after including arrays in aux_data). Calling these static_attributes and dynamic_attributes should hopefully point people in the right direction.

[jakevdp]

This is probably too duplicative of jax.tree_util.register_dataclass. Closing.