Handle pure bodies in nnx.fori_loop

[samanklesaria]

Fixes #5113

Currently, nnx.fori_loop only works with loop bodies that perform some kind of mutation. This is because the the call to _add_fake_index_mapping(pure_init_val) in line 1625 of flax/nnx/transforms/iteration.py adds an outer_index attribute that won't be added in the corresponding pure body. To fix this, this PR checks what the output of the loop function looks like first. If it has an outer index, we add a fake index mapping. Otherwise, we leave it be.

[samanklesaria]

It's unfortunate, but we can't have this error here anymore. The first time we call the body function, we know that the body function's output won't match. We manually fix the structure in nnx.fori_loop, using the output as a guide. In practice, I don't think this matters. The error you get without this handler is still clear- it just refers to mismatched structures in jax.lax.scan rather than nnx.fori_loop.

[vfdev-5]

What would happen if the function changes the graphdef by adding a new attribute like when calling nnx.Module.sow ?
A similar issue happend with nnx.scan on modified graphdef.

[samanklesaria]

What would happen if the function changes the graphdef by adding a new attribute like when calling nnx.Module.sow ? A similar issue happend with nnx.scan on modified graphdef.

If the function changes the graphdef, you'll get an error about the structures not matching. We also get an error even if the changes are only temporary (using nnx.sow with nnx.pop afterwards), which matches what you were talking about with nnx.scan:

class MLP(nnx.Module):
  def __call__(self):
    self.sow(nnx.Intermediate, "blah", jnp.ones(2))

def sow_fn(i, model):
  model()
  nnx.pop(model, nnx.Intermediate)
  return model

def test_sow():
  model = MLP()
  nnx.fori_loop(0, 2, sow_fn, model)

In this case, "blah" appears as a static attribute in the pytree result of sow_fn even though nnx.pop should have removed it, and we get another mismatched structure error. But this is a bug in nnx.pop rather than nnx.fori_loop. If we run

model = MLP()
sow_fn(2, model)
print(nnx.split(model))

we get a NodeDef with attributes=[('_pytree__nodes', Static(value={'_pytree__state': True, '_pytree__nodes': False, 'blah': True})). As this bug has nothing to do with nnx.fori_loop, I propose that we address it in a separate issue and PR. The separate issue is #5146 which is fixed by #5147