flax

1fbe17d8 - Copybara import of the project:

Copybara import of the project: -- 1f27faa8d2e02ae993e2af8f501d29bcee203c83 by Cristian Garcia <cgarcia.e88@gmail.com>: [nnx] improve nnx basics PiperOrigin-RevId: 625789767