flax

f8164dd4 - [nnx] test fori_loop output