Wrap __post_init__ to check or force that Module.__post_init__ is called

[levskaya]

We added a runtime check in #2535 to emit a "IncorrectPostInitOverrideError" when users forget to call super().__post_init__() from a custom __post_init__ method. That implementation only runs the check at the top-level though, so it's not going to catch internal Modules that override __post_init__.

e.g.

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

class Foo(nn.Module):
  def __post_init__(self):
    pass
  @nn.compact
  def __call__(self, x):
    return nn.Dense(12)(x)

class Bar(nn.Module):
  @nn.compact
  def __call__(self, x):
    return Foo()(x)

b = Bar()
b.init(jax.random.PRNGKey(0), x=jnp.ones((10, 3)))

What I think we should do instead is to just wrap __post_init__ on subclass instantiation to check (and possible even force) that Module.__post_init__ is called.

For forcing the only complex case is when users use inheritance and call their own parent-class __post_init__ in some random order inside their child __post_init__ (I've seen this.), but I think there's a way for a wrapper to detect this and to only run in the outermost child-defined __post_init__.

[cgarciae]

I think there's a way for a wrapper to detect this and to only run in the outermost child-defined post_init.

Maybe a context manager that contains that stores a references to the object calling __post_init__ so we don't call it again for the same object in case there are multiple wrapped methods?

[chiamp]

For forcing the only complex case is when users use inheritance and call their own parent-class post_init in some random order inside their child post_init (I've seen this.), but I think there's a way for a wrapper to detect this and to only run in the outermost child-defined post_init.

Is there a short code example of this @levskaya?