flax

63de4c42 - [nnx] simplify flatten/unflatten