flax

9383573c - [nnx] add pytree support