numpyro.contrib.module: only collect mutables in nnx_module() when th… - #2061
Conversation
…ey exist Testing Done: avoided registration of empty mutables in nnx ConvNet Signed-off-by: Eli Sennesh <eli.sennesh@vanderbilt.edu>
|
Thanks! I think it would be great if you could add a small test :) |
| mutable_holder = None | ||
| if eager_other_state_dict: | ||
| mutable_holder = numpyro_mutable(name + "$state") | ||
| if mutable_holder is None: |
There was a problem hiding this comment.
I guess it is better to move this if inside the above if?
There was a problem hiding this comment.
My patch deletes the lower if since by construction, mutable_holder is None is always true here. We don't want to create an empty mutable-state no matter what, since there are inference algorithms and other things that depend upon the nonexistence of mutables (which is how I found this in the first place).
There was a problem hiding this comment.
That makes sense. Thanks for clarifying! Could you adjust the failing test to cover the change? I think without batchnorm, we no longer have mutable state in that test.
There was a problem hiding this comment.
All right! Changed the test and in the no-batchnorm, no-dropout case, we no longer expect to see any 'nn$state' in the keys found.
|
Yeah, of course. I’ve had some fires to put out the past few days, but over the weekend I’m gonna look at what’s up with the test. Wouldn’t BatchNorm actually maintain mutable state throughout execution? So it should have a non-None mutable.
… On Aug 7, 2025, at 23:01, Du Phan ***@***.***> wrote:
@fehiepsi commented on this pull request.
In numpyro/contrib/module.py <#2061 (comment)>:
> @@ -484,8 +484,6 @@ def nnx_module(name, nn_module):
mutable_holder = None
if eager_other_state_dict:
- mutable_holder = numpyro_mutable(name + "$state")
- if mutable_holder is None:
That makes sense. Thanks for clarifying! Could you adjust the failing test to cover the change? I think without batchnorm, we no longer have mutable state in that test.
—
Reply to this email directly, view it on GitHub <#2061 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAFRWFJPCG6MIU7EM55RBT33MQOKBAVCNFSM6AAAAACDISLR4CVHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZTAOJZGM4DIMBQGE>.
You are receiving this because you authored the thread.
|
…out active Signed-off-by: Eli Sennesh <eli.sennesh@vanderbilt.edu>
3 participants
…ey exist
Testing Done: avoided registration of empty mutables in nnx ConvNet