Submodule Sequential must be defined in setup() or in a method wrapped in @compact

[SauravMaheshkar]

I'm working on a Flax implementation for gMLP (gated MLP) from Pay Attention to MLPs. My work so far is in SauravMaheshkar/gMLP and is based on lucidrains' implementations.

I've made a simple test.py to check instantiation using the .init() function. My test script is as follows :

from gmlp import gMLP
import jax
from jax import random

x = jax.random.randint(key=random.PRNGKey(0), minval=0, maxval=20000, shape=(1, 1000))

init_rngs = {"params": random.PRNGKey(0), "layers": random.PRNGKey(1)}

gMLP(num_tokens = 20000,
    dim = 512,
    depth = 4,
    seq_len = 1024).init(init_rngs,x)

And I've been getting this error message

Error Message

---------------------------------------------------------------------------
AssignSubModuleError                      Traceback (most recent call last)
<ipython-input-2-f9407b1608f5> in <module>()
     10     dim = 512,
     11     depth = 4,
---> 12     seq_len = 1024).init(init_rngs,x)

10 frames
/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in init(self, rngs, method, mutable, *args, **kwargs)
    998     _, v_out = self.init_with_output(
    999         rngs, *args,
-> 1000         method=method, mutable=mutable, **kwargs)
   1001     return v_out
   1002 

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in init_with_output(self, rngs, method, mutable, *args, **kwargs)
    967       rngs = {'params': rngs}
    968     return self.apply(
--> 969         {}, *args, rngs=rngs, method=method, mutable=mutable, **kwargs)
    970 
    971   def init(self,

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in apply(self, variables, rngs, method, mutable, capture_intermediates, *args, **kwargs)
    937         method, self,
    938         mutable=mutable, capture_intermediates=capture_intermediates
--> 939     )(variables, *args, **kwargs, rngs=rngs)
    940 
    941   def init_with_output(self,

/usr/local/lib/python3.7/dist-packages/flax/core/scope.py in wrapper(variables, rngs, *args, **kwargs)
    685               **kwargs) -> Union[Any, Tuple[Any, VariableDict]]:
    686     with bind(variables, rngs=rngs, mutable=mutable).temporary() as root:
--> 687       y = fn(root, *args, **kwargs)
    688     if mutable is not False:
    689       return y, root.mutable_variables()

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in scope_fn(scope, *args, **kwargs)
   1176     _context.capture_stack.append(capture_intermediates)
   1177     try:
-> 1178       return fn(module.clone(parent=scope), *args, **kwargs)
   1179     finally:
   1180       _context.capture_stack.pop()

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
    264       self._state.in_setup = True
    265     else:
--> 266       self._try_setup()
    267 
    268     if is_compact_method:

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in _try_setup(self, shallow)
    677             self._register_submodules(field.name, getattr(self, field.name))
    678         if not shallow:
--> 679           self.setup()
    680       finally:
    681         self._state.in_setup = False

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
    273     _context.module_stack.append(self)
    274     try:
--> 275       y = fun(self, *args, **kwargs)
    276       if _context.capture_stack:
    277         filter_fn = _context.capture_stack[-1]

/usr/local/lib/python3.7/dist-packages/gmlp/model.py in setup(self)
     71         self.to_logits = [
     72             Sequential(nn.LayerNorm(), nn.Dense(features=self.num_tokens))
---> 73             if self.num_tokens is not None
     74             else Identity()
     75         ]

/usr/local/lib/python3.7/dist-packages/gmlp/utils.py in __init__(self, layers, parent, name)

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in __post_init__(self)
    580         return
    581       if not self.parent._initialization_allowed:
--> 582         raise errors.AssignSubModuleError(self.__class__.__name__)
    583       # Autonaming of submodules.
    584       if self.name is None:  # pytype: disable=attribute-error

AssignSubModuleError: Submodule Sequential must be defined in `setup()` or in a method wrapped in `@compact` (https://flax.readthedocs.io/en/latest/flax.errors.html#flax.errors.AssignSubModuleError)

The problem lies in the Sequential Module defined as follows:

class Sequential(nn.Module):

    layers: Sequence[nn.Module]

    @nn.compact
    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

Even though the module is defined with @nn.compact, I'm getting the flax.errors.AssignSubModuleError exception.

Any help would be much appreciated 😊.

[asvskartheek]

self.layers is a list of nn.Module objects in your code. Can you make sure that the order is not in one of their definitions i.e, Are they also well-defined with nn.compat decorator?

[SauravMaheshkar]

I tried decorating __call__ with @nn.compact but that doesn't solve the problem. Also, I'm not sure about what you mean by "make sure that the order is not in one of their definitions". I've even tried changing

layers: Sequence[nn.Module]

to 

layers: Sequence[Type(Module)]

but that doesn't solve the problem either.

[SauravMaheshkar]

Tunnel Vision got to me 😅. I wasn't passing a list of modules. I was doing something like Sequential(Dense(), LayerNorm()) instead of Sequential([Dense(), LayerNorm()]). Closing the discussion.

[zaccharieramzi]

I had exactly the same problem: maybe a better error would help though?

[cgarciae]

I am very curious why this original code is not producing a TypeError or similar on construction. Will take a look.

[cgarciae]

I've created #3129 to give an error message if layers is not a Sequence, however, there is a deeper issue with how our dataclasses are defined.