google / google/flax

Linen: Consider raising an error when reading variables from submodule before initialization?

Open
#513 10 comments 0 reactions 0 assignees View on GitHub
Priority: P2 - no schedule Status: pull requests welcome
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

Within a module that uses shape-inference (as most of the built-in Linen modules do), this code is fine:

```
class MyModule(nn.Module):
def __call__(self, x):
conv = nn.Conv(features=3)
y = conv(x)
params = conv.params()
```

But if you instead do:
```
class MyModule(nn.Module):
def __call__(self, x):
conv = nn.Conv(features=3)
params = conv.params()
y = conv(x)
```

Then I believe you get an empty `params` dict (as the parameters of `conv` are only initialized once the input shape is known)

Seems like users may be surprised about this, so instead we could just raise an error if the variables are empty when reading them clarifying what is happening, to guide users to the right direction.

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.