flax

[nnx] fix while_loop/fori_loop bug when sharing references

#4379

Merged