We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
unused `group` parameter in `HGTConv` still documented
### 📚 Describe the documentation issue

#7117 replaces `HGTConv` with the implementation done for the faster `FastHGTConv`. In the process of doing so, the `group` parameter has been removed falling back to the default `sum` operation. (@puririshi98: this was intentional, right?). The docs, however, haven't been updated accordingly.

https://github.com/pyg-team/pytorch_geometric/blob/737707c37fc2bd712a2289b683ec14549926ff49/torch_geometric/nn/conv/hgt_conv.py#L40-L43

### Suggest a potential alternative/fix

Remove the unused parameter from the docs.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torch_geometric/nn/conv/hgt_conv.py`
Content:
```
1 import math
2 from typing import Dict, List, Optional, Tuple, Union
3 
4 import torch
5 from torch import Tensor
6 from torch.nn import Parameter
7 
8 from torch_geometric.nn.conv import MessagePassing
9 from torch_geometric.nn.dense import HeteroDictLinear, HeteroLinear
10 from torch_geometric.nn.inits import ones
11 from torch_geometric.nn.parameter_dict import ParameterDict
12 from torch_geometric.typing import Adj, EdgeType, Metadata, NodeType
13 from torch_geometric.utils import softmax
14 from torch_geometric.utils.hetero import construct_bipartite_edge_index
15 
16 
17 class HGTConv(MessagePassing):
18     r"""The Heterogeneous Graph Transformer (HGT) operator from the
19     `"Heterogeneous Graph Transformer" <https://arxiv.org/abs/2003.01332>`_
20     paper.
21 
22     .. note::
23 
24         For an example of using HGT, see `examples/hetero/hgt_dblp.py
25         <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
26         hetero/hgt_dblp.py>`_.
27 
28     Args:
29         in_channels (int or Dict[str, int]): Size of each input sample of every
30             node type, or :obj:`-1` to derive the size from the first input(s)
31             to the forward method.
32         out_channels (int): Size of each output sample.
33         metadata (Tuple[List[str], List[Tuple[str, str, str]]]): The metadata
34             of the heterogeneous graph, *i.e.* its node and edge types given
35             by a list of strings and a list of string triplets, respectively.
36             See :meth:`torch_geometric.data.HeteroData.metadata` for more
37             information.
38         heads (int, optional): Number of multi-head-attentions.
39             (default: :obj:`1`)
40         group (str, optional): The aggregation scheme to use for grouping node
41             embeddings generated by different relations
42             (:obj:`"sum"`, :obj:`"mean"`, :obj:`"min"`, :obj:`"max"`).
43             (default: :obj:`"sum"`)
44         **kwargs (optional): Additional arguments of
45             :class:`torch_geometric.nn.conv.MessagePassing`.
46     """
47     def __init__(
48         self,
49         in_channels: Union[int, Dict[str, int]],
50         out_channels: int,
51         metadata: Metadata,
52         heads: int = 1,
53         **kwargs,
54     ):
55         super().__init__(aggr='add', node_dim=0, **kwargs)
56 
57         if out_channels % heads != 0:
58             raise ValueError(f"'out_channels' (got {out_channels}) must be "
59                              f"divisible by the number of heads (got {heads})")
60 
61         if not isinstance(in_channels, dict):
62             in_channels = {node_type: in_channels for node_type in metadata[0]}
63 
64         self.in_channels = in_channels
65         self.out_channels = out_channels
66         self.heads = heads
67         self.node_types = metadata[0]
68         self.edge_types = metadata[1]
69         self.edge_types_map = {
70             edge_type: i
71             for i, edge_type in enumerate(metadata[1])
72         }
73 
74         self.dst_node_types = set([key[-1] for key in self.edge_types])
75 
76         self.kqv_lin = HeteroDictLinear(self.in_channels,
77                                         self.out_channels * 3)
78 
79         self.out_lin = HeteroDictLinear(self.out_channels, self.out_channels,
80                                         types=self.node_types)
81 
82         dim = out_channels // heads
83         num_types = heads * len(self.edge_types)
84 
85         self.k_rel = HeteroLinear(dim, dim, num_types, bias=False,
86                                   is_sorted=True)
87         self.v_rel = HeteroLinear(dim, dim, num_types, bias=False,
88                                   is_sorted=True)
89 
90         self.skip = ParameterDict({
91             node_type: Parameter(torch.empty(1))
92             for node_type in self.node_types
93         })
94 
95         self.p_rel = ParameterDict()
96         for edge_type in self.edge_types:
97             edge_type = '__'.join(edge_type)
98             self.p_rel[edge_type] = Parameter(torch.empty(1, heads))
99 
100         self.reset_parameters()
101 
102     def reset_parameters(self):
103         super().reset_parameters()
104         self.kqv_lin.reset_parameters()
105         self.out_lin.reset_parameters()
106         self.k_rel.reset_parameters()
107         self.v_rel.reset_parameters()
108         ones(self.skip)
109         ones(self.p_rel)
110 
111     def _cat(self, x_dict: Dict[str, Tensor]) -> Tuple[Tensor, Dict[str, int]]:
112         """Concatenates a dictionary of features."""
113         cumsum = 0
114         outs: List[Tensor] = []
115         offset: Dict[str, int] = {}
116         for key, x in x_dict.items():
117             outs.append(x)
118             offset[key] = cumsum
119             cumsum += x.size(0)
120         return torch.cat(outs, dim=0), offset
121 
122     def _construct_src_node_feat(
123         self, k_dict: Dict[str, Tensor], v_dict: Dict[str, Tensor],
124         edge_index_dict: Dict[EdgeType, Adj]
125     ) -> Tuple[Tensor, Tensor, Dict[EdgeType, int]]:
126         """Constructs the source node representations."""
127         cumsum = 0
128         num_edge_types = len(self.edge_types)
129         H, D = self.heads, self.out_channels // self.heads
130 
131         # Flatten into a single tensor with shape [num_edge_types * heads, D]:
132         ks: List[Tensor] = []
133         vs: List[Tensor] = []
134         type_list: List[Tensor] = []
135         offset: Dict[EdgeType] = {}
136         for edge_type in edge_index_dict.keys():
137             src = edge_type[0]
138             N = k_dict[src].size(0)
139             offset[edge_type] = cumsum
140             cumsum += N
141 
142             # construct type_vec for curr edge_type with shape [H, D]
143             edge_type_offset = self.edge_types_map[edge_type]
144             type_vec = torch.arange(H, dtype=torch.long).view(-1, 1).repeat(
145                 1, N) * num_edge_types + edge_type_offset
146 
147             type_list.append(type_vec)
148             ks.append(k_dict[src])
149             vs.append(v_dict[src])
150 
151         ks = torch.cat(ks, dim=0).transpose(0, 1).reshape(-1, D)
152         vs = torch.cat(vs, dim=0).transpose(0, 1).reshape(-1, D)
153         type_vec = torch.cat(type_list, dim=1).flatten()
154 
155         k = self.k_rel(ks, type_vec).view(H, -1, D).transpose(0, 1)
156         v = self.v_rel(vs, type_vec).view(H, -1, D).transpose(0, 1)
157 
158         return k, v, offset
159 
160     def forward(
161         self,
162         x_dict: Dict[NodeType, Tensor],
163         edge_index_dict: Dict[EdgeType, Adj]  # Support both.
164     ) -> Dict[NodeType, Optional[Tensor]]:
165         r"""Runs the forward pass of the module.
166 
167         Args:
168             x_dict (Dict[str, torch.Tensor]): A dictionary holding input node
169                 features  for each individual node type.
170             edge_index_dict (Dict[Tuple[str, str, str], torch.Tensor]): A
171                 dictionary holding graph connectivity information for each
172                 individual edge type, either as a :class:`torch.Tensor` of
173                 shape :obj:`[2, num_edges]` or a
174                 :class:`torch_sparse.SparseTensor`.
175 
176         :rtype: :obj:`Dict[str, Optional[torch.Tensor]]` - The output node
177             embeddings for each node type.
178             In case a node type does not receive any message, its output will
179             be set to :obj:`None`.
180         """
181         F = self.out_channels
182         H = self.heads
183         D = F // H
184 
185         k_dict, q_dict, v_dict, out_dict = {}, {}, {}, {}
186 
187         # Compute K, Q, V over node types:
188         kqv_dict = self.kqv_lin(x_dict)
189         for key, val in kqv_dict.items():
190             k, q, v = torch.tensor_split(val, 3, dim=1)
191             k_dict[key] = k.view(-1, H, D)
192             q_dict[key] = q.view(-1, H, D)
193             v_dict[key] = v.view(-1, H, D)
194 
195         q, dst_offset = self._cat(q_dict)
196         k, v, src_offset = self._construct_src_node_feat(
197             k_dict, v_dict, edge_index_dict)
198 
199         edge_index, edge_attr = construct_bipartite_edge_index(
200             edge_index_dict, src_offset, dst_offset, edge_attr_dict=self.p_rel)
201 
202         out = self.propagate(edge_index, k=k, q=q, v=v, edge_attr=edge_attr,
203                              size=None)
204 
205         # Reconstruct output node embeddings dict:
206         for node_type, start_offset in dst_offset.items():
207             end_offset = start_offset + q_dict[node_type].size(0)
208             if node_type in self.dst_node_types:
209                 out_dict[node_type] = out[start_offset:end_offset]
210 
211         # Transform output node embeddings:
212         a_dict = self.out_lin({
213             k:
214             torch.nn.functional.gelu(v) if v is not None else v
215             for k, v in out_dict.items()
216         })
217 
218         # Iterate over node types:
219         for node_type, out in out_dict.items():
220             out = a_dict[node_type]
221 
222             if out.size(-1) == x_dict[node_type].size(-1):
223                 alpha = self.skip[node_type].sigmoid()
224                 out = alpha * out + (1 - alpha) * x_dict[node_type]
225             out_dict[node_type] = out
226 
227         return out_dict
228 
229     def message(self, k_j: Tensor, q_i: Tensor, v_j: Tensor, edge_attr: Tensor,
230                 index: Tensor, ptr: Optional[Tensor],
231                 size_i: Optional[int]) -> Tensor:
232         alpha = (q_i * k_j).sum(dim=-1) * edge_attr
233         alpha = alpha / math.sqrt(q_i.size(-1))
234         alpha = softmax(alpha, index, ptr, size_i)
235         out = v_j * alpha.view(-1, self.heads, 1)
236         return out.view(-1, self.out_channels)
237 
238     def __repr__(self) -> str:
239         return (f'{self.__class__.__name__}(-1, {self.out_channels}, '
240                 f'heads={self.heads})')
241
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/torch_geometric/nn/conv/hgt_conv.py b/torch_geometric/nn/conv/hgt_conv.py
--- a/torch_geometric/nn/conv/hgt_conv.py
+++ b/torch_geometric/nn/conv/hgt_conv.py
@@ -37,10 +37,6 @@
             information.
         heads (int, optional): Number of multi-head-attentions.
             (default: :obj:`1`)
