Eagerly bind submodules #3077
Eagerly bind submodules #3077
Conversation
cgarciae
force-pushed
the
eager-bind-submodules
branch
from
9ee0ba8
to
96fe459
Compare
Codecov Report
@@ Coverage Diff @@
## main #3077 +/- ##
==========================================
+ Coverage 81.94% 82.07% +0.12%
==========================================
Files 55 55
Lines 6037 6058 +21
==========================================
+ Hits 4947 4972 +25
+ Misses 1090 1086 -4
|
cgarciae
force-pushed
the
eager-bind-submodules
branch
3 times, most recently
from
0d554d0
to
0a2ae27
Compare
flax/linen/module.py
Outdated
| for field in dataclasses.fields(self): | ||
| if ( | ||
| field.name not in ('parent', 'name') | ||
| and field.name in self.__dict__ # ignore fields that have not been set |
There was a problem hiding this comment.
why not and hasattr(self, field.name) ?
flax/linen/module.py
Outdated
| and field.name in self.__dict__ # ignore fields that have not been set | ||
| ): | ||
| value = getattr(self, field.name) | ||
| if _tree_has_modules(value): |
There was a problem hiding this comment.
Could you tell me again why we're doing _tree_has_modules when _register_submodules implements a recursion over submodules? is that missing something?
|
also the test |
cgarciae
force-pushed
the
eager-bind-submodules
branch
3 times, most recently
from
f9becd3
to
0bd31b6
Compare
cgarciae
force-pushed
the
eager-bind-submodules
branch
from
0bd31b6
to
8f122f6
Compare
|
Original author of #2864 here, just wanted to pop in and say thanks so much @cgarciae! I've actually been following your series of PRs for the past few months lol, and I can't wait for this to get merged and released so I can finally implement my design that's been blocked by #2864 this whole time. It will simplify huge swaths of code 😌 |
|
@kvablack - so very sorry that this took so long to merge - this is a subtle change that required checking a lot of internal code to make sure that we didn't break anyone. Everything seems good so far. Hopefully I won't be forced to roll back anything. @cgarciae extra thanks for helping plow through this rather annoyingly intricate problem. |
What does this PR do?
Alternative to #2901. Fixes #2864.
__post_init__if ascopeis available using_register_submodule._deep_cloneflag toModule.clonesuch that it correctly clones external submodules from dataclass fields when entering a new context (e.g.init/apply/bind) while preserving shared identities.