Ignore edge type when sampling neighbours

Is it possible to ignore edge types when sampling neighbours in a heterogenous graph? For example, when I use the fanout = [15] for a graph with multiple edge types, I will receive 15 neighbours per edge type. I would like to receive 15 neighbours/edges total, sampled uniformly across the different edges types.

Any advice?

Have you tried converting the heterogeneous graph to a homogeneous graph first before sampling?

Thanks for the suggestion, but (as far as I can tell), it’s not possible to revert the sampled blocks back to Heterogenous graphs before message passing.

I am not sure if this is what you want. Maybe you can revert them before creating blocks.

class HomoNeighborSampler(NeighborSampler):
    def sample_blocks(self, g, seed_nodes, exclude_eids=None):
        output_nodes = seed_nodes
        blocks = []
        # convert to homogeneous graph first
        # NOTE pass the list of feature names in data fields
        homo_g = dgl.to_homogeneous(g, ndata=['nfeats'], edata=['efeats'])
        for fanout in reversed(self.fanouts):
            # sample on the homogeneous graph
            frontier = homo_g.sample_neighbors(
                seed_nodes, fanout, edge_dir=self.edge_dir, prob=self.prob,
                replace=self.replace, output_device=self.output_device,
                exclude_edges=exclude_eids)
            # revert the subgraph to heterogeneous graph before creating the block
            frontier = dgl.to_heterogeneous(frontier, g.ntypes, g.etypes)
            eid = frontier.edata[EID]
            block = to_block(frontier, seed_nodes)
            block.edata[EID] = eid
            seed_nodes = block.srcdata[NID]
            blocks.insert(0, block)

        return seed_nodes, output_nodes, blocks
1 Like

That seems to be working correctly. Thank you for the support!