How do I pass edge_weight to neg_z in DeepGraphInfomax?

[zinuoyi]

Hi, I am using GCN as the encoder in DeepGraphInfomax. The forward function in DeepGraphInfomax is as follows:

def forward(self, *args, **kwargs) -> Tuple[Tensor, Tensor, Tensor]:
        """Returns the latent space for the input arguments, their
        corruptions and their summary representation."""
        pos_z = self.encoder(*args, **kwargs)
        cor = self.corruption(*args, **kwargs)
        cor = cor if isinstance(cor, tuple) else (cor, )
        neg_z = self.encoder(*cor)
        summary = self.summary(pos_z, *args, **kwargs)
        return pos_z, neg_z, summary

Ideally, I pass x, edge_index, and edge_weight to self.encoder and get pos_z, then corrupt the data (here I just shuffle the nodes, to be specific), and pass corrupted_x with the same edge_index and edge_weight to self.encoder again and get neg_z.

The problem is, in the argument list of BasicGNN.forward():

    def forward(
        self,
        x: Tensor,
        edge_index: Adj,
        *,
        edge_weight: OptTensor = None,
        edge_attr: OptTensor = None,
        num_sampled_nodes_per_hop: Optional[List[int]] = None,
        num_sampled_edges_per_hop: Optional[List[int]] = None,
    ) -> Tensor:

It seems that the '*' terminates the positional arguments. But the corrupted data obtained from self.corruption is packed as a tuple. So the unpacked edge_weight in the argument list, as a third positional argument, will later trigger an error:

TypeError: BasicGNN.forward() takes 3 positional arguments but 4 were given

I thought about using keyword arguments, but it seems that the 'cor' would be a tuple anyway...
Does somebody know how to solve this problem? Thanks in advance!

Best,
Zinuo

[rusty1s]

Mh, I see. This is a bit non-trivial to fix. One workaround would be to just wrap GCN into a new module:

class MyGCN(torch.nn.Module):
    def __init__(self):
        self.gcn = GCN(...)
        
    def forward(self, x, edge_index, edge_weight):
        return self.gcn(x, edge_index, edge_weight=edge_weight)

I need to think of better ways to allow this though.

[rusty1s]

I added support for this via #7648.