How to reduce compiling cost of a deep NN?

[YouJiacheng]

It seems that only python for-loop can be used to scan a sequence of flax modules.
The sequence length is much larger than the number of module types, but there are different type of modules.
In other words, most of modules are same type/class.
I don't want to inline all forward passes of them, since same forward passes can be factored out, using a xla custom call.
Does lax.switch help?

[jheek]

You can create a loop over similar layers by using flax.linen.scan(Body, length=num_layers, variable_axes={"params": 0}, split_rngs={"params": True}). Usually Body has a few different types of layers for example Body = MLP() -> Attention() which are then repeated a number of times.

If you would have a non-regular sequence of layers you would indeed need something like a switch. But we don't have a linen version of that yet that can switch between layers

[YouJiacheng]

Thanks!!!
Your example is much simpler than my trial:

class Repeat(nn.Module):
    block: Callable[[], nn.Module]
    n: int
    @nn.compact
    def __call__(self, x):
        class Block(nn.Module):
            @nn.compact
            def __call__(_, x):
                return self.block()(x), None
        return nn.scan(Block, variable_axes={'params': 0}, split_rngs={'params': True}, length=self.n)()(x)[0]

[YouJiacheng]

@jheek Hmmm, I encountered an error with your example.

class Foo(nn.Module):
  num_layers: int = 5

  @nn.compact
  def __call__(self, x):
    def body_fn(mdl, x):
      return nn.Dense(features=x.shape[-1])(x), ()
    x, _ = nn.scan(body_fn, length=self.num_layers, variable_axes={"params": 0}, split_rngs={"params": True})(self, x)
    return x

m = Foo()
x = jnp.ones((3,))
v = m.init(jax.random.PRNGKey(0), x)
print(v) # see Dense_1 in dict
print(m.apply(v, x)) # error: Dense_0 kernel not exists

File "/home/jiacheng/conda/envs/x/lib/python3.10/site-packages/flax/linen/linear.py", line 171, in __call__
    kernel = self.param('kernel',
  File "/home/jiacheng/conda/envs/x/lib/python3.10/site-packages/flax/linen/module.py", line 896, in param
    v = self.scope.param(name, init_fn, *init_args)
  File "/home/jiacheng/conda/envs/x/lib/python3.10/site-packages/flax/core/scope.py", line 642, in param
    raise errors.ScopeParamNotFoundError(name, self.path_text)
jax._src.traceback_util.UnfilteredStackTrace: flax.errors.ScopeParamNotFoundError: No parameter named "kernel" exists in "/scan()/Dense_0". (https://flax.readthedocs.io/en/latest/flax.errors.html#flax.errors.ScopeParamNotFoundError)

[YouJiacheng]

@jheek friendly ping :)

[jheek]

I'm pretty sure that's a bug. With a surprisingly short and simple repro though. let me fix it today.

[jheek]

#2025 this should fix the issue