flax

[bridge] Set _initializing correctly and avoid return RNG states

#4569

Merged