sowing across nnx.while_loop

[jacknewsom]

Hello,

Is it possible to sow values inside an nnx.Module called within nnx.while_loop? The documentation says

Caution: for the NNX internal reference tracing mechanism to work, you cannot change the variable reference structure of init_val inside body_fun.

For example, the snippet below fails:

import flax.nnx as nnx

import jax
import flax.linen as nn
import flax.nnx as nnx
import flax.nnx.bridge as bridge
import jax.numpy as jnp


class Block(nn.Module):
    @nn.compact
    def __call__(self, carry, unused_inputs):
        carry = nn.Dense(100)(carry)
        carry = nn.relu(carry)
        self.sow("intermediates", "carry", carry)
        return carry, None


class MLP(nn.Module):
    depth: int

    @nn.compact
    def __call__(self, x):
        self.sow("intermediates", "x", x)
        scan_block = nn.scan(
            Block,
            variable_axes={"params": 0, "intermediates": 0},
            split_rngs={"params": True},
            length=self.depth,
        )()
        carry, unused_outputs = scan_block(x, None)

        return carry


key = jax.random.key(42)
key, nnx_key = jax.random.split(key, 2)

x = jax.random.normal(key, (1, 100))
model_nnx = bridge.ToNNX(MLP(depth=3), rngs=nnx.Rngs(nnx_key))
bridge.lazy_init(model_nnx, x)

def step(carry):
    model, inputs, step = carry
    outputs = model(inputs, mutable=True)
    return model, outputs, step + 1

def cond(carry):
    _, _, step = carry
    return step < 5

model, _, step = nnx.while_loop(cond, step, (model_nnx, x, 0))

with this error:

