flax

239d4e6b - Reverting unsafe_pytree #4030

Reverting unsafe_pytree #4030 Exploring better alternatives that will not clash with the assumption that graph nodes are not pytrees present in many functions. PiperOrigin-RevId: 651419435