jax

cfbdb65a - add register_pytree_node_class, fixes #2396 (#2400)

add register_pytree_node_class, fixes #2396 (#2400) Co-authored-by: Stephan Hoyer <shoyer@google.com> Co-authored-by: Stephan Hoyer <shoyer@google.com>