Message Passing in Deep Graph Library with UDF

Question

I have the following problem: I am trying to write a message passing function in dgl. After one round of message passing each node should contain a tensor of the form:

[
    [predecessor_node_id, edge_id connecting predecessor and this node, predecessor_node_id], 
    [predecessor_node_id, edge_id connecting predecessor and this node, predecessor_node_id],
    ...
]

I keep on running into issues regarding the dimensionality of my vectors during the reduce phase of my message passing: RuntimeError: Sizes of tensors must match except in dimension 0. Got 1 and 3 in dimension 1 (The offending index is 1) Please enlighten my about my wrong doings.

Minimal Example:

import dgl
import torch
import numpy as np

def initial_send(edges):
    helper = torch.stack((edges._eid, g.find_edges(edges._eid)[0], edges._eid))
    return {"send_message": helper.T}

def initial_reduce(nodes):
    return {"recieved message": nodes.mailbox["send_message"]}

if __name__ == '__main__':
    src = np.array([0, 1, 2, 2, 3, 3, 4, 3])
    dst = np.array([1, 2, 3, 1, 1, 2, 2, 2])
    g = dgl.graph((src, dst))
    g.update_all(initial_send,initial_reduce)