Repository: pyro-ppl/pyro

Issue:
Using PyroModule with torch.nn.ModueList fails for nested modules
Hi,

I already have made a [PR](https://github.com/pyro-ppl/pyro/pull/3339) but it seems like it just does not get any attention without an issue related to it. So here is the issue to it:

While using "nested" `PyroModule`'s: A `PyroModule` with a `PyroModule[torch.nn.ModuleList]` argument, containing another such module (see the test implementation in the PR for a more detailed example), I encountered some errors related to the existence of multiple sample sites with the same name.

I spend some time on tracing the issue and found that the `RuntimeError` was caused by an unlucky combination of `PyroModule` and `torch.nn.ModuleList`:
First: why using `torch.nn.ModuleList`? Well I wanted to model a `PyroModule` that owns a list of sub-pyromodules. My idea was to replace `linear = PyroModule[Linear](5, 2)` from the [modules example](https://pyro.ai/examples/modules.html) by `PyroModule[torch.nn.ModuleList](...)`.
That works fine if there is a
```python
class MyPyroModule(PyroModule):
    def __init__(self, some_module: PyroModule):
        self.my_submodule = PyroModule[torch.nn.ModuleList]([some_module for _ in range(3)])
```
but **can** fail if there is a nested structure, like
```python
my_nested_pyro_module = MyPyroModule(MyPyroModule(...))
```
The "can fail" could be resolved to the following different types if *accessing* the `self.my_submodule` argument:
+ using a index-positioned access, like `self.my_submodule[0](...)` worked fine, even for nested modules
+ using slice-indexing access, like `self.my_submodule[:-1](...)` only works if there is a single `MyPyroModule()` and not a `MyPyroModule(MyPyroModule(...))`

The cause is [this](https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/container.py#L293) line in the `torch.nn.ModuleList` class:
```python
@_copy_to_script_wrapper
    def __getitem__(self, idx: Union[int, slice]) -> Union[Module, 'ModuleList']:
        if isinstance(idx, slice):
            return self.__class__(list(self._modules.values())[idx])
        else:
            return self._modules[self._get_abs_string_index(idx)]
```
which calls `self.__class__`, which, means that for an object of type `PyroModule[torch.nn.ModuleList]`, it calls the `PyroModule.__init__` function, without the context of the parent module. This results in overwriting the names of the module's submodules and their `._pyro_name` attributes, and because of that, during sampling, sample sites may not be unique anymore.

I see two possible fixes for this:
1. I have introduced a `pyro.nn.PyroModuleList` class in this PR that inherits from `torch.nn.ModuleList` can overwrites the problematic `__getitem__` function
2. We may add some example/documentation on this issue, explicitly warning about using `PyroModule[torch.nn.ModuleList]` in combination with slice-indexing (feels a little unsafe to me)
3. Maybe some has another (potentially better) idea of how to deal with this?

I know that this is kind of a "special purpose" usage and may not affect basic `pyro` usage, but in particular the way the [modules example](https://docs.pyro.ai/en/stable/nn.html) motivates the usage of `PyroModule[torch.nn.Something]` could, imho, quickly lure other users into this (and it took me quite some time to find the root issue).

Please let me know what you think about this PR, and whether it needs updates or clarification.

Best regards
Martin


Discussion:

Resolving patch:
diff --git a/pyro/nn/__init__.py b/pyro/nn/__init__.py
--- a/pyro/nn/__init__.py
+++ b/pyro/nn/__init__.py
@@ -9,7 +9,13 @@
     MaskedLinear,
 )
 from pyro.nn.dense_nn import ConditionalDenseNN, DenseNN
-from pyro.nn.module import PyroModule, PyroParam, PyroSample, pyro_method
+from pyro.nn.module import (
+    PyroModule,
+    PyroModuleList,
+    PyroParam,
+    PyroSample,
+    pyro_method,
+)
 
 __all__ = [
     "AutoRegressiveNN",
@@ -21,4 +27,5 @@
     "PyroParam",
     "PyroSample",
     "pyro_method",
+    "PyroModuleList",
 ]
diff --git a/pyro/nn/module.py b/pyro/nn/module.py
--- a/pyro/nn/module.py
+++ b/pyro/nn/module.py
@@ -14,7 +14,21 @@
 """
 import functools
 import inspect
+import warnings
 import weakref
+
+try:
+    from torch._jit_internal import _copy_to_script_wrapper
+except ImportError:
+    warnings.warn(
+        "Cannot find torch._jit_internal._copy_to_script_wrapper", ImportWarning
+    )
+
+    # Fall back to trivial decorator.
+    def _copy_to_script_wrapper(fn):
+        return fn
+
+
 from collections import OrderedDict
 from dataclasses import dataclass
 from types import TracebackType
@@ -902,3 +916,32 @@ def __set__(self, obj: object, value: Any) -> None:
 
 
 PyroModule[torch.nn.RNNBase]._flat_weights = _FlatWeightsDescriptor()  # type: ignore[attr-defined]
+
+
+# pyro module list
+# using pyro.nn.PyroModule[torch.nn.ModuleList] can cause issues when
+# slice-indexing nested PyroModuleLists, so we define a separate PyroModuleList
+# class that overwrites the __getitem__ method to return a torch.nn.ModuleList
+# to not use self.__class__ in __getitem__, as that would call the
+# PyroModule.__init__ without the parent module context, leading to a loss
+# of the parent module's _pyro_name, and eventually, errors during sampling
+# as parameter names may not be unique anymore
+# The scenario is rare but happend.
+# The fix could not be applied in torch directly, which is why we have to deal
+# with it here, see https://github.com/pytorch/pytorch/issues/121008
+class PyroModuleList(torch.nn.ModuleList, PyroModule):
+    def __init__(self, modules):
+        super().__init__(modules)
+
+    @_copy_to_script_wrapper
+    def __getitem__(
+        self, idx: Union[int, slice]
+    ) -> Union[torch.nn.Module, "PyroModuleList"]:
+        if isinstance(idx, slice):
+            # return self.__class__(list(self._modules.values())[idx])
+            return torch.nn.ModuleList(list(self._modules.values())[idx])
+        else:
+            return self._modules[self._get_abs_string_index(idx)]
+
+
+_PyroModuleMeta._pyro_mixin_cache[torch.nn.ModuleList] = PyroModuleList


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
@@ -2,7 +2,9 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import io
+import math
 import warnings
+from typing import Callable, Iterable
 
 import pytest
 import torch
@@ -13,6 +15,7 @@
 import pyro.distributions as dist
 from pyro import poutine
 from pyro.infer import SVI, Trace_ELBO
+from pyro.infer.autoguide.guides import AutoDiagonalNormal
 from pyro.nn.module import PyroModule, PyroParam, PyroSample, clear, to_pyro_module_
 from pyro.optim import Adam
 from tests.common import assert_equal, xfail_param
@@ -844,3 +847,201 @@ def forward(self, x, y):
                 grad_params_func[k], torch.zeros_like(grad_params_func[k])
             ), k
             assert torch.allclose(grad_params_autograd[k], grad_params_func[k]), k
+
+
+class BNN(PyroModule):
+    # this is a vanilla Bayesian neural network implementation, nothing new or exiting here
+    def __init__(
+        self,
+        input_size: int,
+        hidden_layer_sizes: Iterable[int],
+        output_size: int,
+        use_new_module_list_type: bool,
+    ) -> None:
+        super().__init__()
+
+        layer_sizes = (
+            [(input_size, hidden_layer_sizes[0])]
+            + list(zip(hidden_layer_sizes[:-1], hidden_layer_sizes[1:]))
+            + [(hidden_layer_sizes[-1], output_size)]
+        )
+
+        layers = [
+            pyro.nn.module.PyroModule[torch.nn.Linear](in_size, out_size)
+            for in_size, out_size in layer_sizes
+        ]
+        if use_new_module_list_type:
+            self.layers = pyro.nn.module.PyroModuleList(layers)
+        else:
+            self.layers = pyro.nn.module.PyroModule[torch.nn.ModuleList](layers)
+
+        # make the layers Bayesian
+        for layer_idx, layer in enumerate(self.layers):
+            layer.weight = pyro.nn.module.PyroSample(
+                dist.Normal(0.0, 5.0 * math.sqrt(2 / layer_sizes[layer_idx][0]))
+                .expand(
+                    [
+                        layer_sizes[layer_idx][1],
+                        layer_sizes[layer_idx][0],
+                    ]
+                )
+                .to_event(2)
+            )
+            layer.bias = pyro.nn.module.PyroSample(
+                dist.Normal(0.0, 5.0).expand([layer_sizes[layer_idx][1]]).to_event(1)
+            )
+
+        self.activation = torch.nn.Tanh()
+        self.output_size = output_size
+
+    def forward(self, x: torch.Tensor, obs=None) -> torch.Tensor:
+        mean = self.layers[-1](x)
+
+        if obs is not None:
+            with pyro.plate("data", x.shape[0]):
+                pyro.sample(
+                    "obs", dist.Normal(mean, 0.1).to_event(self.output_size), obs=obs
+                )
+
+        return mean
+
+
+class SliceIndexingModuleListBNN(BNN):
+    # I claim that it makes a difference whether slice-indexing is used or whether position-indexing is used
+    # when sub-pyromodule are wrapped in a PyroModule[torch.nn.ModuleList]
+    def __init__(
+        self,
+        input_size: int,
+        hidden_layer_sizes: Iterable[int],
+        output_size: int,
+        use_new_module_list_type: bool,
+    ) -> None:
+        super().__init__(
+            input_size, hidden_layer_sizes, output_size, use_new_module_list_type
+        )
+
+    def forward(self, x: torch.Tensor, obs=None) -> torch.Tensor:
+        for layer in self.layers[:-1]:
+            x = layer(x)
+            x = self.activation(x)
+
+        return super().forward(x, obs=obs)
+
+
+class PositionIndexingModuleListBNN(BNN):
+    # I claim that it makes a difference whether slice-indexing is used or whether position-indexing is used
+    # when sub-pyromodule are wrapped in a PyroModule[torch.nn.ModuleList]
+    def __init__(
+        self,
+        input_size: int,
+        hidden_layer_sizes: Iterable[int],
+        output_size: int,
+        use_new_module_list_type: bool,
+    ) -> None:
+        super().__init__(
+            input_size, hidden_layer_sizes, output_size, use_new_module_list_type
+        )
+
+    def forward(self, x: torch.Tensor, obs=None) -> torch.Tensor:
+        for i in range(len(self.layers) - 1):
+            x = self.layers[i](x)
+            x = self.activation(x)
+
+        return super().forward(x, obs=obs)
+
+
+class NestedBNN(pyro.nn.module.PyroModule):
+    # finally, the issue I want to describe occurs after the second "layer of nesting",
+    # i.e. when a PyroModule[ModuleList] is wrapped in a PyroModule[ModuleList]
+    def __init__(self, bnns: Iterable[BNN], use_new_module_list_type: bool) -> None:
+        super().__init__()
+        if use_new_module_list_type:
+            self.bnns = pyro.nn.module.PyroModuleList(bnns)
+        else:
+            self.bnns = pyro.nn.module.PyroModule[torch.nn.ModuleList](bnns)
+
+    def forward(self, x: torch.Tensor, obs=None) -> torch.Tensor:
+        mean = sum([bnn(x) for bnn in self.bnns]) / len(self.bnns)
+
+        with pyro.plate("data", x.shape[0]):
+            pyro.sample("obs", dist.Normal(mean, 0.1).to_event(1), obs=obs)
+
+        return mean
+
+
+def train_bnn(model: BNN, input_size: int) -> None:
+    pyro.clear_param_store()
+
+    # small numbers for demo purposes
+    num_points = 20
+    num_svi_iterations = 100
+
+    x = torch.linspace(0, 1, num_points).reshape((-1, input_size))
+    y = torch.sin(2 * math.pi * x) + torch.randn(x.size()) * 0.1
+
+    guide = AutoDiagonalNormal(model)
+    adam = pyro.optim.Adam({"lr": 0.03})
+    svi = SVI(model, guide, adam, loss=Trace_ELBO())
+
+    for _ in range(num_svi_iterations):
+        svi.step(x, y)
+
+
+class ModuleListTester:
+    def setup(self, use_new_module_list_type: bool) -> None:
+        self.input_size = 1
+        self.output_size = 1
+        self.hidden_size = 3
+        self.num_hidden_layers = 3
+        self.use_new_module_list_type = use_new_module_list_type
+
+    def get_position_indexing_modulelist_bnn(self) -> PositionIndexingModuleListBNN:
+        return PositionIndexingModuleListBNN(
+            self.input_size,
+            [self.hidden_size] * self.num_hidden_layers,
+            self.output_size,
+            self.use_new_module_list_type,
+        )
+
+    def get_slice_indexing_modulelist_bnn(self) -> SliceIndexingModuleListBNN:
+        return SliceIndexingModuleListBNN(
+            self.input_size,
+            [self.hidden_size] * self.num_hidden_layers,
+            self.output_size,
+            self.use_new_module_list_type,
+        )
+
+    def train_nested_bnn(self, module_getter: Callable[[], BNN]) -> None:
+        train_bnn(
+            NestedBNN(
+                [module_getter() for _ in range(2)],
+                use_new_module_list_type=self.use_new_module_list_type,
+            ),
+            self.input_size,
+        )
+
+
+class TestTorchModuleList(ModuleListTester):
+    def test_with_position_indexing(self) -> None:
+        self.setup(False)
+        self.train_nested_bnn(self.get_position_indexing_modulelist_bnn)
+
+    def test_with_slice_indexing(self) -> None:
+        self.setup(False)
+        # with pytest.raises(RuntimeError):
+        # error no longer gets raised
+        self.train_nested_bnn(self.get_slice_indexing_modulelist_bnn)
+
+
+class TestPyroModuleList(ModuleListTester):
+    def test_with_position_indexing(self) -> None:
+        self.setup(True)
+        self.train_nested_bnn(self.get_position_indexing_modulelist_bnn)
+
+    def test_with_slice_indexing(self) -> None:
+        self.setup(True)
+        self.train_nested_bnn(self.get_slice_indexing_modulelist_bnn)
+
+
+def test_module_list() -> None:
+    assert PyroModule[torch.nn.ModuleList] is pyro.nn.PyroModuleList