Repository: pytorch
Issue: Some module info is missing in nested graph for tensorboard
Labels: triaged,module: tensorboard

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

I'm currently doing some NN visualization project, which needs me to retrieve the precise module structure.

## To Reproduce

The output of the following code snippet is:

```python
import json
import torch
import torch.nn as nn
from torch.utils.tensorboard._pytorch_graph import graph
from google.protobuf import json_format

class MyModule(nn.Module):
  def __init__(self):
    super().__init__()
    self.module = nn.ModuleList([
      nn.Linear(5, 3),
      nn.Linear(3, 1)
    ])

  def forward(self, x):
    x = self.module[0](x)
    x = self.module[1](x)
    return x

module = MyModule()
result, _ = graph(module, torch.randn(5, 5))
result = json_format.MessageToDict(result)
for node in result["node"]:
  node["attr"] = "..."
print(json.dumps(result, sort_keys=True, indent=2))
```

```json
{
  "node": [
    {
      "attr": "...",
      "name": "input/input.1",
      "op": "IO Node"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[1]/46"
      ],
      "name": "output/output.1",
      "op": "IO Node"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[0]/weight/26"
      ],
      "name": "MyModule/ModuleList[module]/Linear[0]/bias/35",
      "op": "prim::GetAttr"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[0]/weight/26"
      ],
      "name": "MyModule/ModuleList[module]/Linear[0]/weight/36",
      "op": "prim::GetAttr"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[0]/weight/36"
      ],
      "name": "MyModule/Linear[0]/37",
      "op": "aten::t"
    },
    {
      "attr": "...",
      "name": "MyModule/Linear[0]/38",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "name": "MyModule/Linear[0]/39",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[0]/bias/35",
        "input/input.1",
        "MyModule/Linear[0]/37",
        "MyModule/Linear[0]/38",
        "MyModule/Linear[0]/39"
      ],
      "name": "MyModule/Linear[0]/input",
      "op": "aten::addmm"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[1]/weight/30"
      ],
      "name": "MyModule/ModuleList[module]/Linear[1]/bias/41",
      "op": "prim::GetAttr"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[1]/weight/30"
      ],
      "name": "MyModule/ModuleList[module]/Linear[1]/weight/42",
      "op": "prim::GetAttr"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[1]/weight/42"
      ],
      "name": "MyModule/Linear[1]/43",
      "op": "aten::t"
    },
    {
      "attr": "...",
      "name": "MyModule/Linear[1]/44",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "name": "MyModule/Linear[1]/45",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[1]/bias/41",
        "MyModule/Linear[0]/input",
        "MyModule/Linear[1]/43",
        "MyModule/Linear[1]/44",
        "MyModule/Linear[1]/45"
      ],
      "name": "MyModule/Linear[1]/46",
      "op": "aten::addmm"
    }
  ],
  "versions": {
    "producer": 22
  }
}
```

Notice that there is no module as `MyModule/Linear[1]/46` in my model. It should be `MyModule/ModuleList[module]/Linear[1]/46`.

## Expected behavior

Actually, the cause of this problem is that there are some levels of modules missing in the jit trace. But since each level of jit trace is a full path, the workaround is easy: modify the last lines of `parse` function of `_pytorch_graph` into:

```python
  for node in nodes_py.nodes_op:
    module_aliases = node.scopeName.split('/')[-1].split('.')
    module_name = ''
    for i, alias in enumerate(module_aliases):
      if i == 0:
        module_name = alias
        node.scopeName = base_name
      else:
        module_name += '.' + alias
        node.scopeName += '/' + (alias_to_name[module_name] if module_name in alias_to_name else alias)

  nodes_py.populate_namespace_from_OP_to_IO()
  return nodes_py.to_proto()
```

The code reads info after the last slash only, and use dot as break to retrieve the info from each level. Now it works as expected.

```json
{
  "node": [
    {
      "attr": "...",
      "name": "input/input.1",
      "op": "IO Node"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[1]/46"
      ],
      "name": "output/output.1",
      "op": "IO Node"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[0]/weight/26"
      ],
      "name": "MyModule/ModuleList[module]/Linear[0]/bias/35",
      "op": "prim::GetAttr"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[0]/weight/26"
      ],
      "name": "MyModule/ModuleList[module]/Linear[0]/weight/36",
      "op": "prim::GetAttr"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[0]/weight/36"
      ],
      "name": "MyModule/ModuleList[module]/Linear[0]/37",
      "op": "aten::t"
    },
    {
      "attr": "...",
      "name": "MyModule/ModuleList[module]/Linear[0]/38",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "name": "MyModule/ModuleList[module]/Linear[0]/39",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[0]/bias/35",
        "input/input.1",
        "MyModule/ModuleList[module]/Linear[0]/37",
        "MyModule/ModuleList[module]/Linear[0]/38",
        "MyModule/ModuleList[module]/Linear[0]/39"
      ],
      "name": "MyModule/ModuleList[module]/Linear[0]/input",
      "op": "aten::addmm"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[1]/weight/30"
      ],
      "name": "MyModule/ModuleList[module]/Linear[1]/bias/41",
      "op": "prim::GetAttr"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[1]/weight/30"
      ],
      "name": "MyModule/ModuleList[module]/Linear[1]/weight/42",
      "op": "prim::GetAttr"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[1]/weight/42"
      ],
      "name": "MyModule/ModuleList[module]/Linear[1]/43",
      "op": "aten::t"
    },
    {
      "attr": "...",
      "name": "MyModule/ModuleList[module]/Linear[1]/44",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "name": "MyModule/ModuleList[module]/Linear[1]/45",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/ModuleList[module]/Linear[1]/bias/41",
        "MyModule/ModuleList[module]/Linear[0]/input",
        "MyModule/ModuleList[module]/Linear[1]/43",
        "MyModule/ModuleList[module]/Linear[1]/44",
        "MyModule/ModuleList[module]/Linear[1]/45"
      ],
      "name": "MyModule/ModuleList[module]/Linear[1]/46",
      "op": "aten::addmm"
    }
  ]
}
```
<!-- A clear and concise description of what you expected to happen. -->

## Environment

It's not important.

## Additional context

<!-- Add any other context about the problem here. -->

I think the root cause lies in jit.trace, which might be worth investigation. Or it's by design, in which case it should be fixed in support for tensorboard.

This issue is similar to #30812, but I double-checked: it's a different issue.