6b8b2026f4501c3d6b62031c79ee38dd64955e1f,torch_geometric/nn/conv/message_passing.py,MessagePassing,propagate,#MessagePassing#Any#Any#,43

Before Change

        message_args = []
        for arg in self.message_args:
            if arg[-2:] == "_i":
                tmp = kwargs[arg[:-2]]
                if tmp is not None:
                    if size[i] is None:
                        size[i] = tmp.size(0)
                    tmp = torch.index_select(tmp, 0, edge_index[i])
                message_args.append(tmp)
            elif arg[-2:] == "_j":
                tmp = kwargs[arg[:-2]]
                if tmp is not None:
                    if size[j] is None:
                        size[j] = tmp.size(0)
                    tmp = torch.index_select(tmp, 0, edge_index[j])
                message_args.append(tmp)
            else:
                message_args.append(kwargs[arg])

        kwargs["size"] = size
        update_args = [kwargs[arg] for arg in self.update_args]

        out = self.message(*message_args)

After Change

            if arg[-2:] in ij.keys():
                tmp = kwargs[arg[:-2]]
                if tmp is not None:
                    if not isinstance(tmp, tuple) and not isinstance(tmp, list):
                        tmp = (tmp, tmp)
                    assert len(tmp) == 2
                    
                    idx = ij[arg[-2:]]
                    if size[idx] is None:
                        size[idx] = tmp[idx].size(0)
                    tmp = torch.index_select(tmp[idx], 0, edge_index[idx])
                message_args.append(tmp)

Instances