Fix PyroModule registration of frozen parameters

pyro/nn/module.py

@@ -629,12 +629,14 @@ def __getattr__(self, name: str) -> Any:
         result = super().__getattr__(name)
 
         # Regular nn.Parameters trigger pyro.param statements.
-        if isinstance(result, torch.nn.Parameter) and not name.endswith(
-            "_unconstrained"
+        if (
+            isinstance(result, torch.nn.Parameter)
+            and not name.endswith("_unconstrained")
+            and self._pyro_context.active
         ):
-            if self._pyro_context.active and not _is_module_local_param_enabled():
+            if result.requires_grad and not _is_module_local_param_enabled():
                 pyro.param(self._pyro_get_fullname(name), result)
-            elif self._pyro_context.active and _is_module_local_param_enabled():
+            elif result.requires_grad:
                 # fake param statement to ensure any handlers of pyro.param are applied,
                 # even though we don't use the contents of the local parameter store
                 fullname = self._pyro_get_fullname(name)
@@ -740,15 +742,23 @@ def __setattr__(
                 delattr(self, name)
             except AttributeError:
                 pass
-            if self._pyro_context.active and not _is_module_local_param_enabled():
+            if (
+                self._pyro_context.active
+                and value.requires_grad
+                and not _is_module_local_param_enabled()
+            ):
                 fullname = self._pyro_get_fullname(name)
                 value = pyro.param(fullname, value)
                 if not isinstance(value, torch.nn.Parameter):
                     # Update PyroModule ---> ParamStore (type only; data is preserved).
                     value = torch.nn.Parameter(detach_provenance(value))
                     _PYRO_PARAM_STORE._params[fullname] = value
                     _PYRO_PARAM_STORE._param_to_name[value] = fullname
-            elif self._pyro_context.active and _is_module_local_param_enabled():
+            elif (
+                self._pyro_context.active
+                and value.requires_grad
+                and _is_module_local_param_enabled()
+            ):
                 # fake param statement to ensure any handlers of pyro.param are applied,
                 # even though we don't use the contents of the local parameter store
                 fullname = self._pyro_get_fullname(name)

tests/nn/test_module.py

@@ -69,6 +69,36 @@ def forward(self, *args, **kwargs):
         svi.step(data)
 
 
+@pytest.mark.parametrize("local_params", [False, True])
+def test_frozen_parameters_are_not_registered(local_params):
+    class Model(PyroModule):
+        def __init__(self):
+            super().__init__()
+            self.trainable = nn.Parameter(torch.zeros(1))
+            self.frozen = nn.Parameter(torch.ones(1), requires_grad=False)
+
+        def forward(self):
+            self.dynamic_frozen = nn.Parameter(
+                torch.full((1,), 2.0), requires_grad=False
+            )
+            return self.trainable + self.frozen + self.dynamic_frozen
+
+    with pyro.settings.context(module_local_params=local_params):
+        model = Model()
+        trace = poutine.trace(model).get_trace()
+
+    assert_equal(trace.nodes["_RETURN"]["value"], torch.tensor([3.0]))
+    assert "trainable" in trace.nodes
+    assert "frozen" not in trace.nodes
+    assert "dynamic_frozen" not in trace.nodes
+    assert model.frozen.requires_grad is False
+    assert model.dynamic_frozen.requires_grad is False
+    if local_params:
+        assert not pyro.get_param_store()
+    else:
+        assert set(pyro.get_param_store()) == {"trainable"}
+
+
 @pytest.mark.parametrize("local_params", [True, False])
 @pytest.mark.parametrize("num_particles", [1, 2])
 @pytest.mark.parametrize("vectorize_particles", [True, False])