Helpful error message when accessing non-existent variables

[freddyaboulton]

Fixes #513. Following the discussion in #513, Scope variables and apply now return variables as instances of NonFinalVariablesDict as opposed to FrozenDict. This will give users a helpful error message that explains that variables (if they exist) should be accessed after shape inference and shows the existing variable names.

#513 explicitly mentions the variables returned by Scope.variables but I thought it would be good to apply this change to the variables returned by Module.apply since users are likely to access those as well and would benefit from the same useful error message. Happy to walk that back though.

[codecov-io]

Codecov Report

Merging #638 (03f72df) into master (0c6a8ba) will increase coverage by 0.05%.
The diff coverage is 91.30%.

@@            Coverage Diff             @@
##           master     #638      +/-   ##
==========================================
+ Coverage   80.59%   80.65%   +0.05%     
==========================================
  Files          56       56              
  Lines        4241     4259      +18     
==========================================
+ Hits         3418     3435      +17     
- Misses        823      824       +1     

Impacted Files	Coverage Δ
flax/core/variables.py	94.73% <90.00%> (-5.27%)	⬇️
flax/core/scope.py	92.30% <100.00%> (-0.06%)	⬇️
flax/linen/module.py	95.25% <0.00%> (+0.36%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 0c6a8ba...03f72df. Read the comment docs.

[avital]

Hi @freddyaboulton -- a few comments:

1. I don't think the variables returned by init or apply should use this new type. The new type is only helpful for reducing user error /within/ modules. By the time a module is done initializing (via init) or new variables come out (via apply), then there is no more opportunity to apply shape inference. So the error is wrong for this case.
2. There should be unit tests for the FrozenVariablesDict class itself.
3. Perhaps the name is wrong, the fact that it represents variables is not the important part, it's the fact that it may be in an intermediate state. So maybe NonFinalVariablesDict?
4. Going forward we can change the error message to point folks to a design note that we'll write about shape inference and lazy submodules (long overdue!).

I also want a second review from @jheek who original proposed this idea on #513

[avital]

(Oh I forgot to say thank you for jumping on this issue so quickly!)

[freddyaboulton]

@avital Thank you for the feedback! I've addressed your comments. The CI passes locally but the github action is stuck installing dependencies. It's quite confusing since this PR does not change any dependencies. Yesterday it failed on a tensorflow_text import which I've seen on other open PRs. I'll keep an eye on it and let you know if anything changes but maybe you have a better idea of what is happening?

[avital]

We fixed the tensorflow_text problem, it was a bad version dependency. If you rebase and push now tests should pass.

[jheek]

I think this is a good idea! I wonder though if we can simplify to make this more widely applicable. Essentially we want a FrozenDict to have a helpful error message on missing keys (and potentially when mutating). However, if we do this by subclasssing we need to make a subclass for each case.

How about adding a missing_key_handler (and perhaps mutation_handler) to a FrozenDict?

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

Let me know what you think

We would probably also want error handlers like this on other frozen structures and an interface like this would allow it to be consistent.

[marcvanzee]

@freddyaboulton any update on this?

[jheek]

Closing due to inactivity