flax

df2ed160 - Copybara import of the project:

Copybara import of the project: -- af408cf5f361b221a748e6dafb22b441675fda77 by Cristian Garcia <cgarcia.e88@gmail.com>: fix List and improve Pytree PiperOrigin-RevId: 862388098