Repository: pyg-team/pytorch_geometric
Pull request #6178 (2022-12-08 00:04:08)

`pyg-lib` integration in `HGTConv`
ready🎉
```
data = FakeHeteroDataset(num_node_types=16, num_edge_types=128).data.to('cuda')
master branch iterative algorithm fwd pass time: 0.0014522361755371093
average fwd pass time(pyg-lib w/ torch<1.14): 0.0004720401763916016
average fwd pass time(pyg-lib w/ torch>=1.14): 0.00046291828155517576
```

3x speedup. As part of my blog I will do a more extensive benchmark to see how runtime for for-loop, segment_matmul, and group_matmul backends of HGT compare agaisnt eachother

smaller heterograph:
```
data = FakeHeteroDataset(num_node_types=8, num_edge_types=16).data.to('cuda')
average fwd pass time: 0.00020465850830078124
average fwd pass time(pyg-lib w/ torch<1.14): 0.00013236522674560546
average fwd pass time(pyg-lib w/ torch>=1.14): 0.00013113975524902343
```
1.6x speedup

scaling:
<img width="641" alt="image" src="https://user-images.githubusercontent.com/20074092/215614068-401066d2-c4db-4685-ad63-f7f772bd4c36.png">


For interactive graph: https://github.com/puririshi98/rgcn_pyg_lib_forward_bench/blob/main/3d_plot.py

