google / google/flax

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

Open
#2,733 2 comments 0 reactions 1 assignee Claimed by @levskaya View on GitHub
Priority: P2 - no schedule
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.