Repository: pytorch
Issue: Some module has incorrect scope in complex naming situations in tensorboard graph export
Labels: triaged,module: tensorboard

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

My example is rather complex. To reproduce, I created a simple example.

```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.weight = nn.Linear(5, 3)
    self.bias = nn.Linear(5, 3)
    self.module = nn.Linear(6, 1)

  def forward(self, x):
    tensors = [self.weight(x), self.bias(x)]
    self.module(torch.cat(tensors, dim=1))
    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))
```

Output is:

```json
{
  "node": [
    {
      "attr": "...",
      "name": "input/input.1",
      "op": "IO Node"
    },
    {
      "attr": "...",
      "input": [
        "input/input.1"
      ],
      "name": "output/output.1",
      "op": "IO Node"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[weight]/weight/35"
      ],
      "name": "MyModule/Linear[weight]/bias/49",
      "op": "prim::GetAttr"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[weight]/weight/35"
      ],
      "name": "MyModule/Linear[weight]/weight/50",
      "op": "prim::GetAttr"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[weight]/weight/50"
      ],
      "name": "MyModule/Linear[weight]/51",
      "op": "aten::t"
    },
    {
      "attr": "...",
      "name": "MyModule/Linear[weight]/52",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "name": "MyModule/Linear[weight]/53",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[weight]/bias/49",
        "input/input.1",
        "MyModule/Linear[weight]/51",
        "MyModule/Linear[weight]/52",
        "MyModule/Linear[weight]/53"
      ],
      "name": "MyModule/Linear[weight]/54",
      "op": "aten::addmm"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[weight]/bias/bias/weight/38"
      ],
      "name": "MyModule/Linear[weight]/bias/bias/55",
      "op": "prim::GetAttr"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[weight]/bias/bias/weight/38"
      ],
      "name": "MyModule/Linear[weight]/bias/bias/weight/56",
      "op": "prim::GetAttr"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[weight]/bias/bias/weight/56"
      ],
      "name": "MyModule/Linear[bias]/57",
      "op": "aten::t"
    },
    {
      "attr": "...",
      "name": "MyModule/Linear[bias]/58",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "name": "MyModule/Linear[bias]/59",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[weight]/bias/bias/55",
        "input/input.1",
        "MyModule/Linear[bias]/57",
        "MyModule/Linear[bias]/58",
        "MyModule/Linear[bias]/59"
      ],
      "name": "MyModule/Linear[bias]/60",
      "op": "aten::addmm"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[weight]/54",
        "MyModule/Linear[bias]/60"
      ],
      "name": "MyModule/23",
      "op": "prim::ListConstruct"
    },
    {
      "attr": "...",
      "name": "MyModule/24",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/23",
        "MyModule/24"
      ],
      "name": "MyModule/input",
      "op": "aten::cat"
    },
    {
      "attr": "...",
      "name": "MyModule/61",
      "op": "prim::Constant"
    }
  ],
  "versions": {
    "producer": 22
  }
}
```

Notice that `bias/bias/weight` there. It's obviously wrong.

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

Code responsible for the error can be found in `_pytorch_graph.py`. It looks like this:

```python
    for node in graph.nodes():
        if node.kind() == GETATTR_KIND:
            attr_name = node.s('name')
            parent = node.input().node()
            if parent.kind() == GETATTR_KIND:  # If the parent node is not the top-level "self" node
                parent_attr_name = parent.s('name')
                parent_scope = attr_to_scope[parent_attr_name]
                attr_scope = parent_scope.split('/')[-1]
                attr_to_scope[attr_name] = '{}/{}.{}'.format(parent_scope, attr_scope, attr_name)
            else:
                attr_to_scope[attr_name] = '__module.{}'.format(attr_name)
            # We don't need classtype nodes; scope will provide this information
            if node.output().type().kind() != CLASSTYPE_KIND:
                node_py = NodePyOP(node)
                node_py.scopeName = attr_to_scope[attr_name]
                nodes_py.append(node_py)
        else:
            nodes_py.append(NodePyOP(node))
```