puririshi98: still need to test and benchmark it but done w/ an initial draft
puririshi98: @rusty1s [this PR ](https://github.com/pyg-team/pyg-lib/pull/183)to pyg-lib is needed for torch >= 1.14
I checked locally that once I build pyg-lib w/ that PR included, the grouped matmul and segment matmul pathways both work except i get the wierd numerical mismatch between the results of using coo edge indices vs csr indices:


```
        for node_type in out_dict1.keys():
>           assert torch.allclose(out_dict1[node_type], out_dict2[node_type],
                                  atol=1e-6)
E           assert False
E            +  where False = <built-in method allclose of type object at 0x7f5817c2abc0>(tensor([[-0.0835, -0.6736,  0.5235,  0.2286, -0.0260,  0.3432, -0.3394, -0.0610,\n          0.6782,  0.3317,  0.3618,  0.3018, -0.1857,  0.1624, -0.4920,  0.0622],\n        [-0.0600,  0.0160,  0.2255,  0.0513,  0.0783, -0.0502,  0.2167,  0.1157,\n          0.0927,  0.0175, -0.0928,  0.0178, -0.0471,  0.2207, -0.1846, -0.3963],\n        [-0.1174,  0.0345,  0.3031,  0.0869,  0.3046,  0.0708, -0.4384,  0.0139,\n          0.0960,  0.5766,  0.0997, -0.3248, -0.3266,  0.0738, -0.2258,  0.3176],\n        [ 0.2621,  0.1540,  0.2881,  0.1190, -0.1480, -0.0665,  0.3536, -0.0015,\n         -0.1051, -0.0060, -0.2733,  0.3685, -0.3494,  0.3862, -0.1167,  0.1122]],\n       grad_fn=<AddBackward0>), tensor([[-0.1288, -0.4715,  0.6094,  0.0883,  0.0351,  0.2979, -0.4755,  0.0921,\n          0.7698,  0.1106,  0.4882,  0.2792, -0.3037,  0.3606, -0.6694,  0.0487],\n        [-0.0075,  0.2320,  0.3542, -0.1559,  0.1909, -0.0801,  0.0948,  0.3497,\n          0.1253, -0.2684,  0.0116,  0.2078, -0.1117,  0.3709, -0.3056, -0.4902],\n        [-0.0718,  0.1115,  0.4493,  0.0243,  0.3455,  0.1374, -0.3731,  0.0014,\n          0.1058,  0.4818,  0.1863, -0.1104, -0.2884,  0.1684, -0.2276,  0.3486],\n        [ 0.2621,  0.1540,  0.2881,  0.1190, -0.1480, -0.0665,  0.3536, -0.0015,\n         -0.1051, -0.0060, -0.2733,  0.3685, -0.3494,  0.3862, -0.1167,  0.1122]],\n       grad_fn=<AddBackward0>), atol=1e-06)
E            +    where <built-in method allclose of type object at 0x7f5817c2abc0> = torch.allclose
```
puririshi98: when i try to run the following benchmark:
```
from torch_geometric.nn.conv import HGTConv
from torch_geometric.datasets import FakeHeteroDataset
import time
data = FakeHeteroDataset(num_node_types=16, num_edge_types=128).data.to('cuda')
net = HGTConv(-1, 32, data.metadata()).to('cuda')
x_dict = data.collect('x')
edge_index_dict = data.collect('edge_index')
for i in range(60):
  if i > 9:
    since = time.time()
  net(x_dict, edge_index_dict)
print("average fwd pass time:", (time.time()-since)/50.0)
```

i get the following error: 
```
Traceback (most recent call last):
  File "x.py", line 11, in <module>
    net(x_dict, edge_index_dict)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1480, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch_geometric-2.3.0-py3.8.egg/torch_geometric/nn/conv/hgt_conv.py", line 395, in forward
    out = self.propagate(e_idx, k=k_out, q=q, v=v_out, rel=p, size=None)
  File "/usr/local/lib/python3.8/dist-packages/torch_geometric-2.3.0-py3.8.egg/torch_geometric/nn/conv/message_passing.py", line 432, in propagate
    coll_dict = self.__collect__(self.__user_args__, edge_index,
  File "/usr/local/lib/python3.8/dist-packages/torch_geometric-2.3.0-py3.8.egg/torch_geometric/nn/conv/message_passing.py", line 304, in __collect__
    data = self.__lift__(data, edge_index, dim)
  File "/usr/local/lib/python3.8/dist-packages/torch_geometric-2.3.0-py3.8.egg/torch_geometric/nn/conv/message_passing.py", line 242, in __lift__
    raise IndexError(
IndexError: Encountered an index error. Please ensure that all indices in 'edge_index' point to valid indices in the interval [0, 130169] (got interval [169034, 259441])
```

to me this implies theres a problem w/ me going from edge_index_dict to the single edge index. will investigate further. lmk if anything is obvious to you
puririshi98: copied your code for making edgeindexdict into and edge index tensor and i still get numerical mismatches btwn adj_t and edge_index use
puririshi98: debugging prints:
```
        # propagate
        print("\n"*2)
        print("k[0,0] =", k_out[0,0])
        print("q[0,0] =", q[0,0])
        print("v[0,0] =", v_out[0,0])
        print("p[0,0] =", p)
        out = self.propagate(e_idx, k=k_out, q=q, v=v_out, rel=p, size=None)
        print("propogate_out[0,0] =", out[0,0])
```
output:
```
coo format:
k[0,0] = tensor([-0.4085, -0.4570, -0.2870,  0.5745,  0.3404,  0.4127,  0.6706,  0.2871],
       grad_fn=<SelectBackward0>)
q[0,0] = tensor([ 0.0614, -0.6190, -1.3580, -0.6144,  0.4729,  0.7143,  0.2142,  0.6712],
       grad_fn=<SelectBackward0>)
v[0,0] = tensor([ 0.1358,  0.2227, -0.3013, -0.0545, -1.0439,  0.2239,  1.1845, -0.0182],
       grad_fn=<SelectBackward0>)
p[0,0] = tensor([2., 2.], grad_fn=<ViewBackward0>)
propogate_out[0,0] = tensor(0.3956, grad_fn=<SelectBackward0>)


csr format:
k[0,0] = tensor([-0.4085, -0.4570, -0.2870,  0.5745,  0.3404,  0.4127,  0.6706,  0.2871],
       grad_fn=<SelectBackward0>)
q[0,0] = tensor([ 0.0614, -0.6190, -1.3580, -0.6144,  0.4729,  0.7143,  0.2142,  0.6712],
       grad_fn=<SelectBackward0>)
v[0,0] = tensor([ 0.1358,  0.2227, -0.3013, -0.0545, -1.0439,  0.2239,  1.1845, -0.0182],
       grad_fn=<SelectBackward0>)
p[0,0] = tensor([2., 2.], grad_fn=<ViewBackward0>)
propogate_out[0,0] = tensor(-0.2069, grad_fn=<SelectBackward0>)
```

notice that the inputs to propagate are identical but the outputs are not. Somehow [converting the edge_indices from coo to csr ](https://github.com/pyg-team/pytorch_geometric/blob/hgtconv_pyglib/torch_geometric/nn/conv/hgt_conv.py#L367-L385)changes the results of propogate for identical inputs.
puririshi98: I have isolated the reason for the test failures where HGT outputs don't match between coo and csr edge_index representation. Notice in the debug output below that k, q, v, and p are identical between both runs, but the edge indices don't match which explains why the output of propagate does not match. There seems to be an issue in converting csr to coo and then back to csr. @rusty1s could you please take a look and see if I am doing anything noticeably wrong?


[debugging prints](https://github.com/pyg-team/pytorch_geometric/blob/hgtconv_pyglib/torch_geometric/nn/conv/hgt_conv.py#L362-L403):

output:
```
coo format:
e_idx[('author', 'writes', 'paper')] = tensor([[1, 2, 0, 1, 1, 1, 0, 3, 1, 3, 2, 1, 0, 3, 1, 2, 1, 3, 0, 1],
        [1, 1, 3, 4, 2, 5, 3, 5, 2, 2, 0, 1, 1, 0, 5, 1, 4, 2, 2, 4]])
e_idx[('paper', 'written_by', 'author')] = tensor([[1, 1, 3, 4, 2, 5, 3, 5, 2, 2, 0, 1, 1, 0, 5, 1, 4, 2, 2, 4],
        [1, 2, 0, 1, 1, 1, 0, 3, 1, 3, 2, 1, 0, 3, 1, 2, 1, 3, 0, 1]])
k[0,0] = tensor([0.6495, 0.3871, 0.5111, 0.0973, 0.1445, 0.5161, 0.4522, 0.1621],
       grad_fn=<SelectBackward0>)
k.shape = torch.Size([10, 2, 8])
q[0,0] = tensor([ 0.1918, -0.7619,  0.7618,  0.1111,  0.4187,  1.5218,  0.3339, -0.9197],
       grad_fn=<SelectBackward0>)
q.shape = torch.Size([10, 2, 8])
v[0,0] = tensor([-0.1232, -0.0434, -0.1867, -0.7753,  0.3347, -0.1346,  0.1867, -0.2690],
       grad_fn=<SelectBackward0>)
v.shape = torch.Size([10, 2, 8])
p[0,0] = tensor([2., 2.], grad_fn=<ViewBackward0>)
e_idx = tensor([[1, 2, 0, 1, 1, 1, 0, 3, 1, 3, 2, 1, 0, 3, 1, 2, 1, 3, 0, 1, 5, 5, 7, 8,
         6, 9, 7, 9, 6, 6, 4, 5, 5, 4, 9, 5, 8, 6, 6, 8],
        [5, 5, 7, 8, 6, 9, 7, 9, 6, 6, 4, 5, 5, 4, 9, 5, 8, 6, 6, 8, 1, 2, 0, 1,
         1, 1, 0, 3, 1, 3, 2, 1, 0, 3, 1, 2, 1, 3, 0, 1]])
propagate_out[0,0] = tensor(-0.0425, grad_fn=<SelectBackward0>)


csr format (converting to coo and then back to csr under the hood to combine adj_t_dict into adj_t):
e_idx[('author', 'writes', 'paper')] = tensor([[0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 4, 4, 4, 5, 5, 5],
        [2, 3, 0, 1, 1, 2, 2, 0, 1, 1, 3, 3, 0, 0, 1, 1, 1, 1, 1, 3]])
e_idx[('paper', 'written_by', 'author')] = tensor([[0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 3, 3, 3, 3],
        [1, 2, 3, 3, 1, 1, 2, 2, 4, 4, 4, 5, 5, 0, 1, 1, 0, 2, 2, 5]])
k[0,0] = tensor([0.6495, 0.3871, 0.5111, 0.0973, 0.1445, 0.5161, 0.4522, 0.1621],
       grad_fn=<SelectBackward0>)
k.shape = torch.Size([10, 2, 8])
q[0,0] = tensor([ 0.1918, -0.7619,  0.7618,  0.1111,  0.4187,  1.5218,  0.3339, -0.9197],
       grad_fn=<SelectBackward0>)
q.shape = torch.Size([10, 2, 8])
v[0,0] = tensor([-0.1232, -0.0434, -0.1867, -0.7753,  0.3347, -0.1346,  0.1867, -0.2690],
       grad_fn=<SelectBackward0>)
v.shape = torch.Size([10, 2, 8])
p[0,0] = tensor([2., 2.], grad_fn=<ViewBackward0>)
e_idx = tensor([[0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5,
         5, 5, 5, 5, 5, 5, 5, 5, 5, 6, 6, 6, 7, 7, 7, 7],
        [6, 7, 4, 5, 5, 6, 6, 4, 5, 5, 7, 7, 4, 4, 1, 2, 3, 3, 5, 5, 5, 1, 1, 2,
         2, 4, 4, 4, 5, 5, 5, 5, 7, 0, 1, 1, 0, 2, 2, 5]])
propagate_out[0,0] = tensor(-0.3010, grad_fn=<SelectBackward0>)
```


puririshi98: figured it out lolz, the src n dst were backwards but since the conversion sorts it was hard to realize that at first glance of the edge indices
codecov[bot]: # [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6178?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#6178](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6178?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (938d44f) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/6d6162dbb0f74a98929d0462500777de78679660?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (6d6162d) will **decrease** coverage by `0.06%`.
> The diff coverage is `79.86%`.

```diff
@@            Coverage Diff             @@
##           master    #6178      +/-   ##
==========================================
- Coverage   85.75%   85.69%   -0.06%     
==========================================
  Files         406      406              
  Lines       22725    22851     +126     
==========================================
+ Hits        19487    19583      +96     
- Misses       3238     3268      +30     
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6178?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/nn/conv/hgt\_conv.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6178?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2NvbnYvaGd0X2NvbnYucHk=) | `85.46% <79.86%> (-11.57%)` | :arrow_down: |

: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)

puririshi98: finished a major cleanup thanks to the review @rusty1s. In doing so I realized I should be parralelizing a_lin as well which further improved the biggest speedboost to 3.6x, see [blog](https://docs.google.com/document/d/1uRB4IpjL5DlS3YG19eNnX9zSvGzgK7JdNX10aXjdDy8/edit?usp=sharing) for updated benchmark plot

puririshi98: thanks for the reviews, lmk if anything else needed to merge @wsad1 @rusty1s 
puririshi98: unfortunately, at the moment, even if i set all params to zero for the new hgt vs old hgt, the outputs dont match
test: https://github.com/puririshi98/rgcn_pyg_lib_forward_bench/blob/main/hgt_correctness.py

puririshi98: basic correctness passes: https://github.com/puririshi98/rgcn_pyg_lib_forward_bench/blob/main/hgt_correctness.py
@rusty1s merge when u can :)

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3b4b66c1c..ba229bc75 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 `HeteroDictLinear` and an optimized `HGTConv` module ([#6178](https://github.com/pyg-team/pytorch_geometric/pull/6178))
 - Added the `DenseGATConv` module ([#6928](https://github.com/pyg-team/pytorch_geometric/pull/6928))
 - Added `trim_to_layer` utility function for more efficient `NeighborLoader` use-cases ([#6661](https://github.com/pyg-team/pytorch_geometric/pull/6661))
 - Added the `DistMult` KGE model ([#6958](https://github.com/pyg-team/pytorch_geometric/pull/6958))
diff --git a/torch_geometric/data/hetero_data.py b/torch_geometric/data/hetero_data.py
index 92f24f373..7d324efcc 100644
--- a/torch_geometric/data/hetero_data.py
+++ b/torch_geometric/data/hetero_data.py
@@ -1006,35 +1006,50 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
 # Helper functions ############################################################
 
 
+def get_node_slices(num_nodes: Dict[str, int]) -> Dict[str, Tuple[int, int]]:
+    r"""Returns the boundaries of each node type in a graph."""
+    node_slices: Dict[NodeType, Tuple[int, int]] = {}
+    cumsum = 0
+    for node_type, N in num_nodes.items():
+        node_slices[node_type] = (cumsum, cumsum + N)
+        cumsum += N
+    return node_slices
+
+
+def offset_edge_index(
+    node_slices: Dict[NodeType, Tuple[int, int]],
+    edge_type: EdgeType,
+    edge_index: Tensor,
+) -> Tensor:
+    r"""Increases the edge indices by the offsets of source and destination
+    node types."""
+    src, _, dst = edge_type
+    offset = [[node_slices[src][0]], [node_slices[dst][0]]]
+    offset = torch.tensor(offset, device=edge_index.device)
+    return edge_index + offset
+
+
 def to_homogeneous_edge_index(
     data: HeteroData,
 ) -> Tuple[Optional[Tensor], Dict[NodeType, Any], Dict[EdgeType, Any]]:
+    r"""Converts a heterogeneous graph into a homogeneous typed graph."""
     # Record slice information per node type:
-    cumsum = 0
-    node_slices: Dict[NodeType, Tuple[int, int]] = {}
-    for node_type, store in data._node_store_dict.items():
-        num_nodes = store.num_nodes
-        node_slices[node_type] = (cumsum, cumsum + num_nodes)
-        cumsum += num_nodes
+    node_slices = get_node_slices(data.num_nodes_dict)
 
     # Record edge indices and slice information per edge type:
     cumsum = 0
     edge_indices: List[Tensor] = []
     edge_slices: Dict[EdgeType, Tuple[int, int]] = {}
-    for edge_type, store in data._edge_store_dict.items():
-        src, _, dst = edge_type
-        offset = [[node_slices[src][0]], [node_slices[dst][0]]]
-        offset = torch.tensor(offset, device=store.edge_index.device)
-        edge_indices.append(store.edge_index + offset)
-
-        num_edges = store.num_edges
-        edge_slices[edge_type] = (cumsum, cumsum + num_edges)
-        cumsum += num_edges
+    for edge_type, edge_index in data.edge_index_dict.items():
+        edge_index = offset_edge_index(node_slices, edge_type, edge_index)
+        edge_indices.append(edge_index)
+        edge_slices[edge_type] = (cumsum, cumsum + edge_index.size(1))
+        cumsum += edge_index.size(1)
 
-    edge_index = None
+    edge_index: Optional[Tensor] = None
     if len(edge_indices) == 1:  # Memory-efficient `torch.cat`:
         edge_index = edge_indices[0]
-    elif len(edge_indices) > 0:
+    elif len(edge_indices) > 1:
         edge_index = torch.cat(edge_indices, dim=-1)
 
     return edge_index, node_slices, edge_slices
diff --git a/torch_geometric/nn/conv/__init__.py b/torch_geometric/nn/conv/__init__.py
index 9410e03f0..a0998c3bf 100644
--- a/torch_geometric/nn/conv/__init__.py
+++ b/torch_geometric/nn/conv/__init__.py
@@ -51,6 +51,7 @@ from .eg_conv import EGConv
 from .pdn_conv import PDNConv
 from .general_conv import GeneralConv
 from .hgt_conv import HGTConv
+from .fast_hgt_conv import FastHGTConv
 from .heat_conv import HEATConv
 from .hetero_conv import HeteroConv
 from .han_conv import HANConv
@@ -118,6 +119,7 @@ __all__ = [
     'PDNConv',
     'GeneralConv',
     'HGTConv',
+    'FastHGTConv',
     'HEATConv',
     'HeteroConv',
     'HANConv',
diff --git a/torch_geometric/nn/conv/fast_hgt_conv.py b/torch_geometric/nn/conv/fast_hgt_conv.py
new file mode 100644
index 000000000..73223b44e
--- /dev/null
+++ b/torch_geometric/nn/conv/fast_hgt_conv.py
@@ -0,0 +1,256 @@
+import math
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn import Parameter
+
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.nn.dense import HeteroDictLinear, HeteroLinear
+from torch_geometric.nn.inits import ones
+from torch_geometric.nn.parameter_dict import ParameterDict
+from torch_geometric.typing import (
+    Adj,
+    EdgeType,
+    Metadata,
+    NodeType,
+    SparseTensor,
+)
+from torch_geometric.utils import softmax
+
+
+class FastHGTConv(MessagePassing):
+    r"""See :class:`HGTConv`."""
+    def __init__(
+        self,
+        in_channels: Union[int, Dict[str, int]],
+        out_channels: int,
+        metadata: Metadata,
+        heads: int = 1,
+        **kwargs,
+    ):
+        super().__init__(aggr='add', node_dim=0, **kwargs)
+
+        if out_channels % heads != 0:
+            raise ValueError(f"'out_channels' (got {out_channels}) must be "
+                             f"divisible by the number of heads (got {heads})")
+
+        if not isinstance(in_channels, dict):
+            in_channels = {node_type: in_channels for node_type in metadata[0]}
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.heads = heads
+
+        self.node_types = metadata[0]
+        self.edge_types = metadata[1]
+        self.src_types = [edge_type[0] for edge_type in self.edge_types]
+
+        self.k_lin = HeteroDictLinear(self.in_channels, self.out_channels)
+        self.q_lin = HeteroDictLinear(self.in_channels, self.out_channels)
+        self.v_lin = HeteroDictLinear(self.in_channels, self.out_channels)
+
+        self.out_lin = HeteroDictLinear(self.out_channels, self.out_channels,
+                                        types=self.node_types)
+
+        dim = out_channels // heads
+        num_types = heads * len(self.edge_types)
+
+        self.k_rel = HeteroLinear(dim, dim, num_types, is_sorted=True,
+                                  bias=False)
+        self.v_rel = HeteroLinear(dim, dim, num_types, is_sorted=True,
+                                  bias=False)
+
+        self.skip = ParameterDict({
+            node_type: Parameter(torch.Tensor(1))
+            for node_type in self.node_types
+        })
+
+        self.p_rel = ParameterDict()
+        for edge_type in self.edge_types:
+            edge_type = '__'.join(edge_type)
+            self.p_rel[edge_type] = Parameter(torch.Tensor(1, heads))
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        super().reset_parameters()
+        self.k_lin.reset_parameters()
+        self.q_lin.reset_parameters()
+        self.v_lin.reset_parameters()
+        self.out_lin.reset_parameters()
+        self.k_rel.reset_parameters()
+        self.v_rel.reset_parameters()
+        ones(self.skip)
+        ones(self.p_rel)
+
+    def _cat(self, x_dict: Dict[str, Tensor]) -> Tuple[Tensor, Dict[str, int]]:
+        """Concatenates a dictionary of features."""
+        cumsum = 0
+        outs: List[Tensor] = []
+        offset: Dict[str, int] = {}
+        for key, x in x_dict.items():
+            outs.append(x)
+            offset[key] = cumsum
+            cumsum += x.size(0)
+        return torch.cat(outs, dim=0), offset
+
+    def _construct_src_node_feat(
+        self,
+        k_dict: Dict[str, Tensor],
+        v_dict: Dict[str, Tensor],
+    ) -> Tuple[Tensor, Tensor, Dict[EdgeType, int]]:
+        """Constructs the source node representations."""
+        count = 0
+        cumsum = 0
+        H, D = self.heads, self.out_channels // self.heads
+
+        # Flatten into a single tensor with shape [num_edge_types * heads, D]:
+        ks: List[Tensor] = []
+        vs: List[Tensor] = []
+        type_list: List[int] = []
+        offset: Dict[EdgeType] = {}
+        for edge_type in self.edge_types:
+            src, _, _ = edge_type
+
+            ks.append(k_dict[src].view(-1, D))
+            vs.append(v_dict[src].view(-1, D))
+
+            N = k_dict[src].size(0)
+            for _ in range(H):
+                type_list.append(torch.full((N, ), count, dtype=torch.long))
+                count += 1
+            offset[edge_type] = cumsum
+            cumsum += N
+
+        type_vec = torch.cat(type_list, dim=0)
+        k = self.k_rel(torch.cat(ks, dim=0), type_vec).view(-1, H, D)
+        v = self.k_rel(torch.cat(vs, dim=0), type_vec).view(-1, H, D)
+
+        return k, v, offset
+
+    def _construct_edge_index(
+        self,
+        edge_index_dict: Dict[EdgeType, Adj],
+        src_offset: Dict[EdgeType, int],
+        dst_offset: [NodeType, int],
+    ) -> Tuple[Adj, Tensor]:
+        """Constructs a tensor of edge indices by concatenating edge indices
+        for each edge type. The edge indices are increased by the offset of the
+        source and destination nodes."""
+        edge_indices: List[Tensor] = []
+        ps: List[Tensor] = []
+
+        for edge_type in self.edge_types:
+            _, _, dst_type = edge_type
+
+            edge_index = edge_index_dict[edge_type]
+
+            # (TODO) Add support for SparseTensor w/o converting.
+            is_sparse = isinstance(edge_index, SparseTensor)
+            if is_sparse:  # Convert to COO
+                dst, src, _ = edge_index.coo()
+                edge_index = torch.stack([src, dst], dim=0)
+            else:
+                edge_index = edge_index.clone()
+
+            p = self.p_rel['__'.join(edge_type)].expand(edge_index.size(1), -1)
+            ps.append(p)
+
+            # Add offset to edge indices:
+            edge_index[0] += src_offset[edge_type]
+            edge_index[1] += dst_offset[dst_type]
+            edge_indices.append(edge_index)
+
+        # Concatenate all edges and edge tensors:
+        p = torch.cat(ps, dim=0)
+        edge_index = torch.cat(edge_indices, dim=1)
+
+        if is_sparse:
+            edge_index = SparseTensor(row=edge_index[1], col=edge_index[0],
+                                      value=p)
+
+        return edge_index, p
+
+    def forward(
+        self,
+        x_dict: Dict[NodeType, Tensor],
+        edge_index_dict: Dict[EdgeType, Adj]  # Support both.
+    ) -> Dict[NodeType, Optional[Tensor]]:
+        r"""Runs the forward pass of the module.
+
+        Args:
+            x_dict (Dict[str, torch.Tensor]): A dictionary holding input node
+                features  for each individual node type.
+            edge_index_dict (Dict[Tuple[str, str, str], torch.Tensor]): A
+                dictionary holding graph connectivity information for each
+                individual edge type, either as a :class:`torch.Tensor` of
+                shape :obj:`[2, num_edges]` or a
+                :class:`torch_sparse.SparseTensor`.
+
+        :rtype: :obj:`Dict[str, Optional[torch.Tensor]]` - The output node
+            embeddings for each node type.
+            In case a node type does not receive any message, its output will
+            be set to :obj:`None`.
+        """
+        H, D = self.heads, self.out_channels // self.heads
+
+        k_dict, q_dict, v_dict = {}, {}, {}
+        out_dict = defaultdict(list)
+
+        # Compute K, Q, V over node types:
+        k_dict = self.k_lin(x_dict)
+        k_dict = {k: v.view(-1, H, D) for k, v in k_dict.items()}
+
+        q_dict = self.q_lin(x_dict)
+        q_dict = {k: v.view(-1, H, D) for k, v in q_dict.items()}
+
+        v_dict = self.v_lin(x_dict)
+        v_dict = {k: v.view(-1, H, D) for k, v in v_dict.items()}
+
+        q, dst_offset = self._cat(q_dict)
+        k, v, src_offset = self._construct_src_node_feat(k_dict, v_dict)
+
+        edge_index, edge_attr = self._construct_edge_index(
+            edge_index_dict, src_offset, dst_offset)
+
+        out = self.propagate(edge_index, k=k, q=q, v=v, edge_attr=edge_attr,
+                             size=None)
+
+        # Reconstruct output node embeddings dict:
+        for node_type, start_offset in dst_offset.items():
+            end_offset = start_offset + q_dict[node_type].size(0)
+            out_dict[node_type] = out[start_offset:end_offset]
+
+        # Transform output node embeddings:
+        a_dict = self.out_lin({k: F.gelu(v) for k, v in out_dict.items()})
+
+        # Iterate over node types:
+        for node_type, out in out_dict.items():
+            if out is None:
+                out_dict[node_type] = None
+                continue
+            else:
+                out = a_dict[node_type]
+
+            if out.size(-1) == x_dict[node_type].size(-1):
+                alpha = self.skip[node_type].sigmoid()
+                out = alpha * out + (1 - alpha) * x_dict[node_type]
+            out_dict[node_type] = out
+
+        return out_dict
+
+    def message(self, k_j: Tensor, q_i: Tensor, v_j: Tensor, edge_attr: Tensor,
+                index: Tensor, ptr: Optional[Tensor],
+                size_i: Optional[int]) -> Tensor:
+        alpha = (q_i * k_j).sum(dim=-1) * edge_attr
+        alpha = alpha / math.sqrt(q_i.size(-1))
+        alpha = softmax(alpha, index, ptr, size_i)
+        out = v_j * alpha.view(-1, self.heads, 1)
+        return out.view(-1, self.out_channels)
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}(-1, {self.out_channels}, '
+                f'heads={self.heads})')
diff --git a/torch_geometric/nn/conv/hgt_conv.py b/torch_geometric/nn/conv/hgt_conv.py
index b56f355d3..1f5bc8692 100644
--- a/torch_geometric/nn/conv/hgt_conv.py
+++ b/torch_geometric/nn/conv/hgt_conv.py
@@ -42,6 +42,11 @@ class HGTConv(MessagePassing):
         <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
         hetero/hgt_dblp.py>`_.
 
+    .. note::
+
+        For a faster alternative, use :class:`FastHGTConv` which does not
+        iterate over individual node and edge types.
+
     Args:
         in_channels (int or Dict[str, int]): Size of each input sample of every
             node type, or :obj:`-1` to derive the size from the first input(s)
diff --git a/torch_geometric/nn/dense/__init__.py b/torch_geometric/nn/dense/__init__.py
index c7ea2a982..5533294c3 100644
--- a/torch_geometric/nn/dense/__init__.py
+++ b/torch_geometric/nn/dense/__init__.py
@@ -1,4 +1,4 @@
-from .linear import Linear, HeteroLinear
+from .linear import Linear, HeteroLinear, HeteroDictLinear
 from .dense_gat_conv import DenseGATConv
 from .dense_sage_conv import DenseSAGEConv
 from .dense_gcn_conv import DenseGCNConv
@@ -11,6 +11,7 @@ from .dmon_pool import DMoNPooling
 __all__ = [
     'Linear',
     'HeteroLinear',
+    'HeteroDictLinear',
     'DenseGCNConv',
     'DenseGINConv',
     'DenseGraphConv',
diff --git a/torch_geometric/nn/dense/linear.py b/torch_geometric/nn/dense/linear.py
index fef1081e1..c8971e1c9 100644
--- a/torch_geometric/nn/dense/linear.py
+++ b/torch_geometric/nn/dense/linear.py
@@ -1,6 +1,6 @@
 import copy
 import math
-from typing import Any, Optional
+from typing import Any, Dict, Optional, Union
 
 import torch
 import torch.nn.functional as F
@@ -218,8 +218,13 @@ class HeteroLinear(torch.nn.Module):
 
         if torch_geometric.typing.WITH_PYG_LIB:
             self.lins = None
-            self.weight = torch.nn.Parameter(
-                torch.Tensor(num_types, in_channels, out_channels))
+            if self.in_channels == -1:
+                self.weight = nn.parameter.UninitializedParameter()
+                self._hook = self.register_forward_pre_hook(
+                    self.initialize_parameters)
+            else:
+                self.weight = torch.nn.Parameter(
+                    torch.Tensor(num_types, in_channels, out_channels))
             if kwargs.get('bias', True):
                 self.bias = Parameter(torch.Tensor(num_types, out_channels))
             else:
@@ -277,7 +282,124 @@ class HeteroLinear(torch.nn.Module):
                 out[mask] = lin(x[mask])
         return out
 
+    @torch.no_grad()
+    def initialize_parameters(self, module, input):
+        if is_uninitialized_parameter(self.weight):
+            self.in_channels = input[0].size(-1)
+            self.weight.materialize(
+                (self.num_types, self.in_channels, self.out_channels))
+            self.reset_parameters()
+        self._hook.remove()
+        delattr(self, '_hook')
+
     def __repr__(self) -> str:
         return (f'{self.__class__.__name__}({self.in_channels}, '
                 f'{self.out_channels}, num_types={self.num_types}, '
                 f'bias={self.kwargs.get("bias", True)})')
+
+
+class HeteroDictLinear(torch.nn.Module):
+    r"""Applies separate linear tranformations to the incoming data dictionary
+
+    .. math::
+        \mathbf{x}^{\prime}_{\kappa} = \mathbf{x}_{\kappa}
+        \mathbf{W}^{\top}_{\kappa} + \mathbf{b}_{\kappa}
+
+    for key :math:`\kappa`.
+    It supports lazy initialization and customizable weight and bias
+    initialization.
+
+    Args:
+        in_channels (int or Dict[Any, int]): Size of each input sample. If
+            passed an integer, :obj:`types` will be a mandatory argument.
+            initialized lazily in case it is given as :obj:`-1`.
+        out_channels (int): Size of each output sample.
+        types (List[Any], optional): The keys of the input dictionary.
+            (default: :obj:`None`)
+        **kwargs (optional): Additional arguments of
+            :class:`torch_geometric.nn.Linear`.
+    """
+    def __init__(
+        self,
+        in_channels: Union[int, Dict[Any, int]],
+        out_channels: int,
+        types: Optional[Any] = None,
+        **kwargs,
+    ):
+        super().__init__()
+
+        if isinstance(in_channels, dict):
+            self.types = list(in_channels.keys())
+
+            if any([i == -1 for i in in_channels.values()]):
+                self._hook = self.register_forward_pre_hook(
+                    self.initialize_parameters)
+
+            if types is not None and set(self.types) != set(types):
+                raise ValueError("The provided 'types' do not match with the "
+                                 "keys in the 'in_channels' dictionary")
+
+        else:
+            if types is None:
+                raise ValueError("Please provide a list of 'types' if passing "
+                                 "'in_channels' as an integer")
+
+            if in_channels == -1:
+                self._hook = self.register_forward_pre_hook(
+                    self.initialize_parameters)
+
+            self.types = types
+            in_channels = {node_type: in_channels for node_type in types}
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kwargs = kwargs
+
+        self.lins = torch.nn.ModuleDict({
+            key: Linear(channels, self.out_channels, **kwargs)
+            for key, channels in self.in_channels.items()
+        })
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        for lin in self.lins.values():
+            lin.reset_parameters()
+
+    def forward(
+        self,
+        x_dict: Dict[Any, Tensor],
+    ) -> Dict[Any, Tensor]:
+        r"""
+        Args:
+            x_dict (Dict[Any, torch.Tensor]): A dictionary holding input
+                features for each individual type.
+        """
+        if torch_geometric.typing.WITH_GMM:
+            xs = [x_dict[key] for key in x_dict.keys()]
+            weights = [self.lins[key].weight.t() for key in x_dict.keys()]
+            if self.kwargs.get('bias', True):
+                biases = [self.lins[key].bias for key in x_dict.keys()]
+            else:
+                biases = None
+
+            outs = pyg_lib.ops.grouped_matmul(xs, weights, biases)
+            return {key: out for key, out in zip(x_dict.keys(), outs)}
+
+        return {key: self.lins[key](x) for key, x in x_dict.items()}
+
+    @torch.no_grad()
+    def initialize_parameters(self, module, input):
+        for key, x in input[0].items():
+            lin = self.lins[key]
+            if is_uninitialized_parameter(lin.weight):
+                self.lins[key].initialize_parameters(None, x)
+        self.reset_parameters()
+        self._hook.remove()
+        self.in_channels = {key: x.size(-1) for key, x in input[0].items()}
+        delattr(self, '_hook')
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.out_channels}, bias={self.kwargs.get("bias", True)})')
diff --git a/torch_geometric/nn/to_hetero_module.py b/torch_geometric/nn/to_hetero_module.py
index 9a7912818..82e3298a9 100644
--- a/torch_geometric/nn/to_hetero_module.py
+++ b/torch_geometric/nn/to_hetero_module.py
@@ -35,7 +35,6 @@ class ToHeteroLinear(torch.nn.Module):
         else:
             raise ValueError(f"Expected 'Linear' module (got '{type(module)}'")
 
-        # TODO: Need to handle `in_channels=-1` case.
         # TODO We currently assume that `x` is sorted according to `type`.
         self.hetero_module = HeteroLinear(
             in_channels,
diff --git a/torch_geometric/typing.py b/torch_geometric/typing.py
index 1c16a22d2..bb9c74e77 100644
--- a/torch_geometric/typing.py
+++ b/torch_geometric/typing.py
@@ -10,14 +10,16 @@ WITH_PT2 = int(torch.__version__.split('.')[0]) >= 2
 try:
     import pyg_lib  # noqa
     WITH_PYG_LIB = True
-    WITH_SAMPLED_OP = WITH_PYG_LIB and hasattr(pyg_lib.ops, 'sampled_add')
-    WITH_INDEX_SORT = WITH_PYG_LIB and hasattr(pyg_lib.ops, 'index_sort')
+    WITH_GMM = WITH_PT2 and hasattr(pyg_lib.ops, 'grouped_matmul')
+    WITH_SAMPLED_OP = hasattr(pyg_lib.ops, 'sampled_add')
+    WITH_INDEX_SORT = hasattr(pyg_lib.ops, 'index_sort')
 except (ImportError, OSError) as e:
     if isinstance(e, OSError):
         warnings.warn(f"An issue occurred while importing 'pyg-lib'. "
                       f"Disabling its usage. Stacktrace: {e}")
     pyg_lib = object
     WITH_PYG_LIB = False
+    WITH_GMM = False
     WITH_SAMPLED_OP = False
     WITH_INDEX_SORT = False
 
diff --git a/torch_geometric/utils/spmm.py b/torch_geometric/utils/spmm.py
index e80f5510d..fc83a34b5 100644
--- a/torch_geometric/utils/spmm.py
+++ b/torch_geometric/utils/spmm.py
@@ -24,9 +24,9 @@ def spmm(src: Adj, other: Tensor, reduce: str = "sum") -> Tensor:
     """Matrix product of sparse matrix with dense matrix.
 
     Args:
-        src (Tensor or torch_sparse.SparseTensor]): The input sparse matrix,
-            either a :class:`torch_sparse.SparseTensor` or a
-            :class:`torch.sparse.Tensor`.
+        src (Tensor or torch_sparse.SparseTensor): The input sparse matrix,
+            either a :pyg:`PyG` :class:`torch_sparse.SparseTensor` or a
+            :pytorch:`PyTorch` :class:`torch.sparse.Tensor`.
         other (Tensor): The input dense matrix.
         reduce (str, optional): The reduce operation to use
             (:obj:`"sum"`, :obj:`"mean"`, :obj:`"min"`, :obj:`"max"`).


Test patch:
diff --git a/test/nn/conv/test_fast_hgt_conv.py b/test/nn/conv/test_fast_hgt_conv.py
new file mode 100644
index 000000000..55d36bf49
--- /dev/null
+++ b/test/nn/conv/test_fast_hgt_conv.py
@@ -0,0 +1,26 @@
+import torch
+
+from torch_geometric.nn import FastHGTConv
+from torch_geometric.testing import get_random_edge_index
+
+
+def test_fast_hgt_conv():
+    x_dict = {
+        'author': torch.randn(4, 16),
+        'paper': torch.randn(6, 16),
+    }
+    edge_index = get_random_edge_index(4, 6, num_edges=20)
+
+    edge_index_dict = {
+        ('author', 'writes', 'paper'): edge_index,
+        ('paper', 'written_by', 'author'): edge_index.flip([0]),
+    }
+
+    metadata = (list(x_dict.keys()), list(edge_index_dict.keys()))
+
+    conv = FastHGTConv(16, 16, metadata, heads=2)
+    assert str(conv) == 'FastHGTConv(-1, 16, heads=2)'
+    out_dict1 = conv(x_dict, edge_index_dict)
+    assert len(out_dict1) == 2
+    assert out_dict1['author'].size() == (4, 16)
+    assert out_dict1['paper'].size() == (6, 16)
diff --git a/test/nn/dense/test_linear.py b/test/nn/dense/test_linear.py
index e7f0ebfc6..d5e9185f1 100644
--- a/test/nn/dense/test_linear.py
+++ b/test/nn/dense/test_linear.py
@@ -5,7 +5,7 @@ import torch
 from torch.nn import Linear as PTLinear
 from torch.nn.parameter import UninitializedParameter
 
-from torch_geometric.nn import HeteroLinear, Linear
+from torch_geometric.nn import HeteroDictLinear, HeteroLinear, Linear
 from torch_geometric.testing import is_full_test, withPackage
 
 weight_inits = ['glorot', 'kaiming_uniform', None]
@@ -116,6 +116,53 @@ def test_hetero_linear():
         assert torch.allclose(jit(x, type_vec), out)
 
 
+def test_lazy_hetero_linear():
+    x = torch.randn(3, 16)
+    type_vec = torch.tensor([0, 1, 2])
+
+    lin = HeteroLinear(-1, 32, num_types=3)
+    assert str(lin) == 'HeteroLinear(-1, 32, num_types=3, bias=True)'
+
+    out = lin(x, type_vec)
+    assert out.size() == (3, 32)
+    assert str(lin) == 'HeteroLinear(16, 32, num_types=3, bias=True)'
+
+
+def test_hetero_dict_linear():
+    x_dict = {'v': torch.randn(3, 16), 'w': torch.randn(2, 8)}
+
+    lin = HeteroDictLinear({'v': 16, 'w': 8}, 32)
+    assert str(lin) == "HeteroDictLinear({'v': 16, 'w': 8}, 32, bias=True)"
+
+    out_dict = lin(x_dict)
+    assert len(out_dict) == 2
+    assert out_dict['v'].size() == (3, 32)
+    assert out_dict['w'].size() == (2, 32)
+
+    x_dict = {'v': torch.randn(3, 16), 'w': torch.randn(2, 16)}
+
+    lin = HeteroDictLinear(16, 32, types=['v', 'w'])
+    assert str(lin) == "HeteroDictLinear({'v': 16, 'w': 16}, 32, bias=True)"
+
+    out_dict = lin(x_dict)
+    assert len(out_dict) == 2
+    assert out_dict['v'].size() == (3, 32)
+    assert out_dict['w'].size() == (2, 32)
+
+
+def test_lazy_hetero_dict_linear():
+    x_dict = {'v': torch.randn(3, 16), 'w': torch.randn(2, 8)}
+
+    lin = HeteroDictLinear(-1, 32, types=['v', 'w'])
+    assert str(lin) == "HeteroDictLinear({'v': -1, 'w': -1}, 32, bias=True)"
+
+    out_dict = lin(x_dict)
+    assert len(out_dict) == 2
+    assert out_dict['v'].size() == (3, 32)
+    assert out_dict['w'].size() == (2, 32)
+    assert str(lin) == "HeteroDictLinear({'v': 16, 'w': 8}, 32, bias=True)"
+
+
 @withPackage('pyg_lib')
 @pytest.mark.parametrize('type_vec', [
     torch.tensor([0, 0, 1, 1, 2, 2]),