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

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.

[jheek]

Still there might be cases where you want to access variables even though they aren’t final. Let’s say I want to access a kernel but I didn’t necessarily define the bias yet. A simple solution would be a helpful reminder to the user. For example variables could be a subclass of FrozenDict that returns a helpful error on KeyError like “Perhaps the variable you are trying to access is not initialized yet…”

[cgarciae]

Based on @jheek’s comment in #638, how about we do the following:

1. Create a simple ErrorHandlers dataclass, e.g:

@dataclass(frozen=True)
class ErrorHandlers:
  missing_key: Optional[Callable[[str], Exception]] = None
  set_item: Optional[Callable[[str, Any], Exception]] = None
  ...

1. Update FrozenDict to have a _error_handlers field and add logic to use the handlers in a couple of places.
2. Add a set_error_handlers method to FrozenDict as proposed by @jheek which takes in the handlers, only modification I would propose is that it returns a new FrozenDict to keep the API immutable.

xs = FrozenDict(variables)
xs = xs.set_error_handlers(missing_key=lambda key: Error(...), set_item=lambda key, value: Error(...)

We could start with a single handler that solve this specific issue.