Repository: pyg-team/pytorch_geometric
Pull request #8247 (2023-10-22 07:56:56)

Allow `to_hetero_with_bases` to be applied on static graphs

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8247?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#8247](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8247?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (4ee682f) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/405ef2cd74212ee061a67cb454d99994f34a95c9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (405ef2c) will **not change** coverage.
> The diff coverage is `93.33%`.

> :exclamation: Current head 4ee682f differs from pull request most recent head 6233437. Consider uploading reports for the commit 6233437 to get more accurate results

```diff
@@           Coverage Diff           @@
##           master    #8247   +/-   ##
=======================================
  Coverage   87.40%   87.40%           
=======================================
  Files         473      473           
  Lines       28632    28632           
=======================================
  Hits        25027    25027           
  Misses       3605     3605           
```


| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8247?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [...geometric/explain/algorithm/attention\_explainer.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8247?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2V4cGxhaW4vYWxnb3JpdGhtL2F0dGVudGlvbl9leHBsYWluZXIucHk=) | `93.22% <100.00%> (ø)` | |
| [torch\_geometric/nn/models/attentive\_fp.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8247?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL21vZGVscy9hdHRlbnRpdmVfZnAucHk=) | `96.00% <100.00%> (ø)` | |
| [...h\_geometric/nn/to\_hetero\_with\_bases\_transformer.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8247?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3RvX2hldGVyb193aXRoX2Jhc2VzX3RyYW5zZm9ybWVyLnB5) | `93.15% <88.88%> (ø)` | |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index afea1dfbe..22598fb75 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added support for `to_hetero_with_bases` on static graphs ([#8247](https://github.com/pyg-team/pytorch_geometric/pull/8247))
 - Added the `RCDD` dataset ([#8196](https://github.com/pyg-team/pytorch_geometric/pull/8196))
 - Added distributed `GAT + ogbn-products` example targeting XPU device ([#8032](https://github.com/pyg-team/pytorch_geometric/pull/8032))
 - Added the option to skip explanations of certain message passing layers via `conv.explain = False` ([#8216](https://github.com/pyg-team/pytorch_geometric/pull/8216))
diff --git a/torch_geometric/nn/to_hetero_with_bases_transformer.py b/torch_geometric/nn/to_hetero_with_bases_transformer.py
index 01410d967..cb787f364 100644
--- a/torch_geometric/nn/to_hetero_with_bases_transformer.py
+++ b/torch_geometric/nn/to_hetero_with_bases_transformer.py
@@ -324,7 +324,7 @@ class HeteroBasisConv(torch.nn.Module):
         # to a materialization of messages.
         def hook(module, inputs, output):
             assert isinstance(module._edge_type, Tensor)
-            if module._edge_type.size(0) != output.size(0):
+            if module._edge_type.size(0) != output.size(-2):
                 raise ValueError(
                     f"Number of messages ({output.size(0)}) does not match "
                     f"with the number of original edges "
@@ -332,7 +332,7 @@ class HeteroBasisConv(torch.nn.Module):
                     f"passing layer create additional self-loops? Try to "
                     f"remove them via 'add_self_loops=False'")
             weight = module.edge_type_weight.view(-1)[module._edge_type]
-            weight = weight.view([-1] + [1] * (output.dim() - 1))
+            weight = weight.view([1] * (output.dim() - 2) + [-1, 1])
             return weight * output
 
         params = list(module.parameters())
@@ -415,7 +415,7 @@ def get_node_offset_dict(
     out: Dict[NodeType, int] = {}
     for key in type2id.keys():
         out[key] = cumsum
-        cumsum += input_dict[key].size(0)
+        cumsum += input_dict[key].size(-2)
     return out
 
 
@@ -433,7 +433,7 @@ def get_edge_offset_dict(
         elif value.dtype == torch.long and value.size(0) == 2:
             cumsum += value.size(-1)
         else:
-            cumsum += value.size(0)
+            cumsum += value.size(-2)
     return out
 
 
@@ -458,7 +458,7 @@ def get_edge_type(
             out = torch.full((value.nnz(), ), i, dtype=torch.long,
                              device=value.device())
         else:
-            out = value.new_full((value.size(0), ), i, dtype=torch.long)
+            out = value.new_full((value.size(-2), ), i, dtype=torch.long)
         outs.append(out)
 
     return outs[0] if len(outs) == 1 else torch.cat(outs, dim=0)
@@ -474,7 +474,7 @@ def group_node_placeholder(input_dict: Dict[NodeType, Tensor],
                            type2id: Dict[NodeType, int]) -> Tensor:
 
     inputs = [input_dict[key] for key in type2id.keys()]
-    return inputs[0] if len(inputs) == 1 else torch.cat(inputs, dim=0)
+    return inputs[0] if len(inputs) == 1 else torch.cat(inputs, dim=-2)
 
 
 def group_edge_placeholder(
@@ -528,7 +528,7 @@ def group_edge_placeholder(
         return torch.stack([row, col], dim=0)
 
     else:
-        return torch.cat(inputs, dim=0)
+        return torch.cat(inputs, dim=-2)
 
 
 ###############################################################################
@@ -542,9 +542,9 @@ def split_output(
     offset_dict: Union[Dict[NodeType, int], Dict[EdgeType, int]],
 ) -> Union[Dict[NodeType, Tensor], Dict[EdgeType, Tensor]]:
 
-    cumsums = list(offset_dict.values()) + [output.size(0)]
+    cumsums = list(offset_dict.values()) + [output.size(-2)]
     sizes = [cumsums[i + 1] - cumsums[i] for i in range(len(offset_dict))]
-    outputs = output.split(sizes)
+    outputs = output.split(sizes, dim=-2)
     return {key: output for key, output in zip(offset_dict, outputs)}
 
 


Test patch:
diff --git a/test/nn/test_to_hetero_transformer.py b/test/nn/test_to_hetero_transformer.py
index 6bf1c269b..2ea69c840 100644
--- a/test/nn/test_to_hetero_transformer.py
+++ b/test/nn/test_to_hetero_transformer.py
@@ -520,10 +520,8 @@ def test_to_hetero_on_static_graphs():
         'author': torch.randn(4, 100, 16),
     }
     edge_index_dict = {
-        ('paper', 'written_by', 'author'):
-        torch.randint(100, (2, 200), dtype=torch.long),
-        ('author', 'writes', 'paper'):
-        torch.randint(100, (2, 200), dtype=torch.long),
+        ('paper', 'written_by', 'author'): torch.randint(100, (2, 200)),
+        ('author', 'writes', 'paper'): torch.randint(100, (2, 200)),
     }
 
     metadata = list(x_dict.keys()), list(edge_index_dict.keys())
diff --git a/test/nn/test_to_hetero_with_bases_transformer.py b/test/nn/test_to_hetero_with_bases_transformer.py
index 6f40f0cb2..3736ad6c6 100644
--- a/test/nn/test_to_hetero_with_bases_transformer.py
+++ b/test/nn/test_to_hetero_with_bases_transformer.py
@@ -293,3 +293,24 @@ def test_to_hetero_with_bases_validate():
 
     with pytest.warns(UserWarning, match="letters, numbers and underscores"):
         model = to_hetero_with_bases(model, metadata, num_bases=4, debug=False)
+
+
+def test_to_hetero_with_bases_on_static_graphs():
+    x_dict = {
+        'paper': torch.randn(4, 100, 16),
+        'author': torch.randn(4, 100, 16),
+    }
+    edge_index_dict = {
+        ('paper', 'written_by', 'author'): torch.randint(100, (2, 200)),
+        ('author', 'writes', 'paper'): torch.randint(100, (2, 200)),
+    }
+
+    metadata = list(x_dict.keys()), list(edge_index_dict.keys())
+    model = to_hetero_with_bases(Net4(), metadata, num_bases=4,
+                                 in_channels={'x0': 16}, debug=False)
+
+    out_dict = model(x_dict, edge_index_dict)
+
+    assert len(out_dict) == 2
+    assert out_dict['paper'].size() == (4, 100, 32)
+    assert out_dict['author'].size() == (4, 100, 32)