Repository: pyro-ppl/pyro

Issue:
PyroModule incompatible with torch.nn.RNN
As pointed out [on the forum](https://forum.pyro.ai/t/pyromodule-pyrosample-pyroparam-and-gru/1731),  `PyroModule` is incompatible with `torch.nn.RNN` due to [internal caching logic](https://github.com/pytorch/pytorch/blob/e90c32f11fc54d017609e68fb80e19350e199864/torch/nn/modules/rnn.py#L88).

I believe a workaround is to call `my_rnn._apply(lambda t: t)` inside a supermodule's `.forward()` method. There may be a more general fix.


Discussion:

Resolving patch:
diff --git a/pyro/nn/module.py b/pyro/nn/module.py
--- a/pyro/nn/module.py
+++ b/pyro/nn/module.py
@@ -97,8 +97,10 @@ def __getitem__(cls, Module):
             return PyroModule
         if Module in _PyroModuleMeta._pyro_mixin_cache:
             return _PyroModuleMeta._pyro_mixin_cache[Module]
+        bases = [PyroModule[b] for b in Module.__bases__
+                 if issubclass(b, torch.nn.Module)]
 
-        class result(Module, PyroModule):
+        class result(Module, *bases):
             # Unpickling helper to load an object of type PyroModule[Module].
             def __reduce__(self):
                 state = getattr(self, '__getstate__', self.__dict__.copy)()
@@ -542,3 +544,20 @@ def to_pyro_module_(m, recurse=True):
         if recurse:
             to_pyro_module_(value)
         setattr(m, name, value)
+
+
+# The following descriptor disables the ._flat_weights cache of
+# torch.nn.RNNBase, forcing recomputation on each access of the ._flat_weights
+# attribute. This is required if any attribute is set to a PyroParam or
+# PyroSample. For motivation, see https://github.com/pyro-ppl/pyro/issues/2390
+class _FlatWeightsDescriptor:
+    def __get__(self, obj, obj_type=None):
+        if obj is None:
+            return self
+        return [getattr(obj, name) for name in obj._flat_weights_names]
+
+    def __set__(self, obj, value):
+        pass  # Ignore value.
+
+
+PyroModule[torch.nn.RNNBase]._flat_weights = _FlatWeightsDescriptor()


Test patch:
diff --git a/tests/nn/test_module.py b/tests/nn/test_module.py
--- a/tests/nn/test_module.py
+++ b/tests/nn/test_module.py
@@ -504,3 +504,28 @@ def forward(self):
     actual_names = {name for name, _ in actual.named_parameters()}
     expected_names = {name for name, _ in module.named_parameters()}
     assert actual_names == expected_names
+
+
+def test_bayesian_gru():
+    input_size = 2
+    hidden_size = 3
+    batch_size = 4
+    seq_len = 5
+
+    # Construct a simple GRU.
+    gru = nn.GRU(input_size, hidden_size)
+    input_ = torch.randn(seq_len, batch_size, input_size)
+    output, _ = gru(input_)
+    assert output.shape == (seq_len, batch_size, hidden_size)
+    output2, _ = gru(input_)
+    assert torch.allclose(output2, output)
+
+    # Make it Bayesian.
+    to_pyro_module_(gru)
+    for name, value in list(gru.named_parameters(recurse=False)):
+        prior = dist.Normal(0, 1).expand(value.shape).to_event(value.dim())
+        setattr(gru, name, PyroSample(prior=prior))
+    output, _ = gru(input_)
+    assert output.shape == (seq_len, batch_size, hidden_size)
+    output2, _ = gru(input_)
+    assert not torch.allclose(output2, output)