We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
HGT implementation bug?
https://github.com/dmlc/dgl/blob/master/examples/pytorch/hgt/model.py#L81

The implemented hgt set value of nodes multiple times when calculating by different edge attention, which will cause the `nodes["v"]` of earlier etype be overwritten by later one. 

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/pytorch/hgt/model.py`
Content:
```
1 import dgl
2 import math
3 import torch
4 import torch.nn as nn
5 import torch.nn.functional as F
6 import dgl.function as fn
7 from dgl.nn.functional import edge_softmax
8 
9 class HGTLayer(nn.Module):
10     def __init__(self,
11                  in_dim,
12                  out_dim,
13                  node_dict,
14                  edge_dict,
15                  n_heads,
16                  dropout = 0.2,
17                  use_norm = False):
18         super(HGTLayer, self).__init__()
19 
20         self.in_dim        = in_dim
21         self.out_dim       = out_dim
22         self.node_dict     = node_dict
23         self.edge_dict     = edge_dict
24         self.num_types     = len(node_dict)
25         self.num_relations = len(edge_dict)
26         self.total_rel     = self.num_types * self.num_relations * self.num_types
27         self.n_heads       = n_heads
28         self.d_k           = out_dim // n_heads
29         self.sqrt_dk       = math.sqrt(self.d_k)
30         self.att           = None
31 
32         self.k_linears   = nn.ModuleList()
33         self.q_linears   = nn.ModuleList()
34         self.v_linears   = nn.ModuleList()
35         self.a_linears   = nn.ModuleList()
36         self.norms       = nn.ModuleList()
37         self.use_norm    = use_norm
38 
39         for t in range(self.num_types):
40             self.k_linears.append(nn.Linear(in_dim,   out_dim))
41             self.q_linears.append(nn.Linear(in_dim,   out_dim))
42             self.v_linears.append(nn.Linear(in_dim,   out_dim))
43             self.a_linears.append(nn.Linear(out_dim,  out_dim))
44             if use_norm:
45                 self.norms.append(nn.LayerNorm(out_dim))
46 
47         self.relation_pri   = nn.Parameter(torch.ones(self.num_relations, self.n_heads))
48         self.relation_att   = nn.Parameter(torch.Tensor(self.num_relations, n_heads, self.d_k, self.d_k))
49         self.relation_msg   = nn.Parameter(torch.Tensor(self.num_relations, n_heads, self.d_k, self.d_k))
50         self.skip           = nn.Parameter(torch.ones(self.num_types))
51         self.drop           = nn.Dropout(dropout)
52 
53         nn.init.xavier_uniform_(self.relation_att)
54         nn.init.xavier_uniform_(self.relation_msg)
55 
56     def forward(self, G, h):
57         with G.local_scope():
58             node_dict, edge_dict = self.node_dict, self.edge_dict
59             for srctype, etype, dsttype in G.canonical_etypes:
60                 sub_graph = G[srctype, etype, dsttype]
61 
62                 k_linear = self.k_linears[node_dict[srctype]]
63                 v_linear = self.v_linears[node_dict[srctype]]
64                 q_linear = self.q_linears[node_dict[dsttype]]
65 
66                 k = k_linear(h[srctype]).view(-1, self.n_heads, self.d_k)
67                 v = v_linear(h[srctype]).view(-1, self.n_heads, self.d_k)
68                 q = q_linear(h[dsttype]).view(-1, self.n_heads, self.d_k)
69 
70                 e_id = self.edge_dict[etype]
71 
72                 relation_att = self.relation_att[e_id]
73                 relation_pri = self.relation_pri[e_id]
74                 relation_msg = self.relation_msg[e_id]
75 
76                 k = torch.einsum("bij,ijk->bik", k, relation_att)
77                 v = torch.einsum("bij,ijk->bik", v, relation_msg)
78 
79                 sub_graph.srcdata['k'] = k
80                 sub_graph.dstdata['q'] = q
81                 sub_graph.srcdata['v'] = v
82 
83                 sub_graph.apply_edges(fn.v_dot_u('q', 'k', 't'))
84                 attn_score = sub_graph.edata.pop('t').sum(-1) * relation_pri / self.sqrt_dk
85                 attn_score = edge_softmax(sub_graph, attn_score, norm_by='dst')
86 
87                 sub_graph.edata['t'] = attn_score.unsqueeze(-1)
88 
89             G.multi_update_all({etype : (fn.u_mul_e('v', 't', 'm'), fn.sum('m', 't')) \
90                                 for etype in edge_dict}, cross_reducer = 'mean')
91 
92             new_h = {}
93             for ntype in G.ntypes:
94                 '''
95                     Step 3: Target-specific Aggregation
96                     x = norm( W[node_type] * gelu( Agg(x) ) + x )
97                 '''
98                 n_id = node_dict[ntype]
99                 alpha = torch.sigmoid(self.skip[n_id])
100                 t = G.nodes[ntype].data['t'].view(-1, self.out_dim)
101                 trans_out = self.drop(self.a_linears[n_id](t))
102                 trans_out = trans_out * alpha + h[ntype] * (1-alpha)
103                 if self.use_norm:
104                     new_h[ntype] = self.norms[n_id](trans_out)
105                 else:
106                     new_h[ntype] = trans_out
107             return new_h
108 
109 class HGT(nn.Module):
110     def __init__(self, G, node_dict, edge_dict, n_inp, n_hid, n_out, n_layers, n_heads, use_norm = True):
111         super(HGT, self).__init__()
112         self.node_dict = node_dict
113         self.edge_dict = edge_dict
114         self.gcs = nn.ModuleList()
115         self.n_inp = n_inp
116         self.n_hid = n_hid
117         self.n_out = n_out
118         self.n_layers = n_layers
119         self.adapt_ws  = nn.ModuleList()
120         for t in range(len(node_dict)):
121             self.adapt_ws.append(nn.Linear(n_inp,   n_hid))
122         for _ in range(n_layers):
123             self.gcs.append(HGTLayer(n_hid, n_hid, node_dict, edge_dict, n_heads, use_norm = use_norm))
124         self.out = nn.Linear(n_hid, n_out)
125 
126     def forward(self, G, out_key):
127         h = {}
128         for ntype in G.ntypes:
129             n_id = self.node_dict[ntype]
130             h[ntype] = F.gelu(self.adapt_ws[n_id](G.nodes[ntype].data['inp']))
131         for i in range(self.n_layers):
132             h = self.gcs[i](G, h)
133         return self.out(h[out_key])
134 
135 class HeteroRGCNLayer(nn.Module):
136     def __init__(self, in_size, out_size, etypes):
137         super(HeteroRGCNLayer, self).__init__()
138         # W_r for each relation
139         self.weight = nn.ModuleDict({
140                 name : nn.Linear(in_size, out_size) for name in etypes
141             })
142 
143     def forward(self, G, feat_dict):
144         # The input is a dictionary of node features for each type
145         funcs = {}
146         for srctype, etype, dsttype in G.canonical_etypes:
147             # Compute W_r * h
148             Wh = self.weight[etype](feat_dict[srctype])
149             # Save it in graph for message passing
150             G.nodes[srctype].data['Wh_%s' % etype] = Wh
151             # Specify per-relation message passing functions: (message_func, reduce_func).
152             # Note that the results are saved to the same destination feature 'h', which
153             # hints the type wise reducer for aggregation.
154             funcs[etype] = (fn.copy_u('Wh_%s' % etype, 'm'), fn.mean('m', 'h'))
155         # Trigger message passing of multiple types.
156         # The first argument is the message passing functions for each relation.
157         # The second one is the type wise reducer, could be "sum", "max",
158         # "min", "mean", "stack"
159         G.multi_update_all(funcs, 'sum')
160         # return the updated node feature dictionary
161         return {ntype : G.nodes[ntype].data['h'] for ntype in G.ntypes}
162 
163 
164 class HeteroRGCN(nn.Module):
165     def __init__(self, G, in_size, hidden_size, out_size):
166         super(HeteroRGCN, self).__init__()
167         # create layers
168         self.layer1 = HeteroRGCNLayer(in_size, hidden_size, G.etypes)
169         self.layer2 = HeteroRGCNLayer(hidden_size, out_size, G.etypes)
170 
171     def forward(self, G, out_key):
172         input_dict = {ntype : G.nodes[ntype].data['inp'] for ntype in G.ntypes}
173         h_dict = self.layer1(G, input_dict)
174         h_dict = {k : F.leaky_relu(h) for k, h in h_dict.items()}
175         h_dict = self.layer2(G, h_dict)
176         # get paper logits
177         return h_dict[out_key]
178
```
--- 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/examples/pytorch/hgt/model.py b/examples/pytorch/hgt/model.py
--- a/examples/pytorch/hgt/model.py
+++ b/examples/pytorch/hgt/model.py
@@ -78,7 +78,7 @@
 
                 sub_graph.srcdata['k'] = k
                 sub_graph.dstdata['q'] = q
-                sub_graph.srcdata['v'] = v
+                sub_graph.srcdata['v_%d' % e_id] = v
 
                 sub_graph.apply_edges(fn.v_dot_u('q', 'k', 't'))
                 attn_score = sub_graph.edata.pop('t').sum(-1) * relation_pri / self.sqrt_dk
@@ -86,8 +86,8 @@
 
                 sub_graph.edata['t'] = attn_score.unsqueeze(-1)
 
-            G.multi_update_all({etype : (fn.u_mul_e('v', 't', 'm'), fn.sum('m', 't')) \
-                                for etype in edge_dict}, cross_reducer = 'mean')
+            G.multi_update_all({etype : (fn.u_mul_e('v_%d' % e_id, 't', 'm'), fn.sum('m', 't')) \
+                                for etype, e_id in edge_dict.items()}, cross_reducer = 'mean')
 
             new_h = {}
             for ntype in G.ntypes:

{"golden_diff": "diff --git a/examples/pytorch/hgt/model.py b/examples/pytorch/hgt/model.py\n--- a/examples/pytorch/hgt/model.py\n+++ b/examples/pytorch/hgt/model.py\n@@ -78,7 +78,7 @@\n \n                 sub_graph.srcdata['k'] = k\n                 sub_graph.dstdata['q'] = q\n-                sub_graph.srcdata['v'] = v\n+                sub_graph.srcdata['v_%d' % e_id] = v\n \n                 sub_graph.apply_edges(fn.v_dot_u('q', 'k', 't'))\n                 attn_score = sub_graph.edata.pop('t').sum(-1) * relation_pri / self.sqrt_dk\n@@ -86,8 +86,8 @@\n \n                 sub_graph.edata['t'] = attn_score.unsqueeze(-1)\n \n-            G.multi_update_all({etype : (fn.u_mul_e('v', 't', 'm'), fn.sum('m', 't')) \\\n-                                for etype in edge_dict}, cross_reducer = 'mean')\n+            G.multi_update_all({etype : (fn.u_mul_e('v_%d' % e_id, 't', 'm'), fn.sum('m', 't')) \\\n+                                for etype, e_id in edge_dict.items()}, cross_reducer = 'mean')\n \n             new_h = {}\n             for ntype in G.ntypes:\n", "issue": "HGT implementation bug?\nhttps://github.com/dmlc/dgl/blob/master/examples/pytorch/hgt/model.py#L81\r\n\r\nThe implemented hgt set value of nodes multiple times when calculating by different edge attention, which will cause the `nodes[\"v\"]` of earlier etype be overwritten by later one. \n", "before_files": [{"content": "import dgl\nimport math\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport dgl.function as fn\nfrom dgl.nn.functional import edge_softmax\n\nclass HGTLayer(nn.Module):\n    def __init__(self,\n                 in_dim,\n                 out_dim,\n                 node_dict,\n                 edge_dict,\n                 n_heads,\n                 dropout = 0.2,\n                 use_norm = False):\n        super(HGTLayer, self).__init__()\n\n        self.in_dim        = in_dim\n        self.out_dim       = out_dim\n        self.node_dict     = node_dict\n        self.edge_dict     = edge_dict\n        self.num_types     = len(node_dict)\n        self.num_relations = len(edge_dict)\n        self.total_rel     = self.num_types * self.num_relations * self.num_types\n        self.n_heads       = n_heads\n        self.d_k           = out_dim // n_heads\n        self.sqrt_dk       = math.sqrt(self.d_k)\n        self.att           = None\n\n        self.k_linears   = nn.ModuleList()\n        self.q_linears   = nn.ModuleList()\n        self.v_linears   = nn.ModuleList()\n        self.a_linears   = nn.ModuleList()\n        self.norms       = nn.ModuleList()\n        self.use_norm    = use_norm\n\n        for t in range(self.num_types):\n            self.k_linears.append(nn.Linear(in_dim,   out_dim))\n            self.q_linears.append(nn.Linear(in_dim,   out_dim))\n            self.v_linears.append(nn.Linear(in_dim,   out_dim))\n            self.a_linears.append(nn.Linear(out_dim,  out_dim))\n            if use_norm:\n                self.norms.append(nn.LayerNorm(out_dim))\n\n        self.relation_pri   = nn.Parameter(torch.ones(self.num_relations, self.n_heads))\n        self.relation_att   = nn.Parameter(torch.Tensor(self.num_relations, n_heads, self.d_k, self.d_k))\n        self.relation_msg   = nn.Parameter(torch.Tensor(self.num_relations, n_heads, self.d_k, self.d_k))\n        self.skip           = nn.Parameter(torch.ones(self.num_types))\n        self.drop           = nn.Dropout(dropout)\n\n        nn.init.xavier_uniform_(self.relation_att)\n        nn.init.xavier_uniform_(self.relation_msg)\n\n    def forward(self, G, h):\n        with G.local_scope():\n            node_dict, edge_dict = self.node_dict, self.edge_dict\n            for srctype, etype, dsttype in G.canonical_etypes:\n                sub_graph = G[srctype, etype, dsttype]\n\n                k_linear = self.k_linears[node_dict[srctype]]\n                v_linear = self.v_linears[node_dict[srctype]]\n                q_linear = self.q_linears[node_dict[dsttype]]\n\n                k = k_linear(h[srctype]).view(-1, self.n_heads, self.d_k)\n                v = v_linear(h[srctype]).view(-1, self.n_heads, self.d_k)\n                q = q_linear(h[dsttype]).view(-1, self.n_heads, self.d_k)\n\n                e_id = self.edge_dict[etype]\n\n                relation_att = self.relation_att[e_id]\n                relation_pri = self.relation_pri[e_id]\n                relation_msg = self.relation_msg[e_id]\n\n                k = torch.einsum(\"bij,ijk->bik\", k, relation_att)\n                v = torch.einsum(\"bij,ijk->bik\", v, relation_msg)\n\n                sub_graph.srcdata['k'] = k\n                sub_graph.dstdata['q'] = q\n                sub_graph.srcdata['v'] = v\n\n                sub_graph.apply_edges(fn.v_dot_u('q', 'k', 't'))\n                attn_score = sub_graph.edata.pop('t').sum(-1) * relation_pri / self.sqrt_dk\n                attn_score = edge_softmax(sub_graph, attn_score, norm_by='dst')\n\n                sub_graph.edata['t'] = attn_score.unsqueeze(-1)\n\n            G.multi_update_all({etype : (fn.u_mul_e('v', 't', 'm'), fn.sum('m', 't')) \\\n                                for etype in edge_dict}, cross_reducer = 'mean')\n\n            new_h = {}\n            for ntype in G.ntypes:\n                '''\n                    Step 3: Target-specific Aggregation\n                    x = norm( W[node_type] * gelu( Agg(x) ) + x )\n                '''\n                n_id = node_dict[ntype]\n                alpha = torch.sigmoid(self.skip[n_id])\n                t = G.nodes[ntype].data['t'].view(-1, self.out_dim)\n                trans_out = self.drop(self.a_linears[n_id](t))\n                trans_out = trans_out * alpha + h[ntype] * (1-alpha)\n                if self.use_norm:\n                    new_h[ntype] = self.norms[n_id](trans_out)\n                else:\n                    new_h[ntype] = trans_out\n            return new_h\n\nclass HGT(nn.Module):\n    def __init__(self, G, node_dict, edge_dict, n_inp, n_hid, n_out, n_layers, n_heads, use_norm = True):\n        super(HGT, self).__init__()\n        self.node_dict = node_dict\n        self.edge_dict = edge_dict\n        self.gcs = nn.ModuleList()\n        self.n_inp = n_inp\n        self.n_hid = n_hid\n        self.n_out = n_out\n        self.n_layers = n_layers\n        self.adapt_ws  = nn.ModuleList()\n        for t in range(len(node_dict)):\n            self.adapt_ws.append(nn.Linear(n_inp,   n_hid))\n        for _ in range(n_layers):\n            self.gcs.append(HGTLayer(n_hid, n_hid, node_dict, edge_dict, n_heads, use_norm = use_norm))\n        self.out = nn.Linear(n_hid, n_out)\n\n    def forward(self, G, out_key):\n        h = {}\n        for ntype in G.ntypes:\n            n_id = self.node_dict[ntype]\n            h[ntype] = F.gelu(self.adapt_ws[n_id](G.nodes[ntype].data['inp']))\n        for i in range(self.n_layers):\n            h = self.gcs[i](G, h)\n        return self.out(h[out_key])\n\nclass HeteroRGCNLayer(nn.Module):\n    def __init__(self, in_size, out_size, etypes):\n        super(HeteroRGCNLayer, self).__init__()\n        # W_r for each relation\n        self.weight = nn.ModuleDict({\n                name : nn.Linear(in_size, out_size) for name in etypes\n            })\n\n    def forward(self, G, feat_dict):\n        # The input is a dictionary of node features for each type\n        funcs = {}\n        for srctype, etype, dsttype in G.canonical_etypes:\n            # Compute W_r * h\n            Wh = self.weight[etype](feat_dict[srctype])\n            # Save it in graph for message passing\n            G.nodes[srctype].data['Wh_%s' % etype] = Wh\n            # Specify per-relation message passing functions: (message_func, reduce_func).\n            # Note that the results are saved to the same destination feature 'h', which\n            # hints the type wise reducer for aggregation.\n            funcs[etype] = (fn.copy_u('Wh_%s' % etype, 'm'), fn.mean('m', 'h'))\n        # Trigger message passing of multiple types.\n        # The first argument is the message passing functions for each relation.\n        # The second one is the type wise reducer, could be \"sum\", \"max\",\n        # \"min\", \"mean\", \"stack\"\n        G.multi_update_all(funcs, 'sum')\n        # return the updated node feature dictionary\n        return {ntype : G.nodes[ntype].data['h'] for ntype in G.ntypes}\n\n\nclass HeteroRGCN(nn.Module):\n    def __init__(self, G, in_size, hidden_size, out_size):\n        super(HeteroRGCN, self).__init__()\n        # create layers\n        self.layer1 = HeteroRGCNLayer(in_size, hidden_size, G.etypes)\n        self.layer2 = HeteroRGCNLayer(hidden_size, out_size, G.etypes)\n\n    def forward(self, G, out_key):\n        input_dict = {ntype : G.nodes[ntype].data['inp'] for ntype in G.ntypes}\n        h_dict = self.layer1(G, input_dict)\n        h_dict = {k : F.leaky_relu(h) for k, h in h_dict.items()}\n        h_dict = self.layer2(G, h_dict)\n        # get paper logits\n        return h_dict[out_key]\n", "path": "examples/pytorch/hgt/model.py"}], "after_files": [{"content": "import dgl\nimport math\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport dgl.function as fn\nfrom dgl.nn.functional import edge_softmax\n\nclass HGTLayer(nn.Module):\n    def __init__(self,\n                 in_dim,\n                 out_dim,\n                 node_dict,\n                 edge_dict,\n                 n_heads,\n                 dropout = 0.2,\n                 use_norm = False):\n        super(HGTLayer, self).__init__()\n\n        self.in_dim        = in_dim\n        self.out_dim       = out_dim\n        self.node_dict     = node_dict\n        self.edge_dict     = edge_dict\n        self.num_types     = len(node_dict)\n        self.num_relations = len(edge_dict)\n        self.total_rel     = self.num_types * self.num_relations * self.num_types\n        self.n_heads       = n_heads\n        self.d_k           = out_dim // n_heads\n        self.sqrt_dk       = math.sqrt(self.d_k)\n        self.att           = None\n\n        self.k_linears   = nn.ModuleList()\n        self.q_linears   = nn.ModuleList()\n        self.v_linears   = nn.ModuleList()\n        self.a_linears   = nn.ModuleList()\n        self.norms       = nn.ModuleList()\n        self.use_norm    = use_norm\n\n        for t in range(self.num_types):\n            self.k_linears.append(nn.Linear(in_dim,   out_dim))\n            self.q_linears.append(nn.Linear(in_dim,   out_dim))\n            self.v_linears.append(nn.Linear(in_dim,   out_dim))\n            self.a_linears.append(nn.Linear(out_dim,  out_dim))\n            if use_norm:\n                self.norms.append(nn.LayerNorm(out_dim))\n\n        self.relation_pri   = nn.Parameter(torch.ones(self.num_relations, self.n_heads))\n        self.relation_att   = nn.Parameter(torch.Tensor(self.num_relations, n_heads, self.d_k, self.d_k))\n        self.relation_msg   = nn.Parameter(torch.Tensor(self.num_relations, n_heads, self.d_k, self.d_k))\n        self.skip           = nn.Parameter(torch.ones(self.num_types))\n        self.drop           = nn.Dropout(dropout)\n\n        nn.init.xavier_uniform_(self.relation_att)\n        nn.init.xavier_uniform_(self.relation_msg)\n\n    def forward(self, G, h):\n        with G.local_scope():\n            node_dict, edge_dict = self.node_dict, self.edge_dict\n            for srctype, etype, dsttype in G.canonical_etypes:\n                sub_graph = G[srctype, etype, dsttype]\n\n                k_linear = self.k_linears[node_dict[srctype]]\n                v_linear = self.v_linears[node_dict[srctype]]\n                q_linear = self.q_linears[node_dict[dsttype]]\n\n                k = k_linear(h[srctype]).view(-1, self.n_heads, self.d_k)\n                v = v_linear(h[srctype]).view(-1, self.n_heads, self.d_k)\n                q = q_linear(h[dsttype]).view(-1, self.n_heads, self.d_k)\n\n                e_id = self.edge_dict[etype]\n\n                relation_att = self.relation_att[e_id]\n                relation_pri = self.relation_pri[e_id]\n                relation_msg = self.relation_msg[e_id]\n\n                k = torch.einsum(\"bij,ijk->bik\", k, relation_att)\n                v = torch.einsum(\"bij,ijk->bik\", v, relation_msg)\n\n                sub_graph.srcdata['k'] = k\n                sub_graph.dstdata['q'] = q\n                sub_graph.srcdata['v_%d' % e_id] = v\n\n                sub_graph.apply_edges(fn.v_dot_u('q', 'k', 't'))\n                attn_score = sub_graph.edata.pop('t').sum(-1) * relation_pri / self.sqrt_dk\n                attn_score = edge_softmax(sub_graph, attn_score, norm_by='dst')\n\n                sub_graph.edata['t'] = attn_score.unsqueeze(-1)\n\n            G.multi_update_all({etype : (fn.u_mul_e('v_%d' % e_id, 't', 'm'), fn.sum('m', 't')) \\\n                                for etype, e_id in edge_dict.items()}, cross_reducer = 'mean')\n\n            new_h = {}\n            for ntype in G.ntypes:\n                '''\n                    Step 3: Target-specific Aggregation\n                    x = norm( W[node_type] * gelu( Agg(x) ) + x )\n                '''\n                n_id = node_dict[ntype]\n                alpha = torch.sigmoid(self.skip[n_id])\n                t = G.nodes[ntype].data['t'].view(-1, self.out_dim)\n                trans_out = self.drop(self.a_linears[n_id](t))\n                trans_out = trans_out * alpha + h[ntype] * (1-alpha)\n                if self.use_norm:\n                    new_h[ntype] = self.norms[n_id](trans_out)\n                else:\n                    new_h[ntype] = trans_out\n            return new_h\n\nclass HGT(nn.Module):\n    def __init__(self, G, node_dict, edge_dict, n_inp, n_hid, n_out, n_layers, n_heads, use_norm = True):\n        super(HGT, self).__init__()\n        self.node_dict = node_dict\n        self.edge_dict = edge_dict\n        self.gcs = nn.ModuleList()\n        self.n_inp = n_inp\n        self.n_hid = n_hid\n        self.n_out = n_out\n        self.n_layers = n_layers\n        self.adapt_ws  = nn.ModuleList()\n        for t in range(len(node_dict)):\n            self.adapt_ws.append(nn.Linear(n_inp,   n_hid))\n        for _ in range(n_layers):\n            self.gcs.append(HGTLayer(n_hid, n_hid, node_dict, edge_dict, n_heads, use_norm = use_norm))\n        self.out = nn.Linear(n_hid, n_out)\n\n    def forward(self, G, out_key):\n        h = {}\n        for ntype in G.ntypes:\n            n_id = self.node_dict[ntype]\n            h[ntype] = F.gelu(self.adapt_ws[n_id](G.nodes[ntype].data['inp']))\n        for i in range(self.n_layers):\n            h = self.gcs[i](G, h)\n        return self.out(h[out_key])\n\nclass HeteroRGCNLayer(nn.Module):\n    def __init__(self, in_size, out_size, etypes):\n        super(HeteroRGCNLayer, self).__init__()\n        # W_r for each relation\n        self.weight = nn.ModuleDict({\n                name : nn.Linear(in_size, out_size) for name in etypes\n            })\n\n    def forward(self, G, feat_dict):\n        # The input is a dictionary of node features for each type\n        funcs = {}\n        for srctype, etype, dsttype in G.canonical_etypes:\n            # Compute W_r * h\n            Wh = self.weight[etype](feat_dict[srctype])\n            # Save it in graph for message passing\n            G.nodes[srctype].data['Wh_%s' % etype] = Wh\n            # Specify per-relation message passing functions: (message_func, reduce_func).\n            # Note that the results are saved to the same destination feature 'h', which\n            # hints the type wise reducer for aggregation.\n            funcs[etype] = (fn.copy_u('Wh_%s' % etype, 'm'), fn.mean('m', 'h'))\n        # Trigger message passing of multiple types.\n        # The first argument is the message passing functions for each relation.\n        # The second one is the type wise reducer, could be \"sum\", \"max\",\n        # \"min\", \"mean\", \"stack\"\n        G.multi_update_all(funcs, 'sum')\n        # return the updated node feature dictionary\n        return {ntype : G.nodes[ntype].data['h'] for ntype in G.ntypes}\n\n\nclass HeteroRGCN(nn.Module):\n    def __init__(self, G, in_size, hidden_size, out_size):\n        super(HeteroRGCN, self).__init__()\n        # create layers\n        self.layer1 = HeteroRGCNLayer(in_size, hidden_size, G.etypes)\n        self.layer2 = HeteroRGCNLayer(hidden_size, out_size, G.etypes)\n\n    def forward(self, G, out_key):\n        input_dict = {ntype : G.nodes[ntype].data['inp'] for ntype in G.ntypes}\n        h_dict = self.layer1(G, input_dict)\n        h_dict = {k : F.leaky_relu(h) for k, h in h_dict.items()}\n        h_dict = self.layer2(G, h_dict)\n        # get paper logits\n        return h_dict[out_key]\n", "path": "examples/pytorch/hgt/model.py"}]}