flax

c48905a2 - [nnx] fix while_loop/fori_loop bug when sharing references