-        group (str, optional): The aggregation scheme to use for grouping node
-            embeddings generated by different relations
-            (:obj:`"sum"`, :obj:`"mean"`, :obj:`"min"`, :obj:`"max"`).
-            (default: :obj:`"sum"`)
         **kwargs (optional): Additional arguments of
             :class:`torch_geometric.nn.conv.MessagePassing`.
     """

{"golden_diff": "diff --git a/torch_geometric/nn/conv/hgt_conv.py b/torch_geometric/nn/conv/hgt_conv.py\n--- a/torch_geometric/nn/conv/hgt_conv.py\n+++ b/torch_geometric/nn/conv/hgt_conv.py\n@@ -37,10 +37,6 @@\n             information.\n         heads (int, optional): Number of multi-head-attentions.\n             (default: :obj:`1`)\n-        group (str, optional): The aggregation scheme to use for grouping node\n-            embeddings generated by different relations\n-            (:obj:`\"sum\"`, :obj:`\"mean\"`, :obj:`\"min\"`, :obj:`\"max\"`).\n-            (default: :obj:`\"sum\"`)\n         **kwargs (optional): Additional arguments of\n             :class:`torch_geometric.nn.conv.MessagePassing`.\n     \"\"\"\n", "issue": "unused `group` parameter in `HGTConv` still documented\n### \ud83d\udcda Describe the documentation issue\n\n#7117 replaces `HGTConv` with the implementation done for the faster `FastHGTConv`. In the process of doing so, the `group` parameter has been removed falling back to the default `sum` operation. (@puririshi98: this was intentional, right?). The docs, however, haven't been updated accordingly.\r\n\r\nhttps://github.com/pyg-team/pytorch_geometric/blob/737707c37fc2bd712a2289b683ec14549926ff49/torch_geometric/nn/conv/hgt_conv.py#L40-L43\n\n### Suggest a potential alternative/fix\n\nRemove the unused parameter from the docs.\n", "before_files": [{"content": "import math\nfrom typing import Dict, List, Optional, Tuple, Union\n\nimport torch\nfrom torch import Tensor\nfrom torch.nn import Parameter\n\nfrom torch_geometric.nn.conv import MessagePassing\nfrom torch_geometric.nn.dense import HeteroDictLinear, HeteroLinear\nfrom torch_geometric.nn.inits import ones\nfrom torch_geometric.nn.parameter_dict import ParameterDict\nfrom torch_geometric.typing import Adj, EdgeType, Metadata, NodeType\nfrom torch_geometric.utils import softmax\nfrom torch_geometric.utils.hetero import construct_bipartite_edge_index\n\n\nclass HGTConv(MessagePassing):\n    r\"\"\"The Heterogeneous Graph Transformer (HGT) operator from the\n    `\"Heterogeneous Graph Transformer\" <https://arxiv.org/abs/2003.01332>`_\n    paper.\n\n    .. note::\n\n        For an example of using HGT, see `examples/hetero/hgt_dblp.py\n        <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/\n        hetero/hgt_dblp.py>`_.\n\n    Args:\n        in_channels (int or Dict[str, int]): Size of each input sample of every\n            node type, or :obj:`-1` to derive the size from the first input(s)\n            to the forward method.\n        out_channels (int): Size of each output sample.\n        metadata (Tuple[List[str], List[Tuple[str, str, str]]]): The metadata\n            of the heterogeneous graph, *i.e.* its node and edge types given\n            by a list of strings and a list of string triplets, respectively.\n            See :meth:`torch_geometric.data.HeteroData.metadata` for more\n            information.\n        heads (int, optional): Number of multi-head-attentions.\n            (default: :obj:`1`)\n        group (str, optional): The aggregation scheme to use for grouping node\n            embeddings generated by different relations\n            (:obj:`\"sum\"`, :obj:`\"mean\"`, :obj:`\"min\"`, :obj:`\"max\"`).\n            (default: :obj:`\"sum\"`)\n        **kwargs (optional): Additional arguments of\n            :class:`torch_geometric.nn.conv.MessagePassing`.\n    \"\"\"\n    def __init__(\n        self,\n        in_channels: Union[int, Dict[str, int]],\n        out_channels: int,\n        metadata: Metadata,\n        heads: int = 1,\n        **kwargs,\n    ):\n        super().__init__(aggr='add', node_dim=0, **kwargs)\n\n        if out_channels % heads != 0:\n            raise ValueError(f\"'out_channels' (got {out_channels}) must be \"\n                             f\"divisible by the number of heads (got {heads})\")\n\n        if not isinstance(in_channels, dict):\n            in_channels = {node_type: in_channels for node_type in metadata[0]}\n\n        self.in_channels = in_channels\n        self.out_channels = out_channels\n        self.heads = heads\n        self.node_types = metadata[0]\n        self.edge_types = metadata[1]\n        self.edge_types_map = {\n            edge_type: i\n            for i, edge_type in enumerate(metadata[1])\n        }\n\n        self.dst_node_types = set([key[-1] for key in self.edge_types])\n\n        self.kqv_lin = HeteroDictLinear(self.in_channels,\n                                        self.out_channels * 3)\n\n        self.out_lin = HeteroDictLinear(self.out_channels, self.out_channels,\n                                        types=self.node_types)\n\n        dim = out_channels // heads\n        num_types = heads * len(self.edge_types)\n\n        self.k_rel = HeteroLinear(dim, dim, num_types, bias=False,\n                                  is_sorted=True)\n        self.v_rel = HeteroLinear(dim, dim, num_types, bias=False,\n                                  is_sorted=True)\n\n        self.skip = ParameterDict({\n            node_type: Parameter(torch.empty(1))\n            for node_type in self.node_types\n        })\n\n        self.p_rel = ParameterDict()\n        for edge_type in self.edge_types:\n            edge_type = '__'.join(edge_type)\n            self.p_rel[edge_type] = Parameter(torch.empty(1, heads))\n\n        self.reset_parameters()\n\n    def reset_parameters(self):\n        super().reset_parameters()\n        self.kqv_lin.reset_parameters()\n        self.out_lin.reset_parameters()\n        self.k_rel.reset_parameters()\n        self.v_rel.reset_parameters()\n        ones(self.skip)\n        ones(self.p_rel)\n\n    def _cat(self, x_dict: Dict[str, Tensor]) -> Tuple[Tensor, Dict[str, int]]:\n        \"\"\"Concatenates a dictionary of features.\"\"\"\n        cumsum = 0\n        outs: List[Tensor] = []\n        offset: Dict[str, int] = {}\n        for key, x in x_dict.items():\n            outs.append(x)\n            offset[key] = cumsum\n            cumsum += x.size(0)\n        return torch.cat(outs, dim=0), offset\n\n    def _construct_src_node_feat(\n        self, k_dict: Dict[str, Tensor], v_dict: Dict[str, Tensor],\n        edge_index_dict: Dict[EdgeType, Adj]\n    ) -> Tuple[Tensor, Tensor, Dict[EdgeType, int]]:\n        \"\"\"Constructs the source node representations.\"\"\"\n        cumsum = 0\n        num_edge_types = len(self.edge_types)\n        H, D = self.heads, self.out_channels // self.heads\n\n        # Flatten into a single tensor with shape [num_edge_types * heads, D]:\n        ks: List[Tensor] = []\n        vs: List[Tensor] = []\n        type_list: List[Tensor] = []\n        offset: Dict[EdgeType] = {}\n        for edge_type in edge_index_dict.keys():\n            src = edge_type[0]\n            N = k_dict[src].size(0)\n            offset[edge_type] = cumsum\n            cumsum += N\n\n            # construct type_vec for curr edge_type with shape [H, D]\n            edge_type_offset = self.edge_types_map[edge_type]\n            type_vec = torch.arange(H, dtype=torch.long).view(-1, 1).repeat(\n                1, N) * num_edge_types + edge_type_offset\n\n            type_list.append(type_vec)\n            ks.append(k_dict[src])\n            vs.append(v_dict[src])\n\n        ks = torch.cat(ks, dim=0).transpose(0, 1).reshape(-1, D)\n        vs = torch.cat(vs, dim=0).transpose(0, 1).reshape(-1, D)\n        type_vec = torch.cat(type_list, dim=1).flatten()\n\n        k = self.k_rel(ks, type_vec).view(H, -1, D).transpose(0, 1)\n        v = self.v_rel(vs, type_vec).view(H, -1, D).transpose(0, 1)\n\n        return k, v, offset\n\n    def forward(\n        self,\n        x_dict: Dict[NodeType, Tensor],\n        edge_index_dict: Dict[EdgeType, Adj]  # Support both.\n    ) -> Dict[NodeType, Optional[Tensor]]:\n        r\"\"\"Runs the forward pass of the module.\n\n        Args:\n            x_dict (Dict[str, torch.Tensor]): A dictionary holding input node\n                features  for each individual node type.\n            edge_index_dict (Dict[Tuple[str, str, str], torch.Tensor]): A\n                dictionary holding graph connectivity information for each\n                individual edge type, either as a :class:`torch.Tensor` of\n                shape :obj:`[2, num_edges]` or a\n                :class:`torch_sparse.SparseTensor`.\n\n        :rtype: :obj:`Dict[str, Optional[torch.Tensor]]` - The output node\n            embeddings for each node type.\n            In case a node type does not receive any message, its output will\n            be set to :obj:`None`.\n        \"\"\"\n        F = self.out_channels\n        H = self.heads\n        D = F // H\n\n        k_dict, q_dict, v_dict, out_dict = {}, {}, {}, {}\n\n        # Compute K, Q, V over node types:\n        kqv_dict = self.kqv_lin(x_dict)\n        for key, val in kqv_dict.items():\n            k, q, v = torch.tensor_split(val, 3, dim=1)\n            k_dict[key] = k.view(-1, H, D)\n            q_dict[key] = q.view(-1, H, D)\n            v_dict[key] = v.view(-1, H, D)\n\n        q, dst_offset = self._cat(q_dict)\n        k, v, src_offset = self._construct_src_node_feat(\n            k_dict, v_dict, edge_index_dict)\n\n        edge_index, edge_attr = construct_bipartite_edge_index(\n            edge_index_dict, src_offset, dst_offset, edge_attr_dict=self.p_rel)\n\n        out = self.propagate(edge_index, k=k, q=q, v=v, edge_attr=edge_attr,\n                             size=None)\n\n        # Reconstruct output node embeddings dict:\n        for node_type, start_offset in dst_offset.items():\n            end_offset = start_offset + q_dict[node_type].size(0)\n            if node_type in self.dst_node_types:\n                out_dict[node_type] = out[start_offset:end_offset]\n\n        # Transform output node embeddings:\n        a_dict = self.out_lin({\n            k:\n            torch.nn.functional.gelu(v) if v is not None else v\n            for k, v in out_dict.items()\n        })\n\n        # Iterate over node types:\n        for node_type, out in out_dict.items():\n            out = a_dict[node_type]\n\n            if out.size(-1) == x_dict[node_type].size(-1):\n                alpha = self.skip[node_type].sigmoid()\n                out = alpha * out + (1 - alpha) * x_dict[node_type]\n            out_dict[node_type] = out\n\n        return out_dict\n\n    def message(self, k_j: Tensor, q_i: Tensor, v_j: Tensor, edge_attr: Tensor,\n                index: Tensor, ptr: Optional[Tensor],\n                size_i: Optional[int]) -> Tensor:\n        alpha = (q_i * k_j).sum(dim=-1) * edge_attr\n        alpha = alpha / math.sqrt(q_i.size(-1))\n        alpha = softmax(alpha, index, ptr, size_i)\n        out = v_j * alpha.view(-1, self.heads, 1)\n        return out.view(-1, self.out_channels)\n\n    def __repr__(self) -> str:\n        return (f'{self.__class__.__name__}(-1, {self.out_channels}, '\n                f'heads={self.heads})')\n", "path": "torch_geometric/nn/conv/hgt_conv.py"}], "after_files": [{"content": "import math\nfrom typing import Dict, List, Optional, Tuple, Union\n\nimport torch\nfrom torch import Tensor\nfrom torch.nn import Parameter\n\nfrom torch_geometric.nn.conv import MessagePassing\nfrom torch_geometric.nn.dense import HeteroDictLinear, HeteroLinear\nfrom torch_geometric.nn.inits import ones\nfrom torch_geometric.nn.parameter_dict import ParameterDict\nfrom torch_geometric.typing import Adj, EdgeType, Metadata, NodeType\nfrom torch_geometric.utils import softmax\nfrom torch_geometric.utils.hetero import construct_bipartite_edge_index\n\n\nclass HGTConv(MessagePassing):\n    r\"\"\"The Heterogeneous Graph Transformer (HGT) operator from the\n    `\"Heterogeneous Graph Transformer\" <https://arxiv.org/abs/2003.01332>`_\n    paper.\n\n    .. note::\n\n        For an example of using HGT, see `examples/hetero/hgt_dblp.py\n        <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/\n        hetero/hgt_dblp.py>`_.\n\n    Args:\n        in_channels (int or Dict[str, int]): Size of each input sample of every\n            node type, or :obj:`-1` to derive the size from the first input(s)\n            to the forward method.\n        out_channels (int): Size of each output sample.\n        metadata (Tuple[List[str], List[Tuple[str, str, str]]]): The metadata\n            of the heterogeneous graph, *i.e.* its node and edge types given\n            by a list of strings and a list of string triplets, respectively.\n            See :meth:`torch_geometric.data.HeteroData.metadata` for more\n            information.\n        heads (int, optional): Number of multi-head-attentions.\n            (default: :obj:`1`)\n        **kwargs (optional): Additional arguments of\n            :class:`torch_geometric.nn.conv.MessagePassing`.\n    \"\"\"\n    def __init__(\n        self,\n        in_channels: Union[int, Dict[str, int]],\n        out_channels: int,\n        metadata: Metadata,\n        heads: int = 1,\n        **kwargs,\n    ):\n        super().__init__(aggr='add', node_dim=0, **kwargs)\n\n        if out_channels % heads != 0:\n            raise ValueError(f\"'out_channels' (got {out_channels}) must be \"\n                             f\"divisible by the number of heads (got {heads})\")\n\n        if not isinstance(in_channels, dict):\n            in_channels = {node_type: in_channels for node_type in metadata[0]}\n\n        self.in_channels = in_channels\n        self.out_channels = out_channels\n        self.heads = heads\n        self.node_types = metadata[0]\n        self.edge_types = metadata[1]\n        self.edge_types_map = {\n            edge_type: i\n            for i, edge_type in enumerate(metadata[1])\n        }\n\n        self.dst_node_types = set([key[-1] for key in self.edge_types])\n\n        self.kqv_lin = HeteroDictLinear(self.in_channels,\n                                        self.out_channels * 3)\n\n        self.out_lin = HeteroDictLinear(self.out_channels, self.out_channels,\n                                        types=self.node_types)\n\n        dim = out_channels // heads\n        num_types = heads * len(self.edge_types)\n\n        self.k_rel = HeteroLinear(dim, dim, num_types, bias=False,\n                                  is_sorted=True)\n        self.v_rel = HeteroLinear(dim, dim, num_types, bias=False,\n                                  is_sorted=True)\n\n        self.skip = ParameterDict({\n            node_type: Parameter(torch.empty(1))\n            for node_type in self.node_types\n        })\n\n        self.p_rel = ParameterDict()\n        for edge_type in self.edge_types:\n            edge_type = '__'.join(edge_type)\n            self.p_rel[edge_type] = Parameter(torch.empty(1, heads))\n\n        self.reset_parameters()\n\n    def reset_parameters(self):\n        super().reset_parameters()\n        self.kqv_lin.reset_parameters()\n        self.out_lin.reset_parameters()\n        self.k_rel.reset_parameters()\n        self.v_rel.reset_parameters()\n        ones(self.skip)\n        ones(self.p_rel)\n\n    def _cat(self, x_dict: Dict[str, Tensor]) -> Tuple[Tensor, Dict[str, int]]:\n        \"\"\"Concatenates a dictionary of features.\"\"\"\n        cumsum = 0\n        outs: List[Tensor] = []\n        offset: Dict[str, int] = {}\n        for key, x in x_dict.items():\n            outs.append(x)\n            offset[key] = cumsum\n            cumsum += x.size(0)\n        return torch.cat(outs, dim=0), offset\n\n    def _construct_src_node_feat(\n        self, k_dict: Dict[str, Tensor], v_dict: Dict[str, Tensor],\n        edge_index_dict: Dict[EdgeType, Adj]\n    ) -> Tuple[Tensor, Tensor, Dict[EdgeType, int]]:\n        \"\"\"Constructs the source node representations.\"\"\"\n        cumsum = 0\n        num_edge_types = len(self.edge_types)\n        H, D = self.heads, self.out_channels // self.heads\n\n        # Flatten into a single tensor with shape [num_edge_types * heads, D]:\n        ks: List[Tensor] = []\n        vs: List[Tensor] = []\n        type_list: List[Tensor] = []\n        offset: Dict[EdgeType] = {}\n        for edge_type in edge_index_dict.keys():\n            src = edge_type[0]\n            N = k_dict[src].size(0)\n            offset[edge_type] = cumsum\n            cumsum += N\n\n            # construct type_vec for curr edge_type with shape [H, D]\n            edge_type_offset = self.edge_types_map[edge_type]\n            type_vec = torch.arange(H, dtype=torch.long).view(-1, 1).repeat(\n                1, N) * num_edge_types + edge_type_offset\n\n            type_list.append(type_vec)\n            ks.append(k_dict[src])\n            vs.append(v_dict[src])\n\n        ks = torch.cat(ks, dim=0).transpose(0, 1).reshape(-1, D)\n        vs = torch.cat(vs, dim=0).transpose(0, 1).reshape(-1, D)\n        type_vec = torch.cat(type_list, dim=1).flatten()\n\n        k = self.k_rel(ks, type_vec).view(H, -1, D).transpose(0, 1)\n        v = self.v_rel(vs, type_vec).view(H, -1, D).transpose(0, 1)\n\n        return k, v, offset\n\n    def forward(\n        self,\n        x_dict: Dict[NodeType, Tensor],\n        edge_index_dict: Dict[EdgeType, Adj]  # Support both.\n    ) -> Dict[NodeType, Optional[Tensor]]:\n        r\"\"\"Runs the forward pass of the module.\n\n        Args:\n            x_dict (Dict[str, torch.Tensor]): A dictionary holding input node\n                features  for each individual node type.\n            edge_index_dict (Dict[Tuple[str, str, str], torch.Tensor]): A\n                dictionary holding graph connectivity information for each\n                individual edge type, either as a :class:`torch.Tensor` of\n                shape :obj:`[2, num_edges]` or a\n                :class:`torch_sparse.SparseTensor`.\n\n        :rtype: :obj:`Dict[str, Optional[torch.Tensor]]` - The output node\n            embeddings for each node type.\n            In case a node type does not receive any message, its output will\n            be set to :obj:`None`.\n        \"\"\"\n        F = self.out_channels\n        H = self.heads\n        D = F // H\n\n        k_dict, q_dict, v_dict, out_dict = {}, {}, {}, {}\n\n        # Compute K, Q, V over node types:\n        kqv_dict = self.kqv_lin(x_dict)\n        for key, val in kqv_dict.items():\n            k, q, v = torch.tensor_split(val, 3, dim=1)\n            k_dict[key] = k.view(-1, H, D)\n            q_dict[key] = q.view(-1, H, D)\n            v_dict[key] = v.view(-1, H, D)\n\n        q, dst_offset = self._cat(q_dict)\n        k, v, src_offset = self._construct_src_node_feat(\n            k_dict, v_dict, edge_index_dict)\n\n        edge_index, edge_attr = construct_bipartite_edge_index(\n            edge_index_dict, src_offset, dst_offset, edge_attr_dict=self.p_rel)\n\n        out = self.propagate(edge_index, k=k, q=q, v=v, edge_attr=edge_attr,\n                             size=None)\n\n        # Reconstruct output node embeddings dict:\n        for node_type, start_offset in dst_offset.items():\n            end_offset = start_offset + q_dict[node_type].size(0)\n            if node_type in self.dst_node_types:\n                out_dict[node_type] = out[start_offset:end_offset]\n\n        # Transform output node embeddings:\n        a_dict = self.out_lin({\n            k:\n            torch.nn.functional.gelu(v) if v is not None else v\n            for k, v in out_dict.items()\n        })\n\n        # Iterate over node types:\n        for node_type, out in out_dict.items():\n            out = a_dict[node_type]\n\n            if out.size(-1) == x_dict[node_type].size(-1):\n                alpha = self.skip[node_type].sigmoid()\n                out = alpha * out + (1 - alpha) * x_dict[node_type]\n            out_dict[node_type] = out\n\n        return out_dict\n\n    def message(self, k_j: Tensor, q_i: Tensor, v_j: Tensor, edge_attr: Tensor,\n                index: Tensor, ptr: Optional[Tensor],\n                size_i: Optional[int]) -> Tensor:\n        alpha = (q_i * k_j).sum(dim=-1) * edge_attr\n        alpha = alpha / math.sqrt(q_i.size(-1))\n        alpha = softmax(alpha, index, ptr, size_i)\n        out = v_j * alpha.view(-1, self.heads, 1)\n        return out.view(-1, self.out_channels)\n\n    def __repr__(self) -> str:\n        return (f'{self.__class__.__name__}(-1, {self.out_channels}, '\n                f'heads={self.heads})')\n", "path": "torch_geometric/nn/conv/hgt_conv.py"}]}