How to pass the SparseTensor and edge_weight into GCNConv at the same time?

[user-AC]

Hi, I'm a beginner in PYG and I have some questions about GCNConv.

In my settings, I need to set GCN's edge_weight to be learnable. At the same time, in order to ensure reproducibility, I need to set

torch.use_deterministic_algorithms(True)

In order to ensure reproducibility, I referred to issue 3175, which converts the input edge_index of GCN into a SparseTensor. In order to ensure the edge_weight is learnable, I referred to issue 2033, i.e.,

self.edge_weight = torch.nn.Parameter(torch.ones(data.num_edges))
...
x = self.conv1(x, edge_index, self.edge_weight.sigmoid()).relu()

However, I found that when edge_index is a SparseTensor, the self.edge_weight will not be used (self.edge_weight will not be updated). This is mainly because of the function gcn_norm() in the GCN source code, it only returns adj_t when edge_index is a SparseTensor!

def gcn_norm(edge_index, edge_weight=None, num_nodes=None, improved=False,
             add_self_loops=True, dtype=None):

    fill_value = 2. if improved else 1.

    if isinstance(edge_index, SparseTensor):
        adj_t = edge_index
        if not adj_t.has_value():
            adj_t = adj_t.fill_value(1., dtype=dtype)
        if add_self_loops:
            adj_t = fill_diag(adj_t, fill_value)
        deg = sparsesum(adj_t, dim=1)
        deg_inv_sqrt = deg.pow_(-0.5)
        deg_inv_sqrt.masked_fill_(deg_inv_sqrt == float('inf'), 0.)
        adj_t = mul(adj_t, deg_inv_sqrt.view(-1, 1))
        adj_t = mul(adj_t, deg_inv_sqrt.view(1, -1))
        return adj_t
....

Question：
So is there any way to ensure that edge_weight is learnable and the code is reproducible (i.e., the algorithm is determined) at the same time?

[rusty1s]

Yes, you can put the edge_weight into the SparseTensor:

adj_t = adj_t.set_value(self.edge_weight)
out = conv(x, adj_t)

[user-AC]

Thank you very much, I'll try it