How and when to use nnx.fori_loop and nnx.scan to stack layers?

[lbergmann1]

Hi everyone,

I am currently trying to optimize my JAX/NNX code and I am wondering how to correctly use nnx.scan and nnx.fori_loop to stack layers.

Here is some example code to create a simple MLP:

import jax
import jax.numpy as jnp
from flax import nnx

######################################################################
# first solution: python for-loop (works) or nnx.fori_loop (errors)  #
######################################################################
class MLP(nnx.Module):
    def __init__(
        self, rngs: nnx.Rngs, in_features: int, hidden_features: list, out_features: int, leaky_relu_negative_slope: float = 0.01
    ) -> None:
        self.leaky_relu_negative_slope = leaky_relu_negative_slope
        self.linear_layers = []

        for i, num_hf in enumerate(hidden_features):
            self.linear_layers.append(
                nnx.Linear(
                    in_features=in_features if i == 0 else hidden_features[i - 1], out_features=num_hf, use_bias=True, rngs=rngs
                )
            )
        self.linear_layers.append(nnx.Linear(in_features=hidden_features[-1], out_features=out_features, use_bias=True, rngs=rngs))

    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
        """Forward pass using a python for-loop."""
        for layer in self.linear_layers[:-1]:
            x = nnx.leaky_relu(layer(x), negative_slope=self.leaky_relu_negative_slope)
        return self.linear_layers[-1](x)
    
    def forward_fori_loop(self, x: jnp.ndarray) -> jnp.ndarray:      
        """Forward pass using nnx.fori_loop."""  
        def forward(i, state):
            layers, x, neg_slope = state
            x = layers[i](x) # <- layers[i] does not work, because i is a tracer 
            x = nnx.leaky_relu(x, negative_slope=neg_slope) 
            new_state = (layers, x, neg_slope)
            return new_state
        
        self.linear_layers, x, _ = nnx.fori_loop(0, len(self.linear_layers), forward, (self.linear_layers, x, self.leaky_relu_negative_slope))

        return self.linear_layers[-1](x)

######################################################################################################################  
# second solution: nnx.scan (most of the code is taken from this issue: https://github.com/jax-ml/jax/issues/27228)  #
######################################################################################################################  
class Model(nnx.Module):
  def __init__(self, layer_keys):
    @nnx.vmap(in_axes=0, out_axes=0)
    def create_layer(key):
        layer_rngs = nnx.Rngs(key)
        return nnx.Linear(32, 32, rngs=layer_rngs)
    
    self.model = create_layer(layer_keys)
    print(f"type(self.model): {type(self.model)}")
    
  def __call__(self, x):
    @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=nnx.Carry)
    def forward(x: jax.Array, model: nnx.Module):
      return model(x)

    return forward(x, self.model)

    
if __name__ == "__main__":
    use_first_solution = True
    
    if use_first_solution:
        # first solution
        rngs = nnx.Rngs(default=0)
                
        mlp = MLP(rngs=rngs, in_features=128, hidden_features=[265,128,64], out_features=10)
        
        x = jnp.ones((42, 128)) # test input
        y_for_loop = mlp(x)
        y_fori_loop = mlp.forward_fori_loop(x) # error: jax.errors.TracerIntegerConversionError: The __index__() method was called on traced array with shape int32[]
        
        assert (y_for_loop == y_fori_loop).all()
    else:
        # second solution
        key = jax.random.PRNGKey(0)
        layer_keys = jax.random.split(key, 2)   # I would expect to get 2 layers, but self.model is a nnx.Linear layer

        model = Model(layer_keys)
        x = jnp.ones((42,32))
        y = model(x)

Here are my questions:

1. First, a very general question: when is it beneficial to use nnx.fori_loop or nnx.scan, especially depending on the number of iterations? If I understand the JAX documentation correctly, these transforms should reduce compilation times for JIT-compiled functions compared to a python for-loop, especially for loops with many iterations. But are there other advantages (performance, memory usage, ...)?
2. The solution using the nnx.fori_loop crashes, because the index variable is a tracer. The NNX documentation says that nnx.Modules can be composed using lists. Is it not intended to use nnx.fori_loop to iterate over nnx.Modules, or am I missing something?
3. The solution using the nnx.scan works, but I would expect to get two linear layers (see self.model). However, self.model is just a single linear layer. So, I assume that the code is not correct? If so, what would be the correct way to use nnx.scan in this example?

Many thanks in advance ☺️

[cgarciae]

Hi @lbergmann1,

The main use case for scan in JAX is to reduce compilation times, I don't think there is a benefit for runtime. However, you have to use something like vmap during initialization to add an additional dimension to the Arrays so you could then scan over it. Check this example: https://flax.readthedocs.io/en/latest/nnx_basics.html#scan-over-layers