Batched heterograph readout

Hi,
How can we apply readout functions (e.g mean nodes) on a batched heterograph ? It is not done the same way as the batched Homogeneous graph

What kind of readout would you like to perform? Do you simply want to compute the average of a node feature for a particular type of nodes?

Yes for the moment only a simple average.
For example : Given a batch of N heterographs such that each heterograph has 2 type of nodes, after applying a RGCN layer on the batch I want apply mean readout such that I get for each graph in the batch and for each nodetype a final representation based on the average embedding of nodes

See if the workaround below is good for you:

import dgl
import dgl.backend as F
import numpy as np
import torch as th

def sum_nodes(bg, node_type, feats):
    batch_size = bg.batch_size
    batch_num_nodes = bg.batch_num_nodes(node_type)

    seg_id = th.from_numpy(np.arange(batch_size, dtype='int64').repeat(batch_num_nodes))
    seg_id = seg_id.to(feats.device)

    return F.unsorted_1d_segment_sum(feats, seg_id, batch_size, 0)

def avg_nodes(bg, node_type, feats):
    batch_node_sum = sum_nodes(bg, node_type, feats)
    batch_num_nodes = torch.tensor(bg.batch_num_nodes(node_type)).float().to(feats.device).reshape(-1, 1)

    return batch_node_sum / batch_num_nodes

g1 = dgl.heterograph({('user', 'plays', 'game'): [(0, 0), (1, 0)]})
g2 = dgl.heterograph({('user', 'plays', 'game'): [(0, 0)]})
g1.nodes['user'].data['h1'] = th.tensor([[1., 2.], [2., 4.]])
g2.nodes['user'].data['h1'] = th.tensor([[3., 6.]])
bg = dgl.batch_hetero([g1, g2])
bg_repr = avg_nodes(bg, 'user', bg.nodes['user'].data['h1'])

Thanks for your reply it answers perfectly my question

Does it still work in 0.4.3? I get a error "module ‘dgl.backend’ has no attribute ‘unsorted_1d_segment_sum’ "

@wdyreborn The code @mufeili provided should still work in 0.4.3 but not in the master branch because that function is deleted. Did you check your DGL’s path? Maybe there are two versions in your system.