[Bug] Semantics of ndata and edata is confusing when updates are involved
- Dominant language
- Python
- Stars
- 14.3k
- Forks
- 3.1k
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
Make `g.ndata` and `g.edata` dict-of-dicts for heterogeneous graphs with the following structures:
```
g.ndata[feature_name][ntype]
g.edata[feature_name][etype]
```
Meanwhile, deprecate `g.nodes[ntype].data` and `g.edges[etype].data`.
## Motivation
Currently `g.ndata/g.edata` is already a dict-of-dict with inner keys node/edge types and outer keys feature names.
```python
>>> g = dgl.heterograph({
... ('A', 'AB', 'B'): (torch.randint(0, 3, (3,)), torch.randint(0, 4, (3,))),
... ('A', 'AC', 'C'): (torch.randint(0, 5, (6,)), torch.randint(0, 5, (6,)))})
>>> g.nodes['A'].data['x'] = torch.randn(g.num_nodes('A'), 4)
>>> g.nodes['B'].data['x'] = torch.randn(g.num_nodes('B'), 5)
>>> g.nodes['C'].data['y'] = torch.randn(g.num_nodes('C'), 6)
>>> g.ndata
defaultdict(, {'x': {'A': tensor([[ 0.5082, 0.3065, 0.2581, -1.9657],
[-0.5501, -0.1200, -0.1060, -0.7692],
[ 1.7631, 1.1369, -0.6438, 2.7532],
[-0.3079, -1.1753, -1.6741, -1.5737]]), 'B': tensor([[ 1.0978, 0.9534, -0.6323, -0.4815, -0.8591],
[ 0.6285, -1.0719, 1.5144, 0.4381, -0.6814],
[ 0.2598, -0.2076, 0.8999, -0.5944, 0.2032],
[-0.1002, 0.3376, -0.3126, 0.6450, 1.5637]])}, 'y': {'C': tensor([[-1.0204, -0.1597, -1.3869, 0.1972, -0.2861, 1.3119],
[-0.2160, -1.8504, -0.9787, -1.4258, 0.8731, 0.6307],
[-0.2794, -1.3630, -1.4429, -0.2576, -0.4704, 1.9810],
[-1.8767, 1.6221, -0.9193, -0.2034, -0.0081, -0.7879],
[-1.8845, 0.0457, -2.3964, 0.1123, 0.1346, -1.2281]])}})
```
This is very good since it can conveniently retrieve the features of different node/edge types with the same name:
```
>>> g.ndata['x']
{'A': tensor([[ 0.5082, 0.3065, 0.2581, -1.9657],
[-0.5501, -0.1200, -0.1060, -0.7692],
[ 1.7631, 1.1369, -0.6438, 2.7532],
[-0.3079, -1.1753, -1.6741, -1.5737]]),
'B': tensor([[ 1.0978, 0.9534, -0.6323, -0.4815, -0.8591],
[ 0.6285, -1.0719, 1.5144, 0.4381, -0.6814],
[ 0.2598, -0.2076, 0.8999, -0.5944, 0.2032],
[-0.1002, 0.3376, -0.3126, 0.6450, 1.5637]])}
```
which aligns well with our NN modules that takes in a dictionary of features whose keys are also types, especially `HeteroGraphConv`:
```python
module = nn.HeteroGraphConv(...)
module(g, g.ndata['x']) # good
```
The problem is, when people think that `g.ndata` is a dict-of-dict and inserts elements directly, they will find that this is not allowed.
```python
>>> g.ndata['x']['C'] = torch.randn(g.num_nodes('C'), 2)
>>> g.ndata['x'] # ?
{'A': tensor([[ 0.5082, 0.3065, 0.2581, -1.9657],
[-0.5501, -0.1200, -0.1060, -0.7692],
[ 1.7631, 1.1369, -0.6438, 2.7532],
[-0.3079, -1.1753, -1.6741, -1.5737]]),
'B': tensor([[ 1.0978, 0.9534, -0.6323, -0.4815, -0.8591],
[ 0.6285, -1.0719, 1.5144, 0.4381, -0.6814],
[ 0.2598, -0.2076, 0.8999, -0.5944, 0.2032],
[-0.1002, 0.3376, -0.3126, 0.6450, 1.5637]])}
```
However, `g.ndata.update` works:
```
>>> g.ndata.update({'x': {'C': torch.randn(g.num_nodes('C'), 2)}})
>>> g.ndata['x']
{'A': tensor([[ 0.5082, 0.3065, 0.2581, -1.9657],
[-0.5501, -0.1200, -0.1060, -0.7692],
[ 1.7631, 1.1369, -0.6438, 2.7532],
[-0.3079, -1.1753, -1.6741, -1.5737]]),
'B': tensor([[ 1.0978, 0.9534, -0.6323, -0.4815, -0.8591],
[ 0.6285, -1.0719, 1.5144, 0.4381, -0.6814],
[ 0.2598, -0.2076, 0.8999, -0.5944, 0.2032],
[-0.1002, 0.3376, -0.3126, 0.6450, 1.5637]]),
'C': tensor([[ 1.5628, 0.7909],
[ 0.8747, 0.6824],
[-0.1325, 1.5887],
[ 1.2232, 0.2412],
[ 0.7951, -0.1124]])}
```
The reason is that when calling `g.ndata.__getitem__` in a heterogeneous graph, the returned dictionary containing the node types and tensors is constructed on-the-fly. So `g.ndata['x'].__setitem__` will not reflect into `g.ndata`.
Moreover, `g.ndata` does not support `keys()`, `values()`, `__len__()` when it's a heterogeneous graph (it does support `items()`). So it's not entirely a `defaultdict`, which is confusing since `__repr__` indeed shows a `defaultdict`.
## Alternatives
Another way is to change `__repr__` so that it's not outputting a `defaultdict`. I'd consider this a much inferior solution since the usage of `g.ndata['x'][ntype] = value` is quite intuitive.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.