It used `attr_name` as key to find scope. However `attr_name` is often `weight` and `bias` and very easy to induce collision.

I tried to use node id as key, but failed to find any docs (nor related source code). So I use the first part of str (which is `%xxx`) as key. The modified code is:

```python
    attr_to_scope = dict()
    node_to_name = lambda d: str(d).split(":")[0].strip()
    for node in graph.nodes():
        if node.kind() == GETATTR_KIND:
            attr_name = node.s('name')
            node_name = node_to_name(node)
            parent = node.input().node()
            if parent.kind() == GETATTR_KIND:  # If the parent node is not the top-level "self" node
                parent_attr_name = parent.s('name')
                parent_scope = attr_to_scope[node_to_name(parent)]
                attr_scope = parent_scope.split('/')[-1]
                attr_to_scope[node_name] = '{}/{}.{}'.format(parent_scope, attr_scope, attr_name)
            else:
                attr_to_scope[node_name] = '__module.{}'.format(attr_name)
            # We don't need classtype nodes; scope will provide this information
            if node.output().type().kind() != CLASSTYPE_KIND:
                node_py = NodePyOP(node)
                node_py.scopeName = attr_to_scope[node_name]
                nodes_py.append(node_py)
        else:
            nodes_py.append(NodePyOP(node))
```

Now the result is:

```json
{
  "node": [
    {
      "attr": "...",
      "name": "input/input.1",
      "op": "IO Node"
    },
    {
      "attr": "...",
      "input": [
        "input/input.1"
      ],
      "name": "output/output.1",
      "op": "IO Node"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[weight]/weight/35"
      ],
      "name": "MyModule/Linear[weight]/bias/49",
      "op": "prim::GetAttr"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[weight]/weight/35"
      ],
      "name": "MyModule/Linear[weight]/weight/50",
      "op": "prim::GetAttr"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[weight]/weight/50"
      ],
      "name": "MyModule/Linear[weight]/51",
      "op": "aten::t"
    },
    {
      "attr": "...",
      "name": "MyModule/Linear[weight]/52",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "name": "MyModule/Linear[weight]/53",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[weight]/bias/49",
        "input/input.1",
        "MyModule/Linear[weight]/51",
        "MyModule/Linear[weight]/52",
        "MyModule/Linear[weight]/53"
      ],
      "name": "MyModule/Linear[weight]/54",
      "op": "aten::addmm"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[bias]/weight/38"
      ],
      "name": "MyModule/Linear[bias]/bias/55",
      "op": "prim::GetAttr"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[bias]/weight/38"
      ],
      "name": "MyModule/Linear[bias]/weight/56",
      "op": "prim::GetAttr"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[bias]/weight/56"
      ],
      "name": "MyModule/Linear[bias]/57",
      "op": "aten::t"
    },
    {
      "attr": "...",
      "name": "MyModule/Linear[bias]/58",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "name": "MyModule/Linear[bias]/59",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[bias]/bias/55",
        "input/input.1",
        "MyModule/Linear[bias]/57",
        "MyModule/Linear[bias]/58",
        "MyModule/Linear[bias]/59"
      ],
      "name": "MyModule/Linear[bias]/60",
      "op": "aten::addmm"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/Linear[weight]/54",
        "MyModule/Linear[bias]/60"
      ],
      "name": "MyModule/23",
      "op": "prim::ListConstruct"
    },
    {
      "attr": "...",
      "name": "MyModule/24",
      "op": "prim::Constant"
    },
    {
      "attr": "...",
      "input": [
        "MyModule/23",
        "MyModule/24"
      ],
      "name": "MyModule/input",
      "op": "aten::cat"
    },
    {
      "attr": "...",
      "name": "MyModule/61",
      "op": "prim::Constant"
    }
  ],
  "versions": {
    "producer": 22
  }
}
```

## Environment

Not important.

## Additional context

<!-- Add any other context about the problem here. -->

Looks like this parse function is buggy. #33670