numpyro.contrib.module: only collect mutables in nnx_module() when th… (#2061)

* numpyro.contrib.module: only collect mutables in nnx_module() when they exist

Testing Done: avoided registration of empty mutables in nnx ConvNet

Signed-off-by: Eli Sennesh <[email protected]>

* test.contrib.test_module: expect mutables only when batchnorm or dropout active

Signed-off-by: Eli Sennesh <[email protected]>

---------

Signed-off-by: Eli Sennesh <[email protected]>

Author: esennesh

numpyro/contrib/module.py

@@ -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:
         mutable_holder = numpyro_mutable(
             name + "$state", {"state": eager_other_state_dict}
         )

test/contrib/test_module.py

@@ -424,8 +424,12 @@ def model():
     with handlers.trace(model) as tr, handlers.seed(rng_seed=0):
         model()
 
-    assert set(tr.keys()) == {"nn$params", "nn$state", "x", "y"}
-    assert tr["nn$state"]["type"] == "mutable"
+    key_set = {"nn$params", "x", "y"}
+    if batchnorm or dropout:
+        key_set.add("nn$state")
+    assert set(tr.keys()) == key_set
+    if batchnorm or dropout:
+        assert tr["nn$state"]["type"] == "mutable"
 
     # test svi
     guide = AutoDelta(model)