ubuntu@129-146-27-29:~/us-west-2-fs1/steering$ python scripts/looptest.py 
Traceback (most recent call last):
  File "/lambda/nfs/us-west-2-fs1/steering/openpi/.venv/lib/python3.11/site-packages/flax/nnx/transforms/iteration.py", line 1398, in __call__
    jax.tree.map(lambda a, b: None, pure_val, pure_out)
  File "/lambda/nfs/us-west-2-fs1/steering/openpi/.venv/lib/python3.11/site-packages/jax/_src/tree.py", line 155, in map
    return tree_util.tree_map(f, tree, *rest, is_leaf=is_leaf)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lambda/nfs/us-west-2-fs1/steering/openpi/.venv/lib/python3.11/site-packages/jax/_src/tree_util.py", line 357, in tree_map
    all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lambda/nfs/us-west-2-fs1/steering/openpi/.venv/lib/python3.11/site-packages/jax/_src/tree_util.py", line 357, in <listcomp>
    all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
                             ^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: Mismatch custom node data: ('ScanBlock_0', 'rngs') != ('ScanBlock_0', 'rngs', 'x'); value: State({
  'ScanBlock_0': {
    'Dense_0': {
      'bias': VariableState(
        type=Param,
        value=Traced<ShapedArray(float32[3,100])>with<DynamicJaxprTrace>
      ),
      'kernel': VariableState(
        type=Param,
        value=Traced<ShapedArray(float32[3,100,100])>with<DynamicJaxprTrace>
      )
    },
    'carry': {
      0: VariableState(
        type=Intermediate,
        value=Traced<ShapedArray(float32[3,1,100])>with<DynamicJaxprTrace>
      )
    }
  },
  'rngs': {
    'default': {
      'count': VariableState(
        type=RngCount,
        value=Traced<ShapedArray(uint32[])>with<DynamicJaxprTrace>,
        tag='default'
      ),
      'key': VariableState(
        type=RngKey,
        value=Traced<ShapedArray(key<fry>[])>with<DynamicJaxprTrace>,
        tag='default'
      )
    }
  },
  'x': {
    0: VariableState(
      type=Intermediate,
      value=Traced<ShapedArray(float32[1,100])>with<DynamicJaxprTrace>
    )
  }
}).

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/lambda/nfs/us-west-2-fs1/steering/scripts/looptest.py", line 52, in <module>
    model, _, step = nnx.while_loop(cond, step, (model_nnx, x, 0))
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lambda/nfs/us-west-2-fs1/steering/openpi/.venv/lib/python3.11/site-packages/flax/nnx/graph.py", line 1081, in update_context_manager_wrapper
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/lambda/nfs/us-west-2-fs1/steering/openpi/.venv/lib/python3.11/site-packages/flax/nnx/transforms/iteration.py", line 1451, in while_loop
    pure_out = jax.lax.while_loop(
               ^^^^^^^^^^^^^^^^^^^
  File "/lambda/nfs/us-west-2-fs1/steering/openpi/.venv/lib/python3.11/site-packages/flax/nnx/graph.py", line 1081, in update_context_manager_wrapper
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/lambda/nfs/us-west-2-fs1/steering/openpi/.venv/lib/python3.11/site-packages/flax/nnx/transforms/iteration.py", line 1406, in __call__
    raise ValueError(msg)
ValueError: nnx.while_loop requires body function's input and output to have the same reference and pytree structure, but they differ. If the mismatch comes from `index_mapping` field, you might have modified reference structure within the body function, which is not allowed.Detail of the mismatch: 
 Mismatch custom node data: ('ScanBlock_0', 'rngs') != ('ScanBlock_0', 'rngs', 'x'); value: State({
  'ScanBlock_0': {
    'Dense_0': {
      'bias': VariableState(
        type=Param,
        value=Traced<ShapedArray(float32[3,100])>with<DynamicJaxprTrace>
      ),
      'kernel': VariableState(
        type=Param,
        value=Traced<ShapedArray(float32[3,100,100])>with<DynamicJaxprTrace>
      )
    },
    'carry': {
      0: VariableState(
        type=Intermediate,
        value=Traced<ShapedArray(float32[3,1,100])>with<DynamicJaxprTrace>
      )
    }
  },
  'rngs': {
    'default': {
      'count': VariableState(
        type=RngCount,
        value=Traced<ShapedArray(uint32[])>with<DynamicJaxprTrace>,
        tag='default'
      ),
      'key': VariableState(
        type=RngKey,
        value=Traced<ShapedArray(key<fry>[])>with<DynamicJaxprTrace>,
        tag='default'
      )
    }
  },
  'x': {
    0: VariableState(
      type=Intermediate,
      value=Traced<ShapedArray(float32[1,100])>with<DynamicJaxprTrace>
    )
  }
}).
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

Is there any way around this that lets me use the JIT and sow in a loop?

[cgarciae]

@jacknewsom to effectively sow inside a while loop you need to create an array with a max size and add the value you want to sow at a specific index that should increase every time you sow. To do this you can use sow's init_fn and reduce_fn:

max_size: int = ...

def init_fn(x):
  return dict(
    i=jnp.array(0),
    x=jnp.zeros((max_size, *x.shape))
  )
  
def reduce_fn(acc, x):
  return dict(
    i=acc['i'] + 1,
    x=acc['x].at[i].set(x),
  )

class MLP(nn.Module):
  def __call__(self, x):
    self.sow("intermediates", "x", x, init_fn=init_fn, reduce_fn=reduce_fn)
    ...

The sow inside scan is a trickier, as you need to initialize the state before calling scan, you can do something like this:

class MLP(nn.Module):
  def __call__(self, x):
    self.sow("intermediates", "x", x, init_fn=init_fn, reduce_fn=reduce_fn)
    block = Block()
    if self.is_initializing():
      carry = ... # TODO
      block.sow("intermediates", "carry", carry, init_fn=init_fn, reduce_fn=reduce_fn)
    
    carry, _ = nn.scan(
        Block.__call__,
        variable_axes={"params": 0, "intermediates": 0},
        split_rngs={"params": True},
        length=self.depth,
    )(block, x, None)
    return carry