flax

1f27faa8 - [nnx] improve nnx basics