Edge weight support in heteregenous graphs with HGTConv

[SHamda]

I have the following graph structure:

data = HeteroData()
data['word'].x = words_features.float()
#data['root'].x = root_features.float()
#data['lemma'].x = lemma_features.float()
data['user'].x = user_features.float()
data['tweet'].x = tweet_features.float()
data['user', 'writes', 'tweet'].edge_index = edge_index_t_u
data['word', 'links', 'lemma'].edge_index = edge_index_w_l
data['lemma', 'links', 'root'].edge_index = edge_index_l_r
data['word', 'in', 'tweet'].edge_index = edge_index_w_t
data['user'].y = None
data['word'].y = None
data['lemma'].y = None
data['root'].y = None
data['tweet'].y = torch.tensor(df_tweets['label'], dtype=torch.long)

For each edge type, I would like to add an edge weight, I cant find the documentation to add edgeweight in the case of a heterodata object. Furthermore, how can I use the added edge weight knowing that I'm using HGTconv as GNN layer.

NN architecutre:

import torch_geometric.transforms as T
from torch_geometric.datasets import OGB_MAG
from torch_geometric.nn import HGTConv, Linear
h_c = 256
class HGT(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels, num_heads, num_layers):
        super().__init__()

        self.lin_dict = torch.nn.ModuleDict()
        for node_type in data.node_types:
            self.lin_dict[node_type] = Linear(-1, hidden_channels)

        self.convs = torch.nn.ModuleList()
        for _ in range(num_layers):
            conv = HGTConv(hidden_channels, hidden_channels, data.metadata(),
                           num_heads, group='sum')
            self.convs.append(conv)

        self.linear1 = Linear(hidden_channels, h_c)
        self.dropout = torch.nn.Dropout(p=0.7)
        self.linear2 = Linear(h_c, 128)
        self.dropout2 = torch.nn.Dropout(p=0.5)
        self.linear3 = Linear(128, out_channels)

    def forward(self, x_dict, edge_index_dict):
        for node_type, x in x_dict.items():
            x_dict[node_type] = self.lin_dict[node_type](x).tanh_()

        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict)

        x = x_dict['tweet']
        x = self.dropout(x)
        x = torch.relu(self.linear1(x))
        x = self.dropout2(x)
        x = torch.relu(self.linear2(x))
        x = self.dropout2(x)
        x = self.linear3(x)

        return x
model = HGT(hidden_channels=512, out_channels=2,
            num_heads=8, num_layers=2)

I'm aware of the edge_weight support issues shown in this link https://pytorch-geometric.readthedocs.io/en/latest/notes/cheatsheet.html
So what would be the standard way to approach such an issue.

[rusty1s]

You can simply add the edge_weight to your data object, e.g.:

data['user', 'writes', 'tweet'].edge_weight = edge_weight

Note that HGTConv doesn't support edge weights, so you would need to use a different layer in this case. For example, you can construct a heterogeneous layer via HeteroConv by making use of GNN layers in PyG that support edge weights.