Message passing for heterogenous GraphSAGE

Hi,

I’m trying to implement SAGEConv for heterogenous graphs(multiple node types, multiple edge types).

This is the equation I want to use -
gif

I’m trying to modify code from this example

for srctype, etype, dsttype in G.canonical_etypes:
    h = feat_dict[srctype]
    G.nodes[srctype].data['h_%s' % etype] = h
    funcs[etype] = (fn.copy_u('h_%s' % etype, 'm'), fn.mean('m', 'neigh'))
    # Now I want to multiply the mean of neighborhood embedding('neigh')
    # with self.weight[etype](FC layer) and store that in 'h'
G.multi_update_all(funcs, 'sum')

# Add self embedding
G.nodes.data['h'] = G.nodes.data['h'] + self.fc_self * h_self
return {ntype : G.nodes[ntype].data['h'] for ntype in G.ntypes}

How do I multiply the mean neighborhood embedding with W_r in the above code? Is there an easier way?

Thanks,
Sachin.

Either you can define a custom message/reduce function or you can perform update_all separately for each edge type.