Compatibility between submodule declaration and @nn.compact

[rolandgvc]

Wouldn't it be handy to be able to use smaller blocks defined with @nn.compact as submodules of a Model? Something like:

class ConvBlock(nn.Module):
  out_channels: int = 32
  kernel_size: Sequence[int] = (3,3)
  strides: Optional[Sequence[int]] = None
  padding: Union[str, Sequence[Tuple[int, int]]] = 'VALID'
  train: bool = False
  dtype: Any = jnp.float32

  @nn.compact
  def __call__(self, x):
    x = nn.Conv(self.out_channels, kernel_size=self.kernel_size, strides=self.strides,
                padding=self.padding, use_bias=False, name='conv', dtype=self.dtype)(x)
    x = nn.BatchNorm(use_running_average=not self.train, name='bn', dtype=self.dtype)(x)
    return nn.relu(x)


class Model(nn.Module):
  train: bool = False
  dtype: Any = jnp.float32

  def setup(self):
    self.conv_block = functools.partial(ConvBlock, train=self.train, dtype=self.dtype)

  def __call__(self, x):
    x = self.conv_block(32, kernel_size=(3, 3), strides=(2, 2))(x)
    return x

Is there a technical reason why this is not possible?

[avital]

Hi @rolandgvc -- to clarify, are you saying that the code you wrote above doesn't work? (I guess our __setattr__ system for detecting submodule assignment wouldn't understand that functools.partial(ConvBlock, ...) is a module.)

What actually currently happens when you run this code?

[rolandgvc]

ValueError: Submodules must be defined in setup() or in a method wrapped in @compact
You can check it out in this colab: https://colab.research.google.com/drive/1ydIe8pwFbkEXouW0onIdxesj2uvgwD8r?usp=sharing

[rolandgvc]

Yes, apparently this happens only when I use partial. Now I see why. Thanks!

[avital]

Can you please file a bug with a repro?

[rolandgvc]

A bug about the __setattr__ system not understanding modules with functools.partial?

[avital]

yes

[marcvanzee]

I think the problem here is not related to __setattr__. The problem is that in Model.__call__() you call self.conv_block(...), which constructs the Module ConvBlock (because partial inside Model.setup() does not create it yet). The __post_init__ call in ConvBlock will then check where in the parent (which is Model) the submodule was created, and realize this was outside of setup.

I don't think this is a bug, because you aren't allowed to create submodules in this way. However, I agree this error message is confusing. I am adding this example to the error messages in #1072 .