flax
replace lazy_init implementation to avoid jax partial_eval
#5551
Merged

Loading