How to assign methods as Flax Module attributes?

[marcvanzee]

Original question from @jacobaustin123:

I've been running into issues with Flax when assigning methods as dataclass attributes. Here's a minimal example:

class UNet(nn.Module):
  nonlinearity = nn.swish

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

This raises an error silu() takes 1 positional argument but 2 were given. From some debugging, my guess is that dataclasses.dataclass and flax.struct.dataclass make Callables assigned as attributes into bound methods, so when they're called, self is passed as the first argument? This makes common idioms this impossible, because self is passed to self.nonlinearity. Is there a workaround or different idiom that can be used in Flax? And even, maybe it's worth discussing if this behavior is desirable (or maybe a warning should be added)?

[marcvanzee]

Answer by @hexahedria:

Dataclass actually reads the type annotation, so the code below works fine (even with Any), but fails without any type annotation.

class Foo(nn.Module):
  go : Callable = somefunc

So moral of the story, put type annotations for the things you want to be dataclass attributes

[marcvanzee]

Some more details by @levskaya:

dataclasses are a class transform that depend on the __annotations__ defined on class attributes, if you don't have a type annotation, it doesn't show up in __annotations__ and isn't processed by dataclass.

methods in python are actually a "non-data descriptor", descriptors are a core mechanism that extends attribute lookup on object (usually things just live in obj.dict, but descriptors are an extension of that, see: https://docs.python.org/3/howto/descriptor.html#functions-and-methods)