Wrap `__post_init__` to check or force that `Module.__post_init__` is called
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
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__`.
Contributor guide
Assessment
This issue has not been assessed yet.