Here is the full message passing code:
import inspect
import torch
from torch_geometric.utils import scatter_
class MessagePassing(torch.nn.Module):
def init(self, aggr=‘add’, flow=‘source_to_target’):
super(MessagePassing, self).init()
    self.aggr = aggr
    assert self.aggr in ['add', 'mean', 'max']

    self.flow = flow
    assert self.flow in ['source_to_target', 'target_to_source']

    self.message_args = inspect.getargspec(self.message)[0][1:]
    self.update_args = inspect.getargspec(self.update)[0][2:]

def propagate(self, edge_index, size=None, **kwargs):
    r"""The initial call to start propagating messages.

    Args:
        edge_index (Tensor): The indices of a general (sparse) assignment
            matrix with shape :obj:`[N, M]` (can be directed or
            undirected).
        size (list or tuple, optional): The size :obj:`[N, M]` of the
            assignment matrix. If set to :obj:`None`, the size is tried to
            get automatically inferrred. (default: :obj:`None`)
        **kwargs: Any additional data which is needed to construct messages
            and to update node embeddings.
    """

    size = [None, None] if size is None else list(size)
    assert len(size) == 2

    kwargs['edge_index'] = edge_index
    i, j = (0, 1) if self.flow == 'target_to_source' else (1, 0)
    ij = {"_i": i, "_j": j}
    
    message_args = []
    for arg in self.message_args:
        if arg[-2:] in ij.keys():
            tmp = kwargs[arg[:-2]]
            if tmp is None:
                message_args.append(tmp)
            else:
                idx = ij[arg[-2:]]
                if isinstance(tmp, tuple):
                    assert len(tmp) == 2
                    if size[1 - idx] is None:
                        size[1 - idx] = tmp[1 - idx].size(0)
                    tmp = tmp[idx]

                if size[idx] is None:
                    size[idx] = tmp.size(0)
                #print(edge_index)
                #print(edge_index[idx])
                #print(idx)
                tmp = torch.index_select(tmp, 0, edge_index[idx])
                message_args.append(tmp)
        else:
            message_args.append(kwargs[arg])

    size[0] = size[1] if size[0] is None else size[0]
    size[1] = size[0] if size[1] is None else size[1]

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

    out = self.message(*message_args)
    if self.aggr in ["add", "mean", "max"]:
        out = scatter_(self.aggr, out, edge_index[i], dim_size=size[i])
    else:
        pass
    out = self.update(out, *update_args)

    return out

def message(self, x_j):  # pragma: no cover
    return x_j

def update(self, aggr_out):  # pragma: no cover
    return aggr_out