flax

bf36dec6 - [nnx] improving performance of `_graph_flatten` and `_graph_unflatten` functions