flax

e3789de2 - Copybara import of the project:

Copybara import of the project: -- 4d8c9bbeff254016a81b5a286eb97dd02e9b7a0e by Cristian Garcia <cgarcia.e88@gmail.com>: [nnx] pytrees are graph nodes PiperOrigin-RevId: 734264536