Forward Pass: While Loop

[pharringtonp19]

How should I go about writing a while loop in the forward pass of a neural network? In my naive implementation, I run into the error that " Module instance is frozen outside of setup method."

import jax 
import jax.numpy as jnp 
from typing import Sequence
import flax.linen as nn

class MLP(nn.Module):
  tol : float
  max_iter : int
  iterations : int = 0 

  @nn.compact
  def __call__(self, x):
    z = jnp.zeros_like(x)
    while self.iterations < self.max_iter:
      z_next = nn.tanh(nn.Dense(x.shape[0])(x) + x)
      err = jnp.linalg.norm(z-z_next)
      z = z_next 
      self.iterations += 1 
      if err < self.tol:
        break 
    return z 

model = MLP(1e-4, 1000)
x = jax.random.normal(jax.random.PRNGKey(1), shape=(10,1))
variables = model.init(jax.random.PRNGKey(0), x)

SetAttributeFrozenModuleError: Can't set iterations=1 for Module of type MLP: Module instance is frozen outside of setup method. (https://flax.readthedocs.io/en/latest/flax.errors.html#flax.errors.SetAttributeFrozenModuleError)

[jheek]

Are you sure you want iterations to be a mutating field of MLP instead of a local variable in __call__?
Also you are calling nn.Dense inside the while loop which means you create a new Dense (with it's own params) each iteration. Are you sure you don't want param sharing and create the nn.Dense instance before entering the while loop?

[pharringtonp19]

@jheek You are correct on both points. Thanks!

import jax 
import jax.numpy as jnp 
from typing import Sequence
import flax.linen as nn

class MLP(nn.Module):
  features: int
  tol : float
  max_iter : int

  @nn.compact
  def __call__(self, x):

    # Repeated Layer
    affine_map = nn.Dense(self.features)
    layer = lambda z: nn.tanh(affine_map(z) + x)

    # Fwd Pass
    def cond_fun(carry):
      z_prev, z, _ = carry 
      return jnp.linalg.norm(z_prev-z) > self.tol 
    
    def body_fun(carry):
      _, z, iter = carry 
      return z, layer(z), iter + 1
    
    init_carry = (jnp.zeros_like(x), layer(jnp.zeros_like(x)), 0)
    return  jax.lax.while_loop(cond_fun, body_fun, init_carry)[1:]

data_dim = 10 
model = MLP(data_dim, 1e-4, 1000)
x = jax.random.normal(jax.random.PRNGKey(1), shape=(1, data_dim))
variables = model.init(jax.random.PRNGKey(0), x)
output = model.apply(variables, x)
print(output)