Repository: dmlc/dgl

Issue:
Edge type specific fanout in the distributed dgl.
## 🚀 Feature
<!-- A brief description of the feature proposal -->

We need to have edge type specific fanout support in the distributed dgl. This will unlock a set of new applications that will require less samples from specific edge types that may be noisy.

We would like the same support as here
https://docs.dgl.ai/en/0.6.x/generated/dgl.sampling.sample_neighbors.html
for the distributed sample neighbor function
https://docs.dgl.ai/en/0.6.x/generated/dgl.distributed.graph_services.sample_neighbors.html#dgl.distributed.graph_services.sample_neighbors

<!-- Please outline the motivation for the proposal. Is your feature request
related to a problem? e.g., I'm always frustrated when [...]. If this is
related to another GitHub issue, please link here too -->

## Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

## Pitch

<!-- A clear and concise description of what you want to happen. -->

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->



Discussion:

Resolving patch:
diff --git a/python/dgl/dataloading/neighbor.py b/python/dgl/dataloading/neighbor.py
--- a/python/dgl/dataloading/neighbor.py
+++ b/python/dgl/dataloading/neighbor.py
@@ -132,9 +132,6 @@ def sample_frontier(self, block_id, g, seed_nodes, exclude_eids=None):
         fanout = self.fanouts[block_id]
         if isinstance(g, distributed.DistGraph):
             if len(g.etypes) > 1: # heterogeneous distributed graph
-                # The edge type is stored in g.edata[dgl.ETYPE]
-                assert isinstance(fanout, int), "For distributed training, " \
-                    "we can only sample same number of neighbors for each edge type"
                 frontier = distributed.sample_etype_neighbors(
                     g, seed_nodes, ETYPE, fanout, replace=self.replace)
             else:
diff --git a/python/dgl/distributed/graph_services.py b/python/dgl/distributed/graph_services.py
--- a/python/dgl/distributed/graph_services.py
+++ b/python/dgl/distributed/graph_services.py
@@ -450,8 +450,9 @@ def sample_etype_neighbors(g, nodes, etype_field, fanout, edge_dir='in', prob=No
         one key-value pair to make this API consistent with dgl.sampling.sample_neighbors.
     etype_field : string
         The field in g.edata storing the edge type.
-    fanout : int
-        The number of edges to be sampled for each node per edge type.
+    fanout : int or dict[etype, int]
+        The number of edges to be sampled for each node per edge type.  If an integer
+        is given, DGL assumes that the same fanout is applied to every edge type.
 
         If -1 is given, all of the neighbors will be selected.
     edge_dir : str, optional
@@ -479,6 +480,11 @@ def sample_etype_neighbors(g, nodes, etype_field, fanout, edge_dir='in', prob=No
     DGLGraph
         A sampled subgraph containing only the sampled neighboring edges.  It is on CPU.
     """
+    if isinstance(fanout, int):
+        fanout = F.full_1d(len(g.etypes), fanout, F.int64, F.cpu())
+    else:
+        fanout = F.tensor([fanout[etype] for etype in g.etypes], dtype=F.int64)
+
     gpb = g.get_partition_book()
     if isinstance(nodes, dict):
         homo_nids = []
diff --git a/python/dgl/sampling/neighbor.py b/python/dgl/sampling/neighbor.py
--- a/python/dgl/sampling/neighbor.py
+++ b/python/dgl/sampling/neighbor.py
@@ -36,14 +36,11 @@ def sample_etype_neighbors(g, nodes, etype_field, fanout, edge_dir='in', prob=No
         If a single tensor is given, the graph must only have one type of nodes.
     etype_field : string
         The field in g.edata storing the edge type.
-    fanout : int
-        The number of edges to be sampled for each node on each edge type.
-
-        This argument can only take a single int. DGL will sample this number of edges for
-        each node for every edge type.
+    fanout : Tensor
+        The number of edges to be sampled for each node per edge type.  Must be a
+        1D tensor with the number of elements same as the number of edge types.
 
-        If -1 is given for a single edge type, all the neighboring edges with that edge
-        type will be selected.
+        If -1 is given, all of the neighbors will be selected.
     edge_dir : str, optional
         Determines whether to sample inbound or outbound edges.
 
@@ -99,15 +96,19 @@ def sample_etype_neighbors(g, nodes, etype_field, fanout, edge_dir='in', prob=No
     if etype_field not in g.edata:
         raise DGLError("The graph should have {} in the edge data" \
                        "representing the edge type.".format(etype_field))
-    if isinstance(fanout, int) is False:
-        raise DGLError("The fanout should be an integer")
-    if isinstance(nodes, dict) is True:
+    # (BarclayII) because the homogenized graph no longer contains the *name* of edge
+    # types, the fanout argument can no longer be a dict of etypes and ints, as opposed
+    # to sample_neighbors.
+    if not F.is_tensor(fanout):
+        raise DGLError("The fanout should be a tensor")
+    if isinstance(nodes, dict):
         assert len(nodes) == 1, "The input graph should not have node types"
         nodes = list(nodes.values())[0]
     nodes = F.to_dgl_nd(utils.prepare_tensor(g, nodes, 'nodes'))
     # treat etypes as int32, it is much cheaper than int64
     # TODO(xiangsx): int8 can be a better choice.
     etypes = F.to_dgl_nd(F.astype(g.edata[etype_field], ty=F.int32))
+    fanout = F.to_dgl_nd(fanout)
 
     if prob is None:
         prob_array = nd.array([], ctx=nd.cpu())


Test patch:
diff --git a/tests/compute/test_sampling.py b/tests/compute/test_sampling.py
--- a/tests/compute/test_sampling.py
+++ b/tests/compute/test_sampling.py
@@ -725,106 +725,75 @@ def check_num(nodes, tag):
         check_num(subg.edges()[1], tag)
 
 @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sample neighbors not implemented")
-def test_sample_neighbors_etype_homogeneous():
+@pytest.mark.parametrize('format_', ['coo', 'csr', 'csc'])
+@pytest.mark.parametrize('direction', ['in', 'out'])
+@pytest.mark.parametrize('replace', [False, True])
+def test_sample_neighbors_etype_homogeneous(format_, direction, replace):
     num_nodes = 100
     rare_cnt = 4
     g = create_etype_test_graph(100, 30, rare_cnt)
     h_g = dgl.to_homogeneous(g)
     seed_ntype = g.get_ntype_id("u")
     seeds = F.nonzero_1d(h_g.ndata[dgl.NTYPE] == seed_ntype)
-
-    def check_num(nodes, replace):
-        nodes = F.asnumpy(nodes)
-        cnt = [sum(nodes == i) for i in range(num_nodes)]
-
-        for i in range(20):
-            if i < rare_cnt:
-                if replace is False:
-                    assert cnt[i] == 22
-                else:
-                    assert cnt[i] == 30
+    fanouts = F.tensor([6, 5, 4, 3, 2], dtype=F.int64)
+
+    def check_num(h_g, all_src, all_dst, subg, replace, fanouts, direction):
+        src, dst = subg.edges()
+        num_etypes = F.asnumpy(h_g.edata[dgl.ETYPE]).max()
+        etype_array = F.asnumpy(subg.edata[dgl.ETYPE])
+        src = F.asnumpy(src)
+        dst = F.asnumpy(dst)
+        fanouts = F.asnumpy(fanouts)
+
+        all_etype_array = F.asnumpy(h_g.edata[dgl.ETYPE])
+        all_src = F.asnumpy(all_src)
+        all_dst = F.asnumpy(all_dst)
+
+        src_per_etype = []
+        dst_per_etype = []
+        for etype in range(num_etypes):
+            src_per_etype.append(src[etype_array == etype])
+            dst_per_etype.append(dst[etype_array == etype])
+
+        if replace:
+            if direction == 'in':
+                in_degree_per_etype = [np.bincount(d) for d in dst_per_etype]
+                for in_degree, fanout in zip(in_degree_per_etype, fanouts):
+                    assert np.all(in_degree == fanout)
             else:
-                if replace is False:
-                    assert cnt[i] == 12
-                else:
-                    assert cnt[i] == 20
-
-    # graph with coo format
-    coo_g = h_g.formats('coo')
-    for _ in range(5):
-        subg = dgl.sampling.sample_etype_neighbors(coo_g, seeds, dgl.ETYPE, 10, replace=False)
-        check_num(subg.edges()[1], False)
-
-    for _ in range(5):
-        subg = dgl.sampling.sample_etype_neighbors(coo_g, seeds, dgl.ETYPE, 10, replace=True)
-        check_num(subg.edges()[1], True)
-
-    # graph with csr format
-    csr_g = h_g.formats('csr')
-    csr_g = csr_g.formats(['csr','csc','coo'])
-    for _ in range(5):
-        subg = dgl.sampling.sample_etype_neighbors(csr_g, seeds, dgl.ETYPE, 10, replace=False)
-        check_num(subg.edges()[1], False)
-
-    for _ in range(5):
-        subg = dgl.sampling.sample_etype_neighbors(csr_g, seeds, dgl.ETYPE, 10, replace=True)
-        check_num(subg.edges()[1], True)
-
-    # graph with csc format
-    csc_g = h_g.formats('csc')
-    for _ in range(5):
-        subg = dgl.sampling.sample_etype_neighbors(csc_g, seeds, dgl.ETYPE, 10, replace=False)
-        check_num(subg.edges()[1], False)
-
-    for _ in range(5):
-        subg = dgl.sampling.sample_etype_neighbors(csc_g, seeds, dgl.ETYPE, 10, replace=True)
-        check_num(subg.edges()[1], True)
-
-    def check_num2(nodes, replace):
-        nodes = F.asnumpy(nodes)
-        cnt = [sum(nodes == i) for i in range(num_nodes)]
-
-        for i in range(20):
-            if replace is False:
-                assert cnt[i] == 7
+                out_degree_per_etype = [np.bincount(s) for s in src_per_etype]
+                for out_degree, fanout in zip(out_degree_per_etype, fanouts):
+                    assert np.all(out_degree == fanout)
+        else:
+            if direction == 'in':
+                for v in set(dst):
+                    u = src[dst == v]
+                    et = etype_array[dst == v]
+                    all_u = all_src[all_dst == v]
+                    all_et = all_etype_array[all_dst == v]
+                    for etype in set(et):
+                        u_etype = set(u[et == etype])
+                        all_u_etype = set(all_u[all_et == etype])
+                        assert (len(u_etype) == fanouts[etype]) or (u_etype == all_u_etype)
             else:
-                assert cnt[i] == 10
-
-    # edge dir out
-    # graph with coo format
-    coo_g = h_g.formats('coo')
-    for _ in range(5):
-        subg = dgl.sampling.sample_etype_neighbors(
-            coo_g, seeds, dgl.ETYPE, 5, edge_dir='out', replace=False)
-        check_num2(subg.edges()[0], False)
-    for _ in range(5):
-        subg = dgl.sampling.sample_etype_neighbors(
-            coo_g, seeds, dgl.ETYPE, 5, edge_dir='out', replace=True)
-        check_num2(subg.edges()[0], True)
-    # graph with csr format
-    csr_g = h_g.formats('csr')
-    for _ in range(5):
-        subg = dgl.sampling.sample_etype_neighbors(
-            csr_g, seeds, dgl.ETYPE, 5, edge_dir='out', replace=False)
-        check_num2(subg.edges()[0], False)
-
-    for _ in range(5):
-        subg = dgl.sampling.sample_etype_neighbors(
-            csr_g, seeds, dgl.ETYPE, 5, edge_dir='out', replace=True)
-        check_num2(subg.edges()[0], True)
-
-    # graph with csc format
-    csc_g = h_g.formats('csc')
-    csc_g = csc_g.formats(['csc','csr','coo'])
-    for _ in range(5):
-        subg = dgl.sampling.sample_etype_neighbors(
-            csc_g, seeds, dgl.ETYPE, 5, edge_dir='out', replace=False)
-        check_num2(subg.edges()[0], False)
-
+                for u in set(src):
+                    v = dst[src == u]
+                    et = etype_array[src == u]
+                    all_v = all_dst[all_src == u]
+                    all_et = all_etype_array[all_src == u]
+                    for etype in set(et):
+                        v_etype = set(v[et == etype])
+                        all_v_etype = set(all_v[all_et == etype])
+                        assert (len(v_etype) == fanouts[etype]) or (v_etype == all_v_etype)
+
+    all_src, all_dst = h_g.edges()
+    h_g = h_g.formats(format_)
+    if (direction, format_) in [('in', 'csr'), ('out', 'csc')]:
+        h_g = h_g.formats(['csc', 'csr', 'coo'])
     for _ in range(5):
         subg = dgl.sampling.sample_etype_neighbors(
-            csc_g, seeds, dgl.ETYPE, 5, edge_dir='out', replace=True)
-        check_num2(subg.edges()[0], True)
+            h_g, seeds, dgl.ETYPE, fanouts, replace=replace, edge_dir=direction)
+        check_num(h_g, all_src, all_dst, subg, replace, fanouts, direction)
 
 @pytest.mark.parametrize('dtype', ['int32', 'int64'])
 @unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sample neighbors not implemented")
@@ -920,7 +889,9 @@ def test_sample_neighbors_exclude_edges_homoG(dtype):
 
 
 if __name__ == '__main__':
-    test_sample_neighbors_etype_homogeneous()
+    from itertools import product
+    for args in product(['coo', 'csr', 'csc'], ['in', 'out'], [False, True]):
+        test_sample_neighbors_etype_homogeneous(*args)
     test_random_walk()
     test_pack_traces()
     test_pinsage_sampling()
diff --git a/tests/cpp/test_rowwise.cc b/tests/cpp/test_rowwise.cc
--- a/tests/cpp/test_rowwise.cc
+++ b/tests/cpp/test_rowwise.cc
@@ -264,11 +264,11 @@ void _TestCSRPerEtypeSampling(bool has_data) {
       NDArray::FromVector(std::vector<int32_t>({3, 1, 3, 3, 2, 3, 0})) :
       NDArray::FromVector(std::vector<int32_t>({3, 3, 3, 0, 3, 1, 2}));
   for (int k = 0; k < 10; ++k) {
-    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, 2, prob, true);
+    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, true);
     CheckSampledPerEtypeReplaceResult<Idx>(rst, rows, has_data);
   }
   for (int k = 0; k < 10; ++k) {
-    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, 2, prob, false);
+    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, false);
     CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
     auto eset = ToEdgeSet<Idx>(rst);
     if (has_data) {
@@ -316,7 +316,7 @@ void _TestCSRPerEtypeSampling(bool has_data) {
     NDArray::FromVector(
       std::vector<FloatType>({.0, .5, .0, .5, .0, .5, .5}));
   for (int k = 0; k < 10; ++k) {
-    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, 2, prob, true);
+    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, true);
     CheckSampledPerEtypeReplaceResult<Idx>(rst, rows, has_data);
     auto eset = ToEdgeSet<Idx>(rst);
     if (has_data) {
@@ -339,11 +339,11 @@ void _TestCSRPerEtypeSamplingSorted(bool has_data, bool etype_sorted) {
       NDArray::FromVector(std::vector<int32_t>({0, 1, 0, 0, 2, 0, 3})) :
       NDArray::FromVector(std::vector<int32_t>({0, 0, 0, 3, 0, 1, 2}));
   for (int k = 0; k < 10; ++k) {
-    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, 2, prob, true, etype_sorted);
+    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, true, etype_sorted);
     CheckSampledPerEtypeReplaceResult<Idx>(rst, rows, has_data);
   }
   for (int k = 0; k < 10; ++k) {
-    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, 2, prob, false, etype_sorted);
+    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, false, etype_sorted);
     CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
     auto eset = ToEdgeSet<Idx>(rst);
     if (has_data) {
@@ -391,7 +391,7 @@ void _TestCSRPerEtypeSamplingSorted(bool has_data, bool etype_sorted) {
     NDArray::FromVector(
       std::vector<FloatType>({.0, .5, .0, .5, .0, .5, .5}));
   for (int k = 0; k < 10; ++k) {
-    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, 2, prob, true, etype_sorted);
+    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, true, etype_sorted);
     CheckSampledPerEtypeReplaceResult<Idx>(rst, rows, has_data);
     auto eset = ToEdgeSet<Idx>(rst);
     if (has_data) {
@@ -440,12 +440,12 @@ void _TestCSRPerEtypeSamplingUniform(bool has_data) {
       NDArray::FromVector(std::vector<int32_t>({3, 1, 3, 3, 2, 3, 0})) :
       NDArray::FromVector(std::vector<int32_t>({3, 3, 3, 0, 3, 1, 2}));
   for (int k = 0; k < 10; ++k) {
-    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, 2, prob, true);
+    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, true);
     CheckSampledPerEtypeReplaceResult<Idx>(rst, rows, has_data);
   }
 
   for (int k = 0; k < 10; ++k) {
-    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, 2, prob, false);
+    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, false);
     CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
     auto eset = ToEdgeSet<Idx>(rst);
     if (has_data) {
@@ -497,12 +497,12 @@ void _TestCSRPerEtypeSamplingUniformSorted(bool has_data, bool etype_sorted) {
       NDArray::FromVector(std::vector<int32_t>({0, 1, 0, 0, 2, 0, 3})) :
       NDArray::FromVector(std::vector<int32_t>({0, 0, 0, 3, 0, 1, 2}));
   for (int k = 0; k < 10; ++k) {
-    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, 2, prob, true, etype_sorted);
+    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, true, etype_sorted);
     CheckSampledPerEtypeReplaceResult<Idx>(rst, rows, has_data);
   }
 
   for (int k = 0; k < 10; ++k) {
-    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, 2, prob, false, etype_sorted);
+    auto rst = CSRRowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, false, etype_sorted);
     CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
     auto eset = ToEdgeSet<Idx>(rst);
     if (has_data) {
@@ -674,11 +674,11 @@ void _TestCOOerEtypeSampling(bool has_data) {
       NDArray::FromVector(std::vector<int32_t>({3, 1, 3, 3, 2, 3, 0})) :
       NDArray::FromVector(std::vector<int32_t>({3, 3, 3, 0, 3, 1, 2}));
   for (int k = 0; k < 10; ++k) {
-    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, 2, prob, true);
+    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, true);
     CheckSampledPerEtypeReplaceResult<Idx>(rst, rows, has_data);
   }
   for (int k = 0; k < 10; ++k) {
-    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, 2, prob, false);
+    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, false);
     CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
     auto eset = ToEdgeSet<Idx>(rst);
     if (has_data) {
@@ -726,7 +726,7 @@ void _TestCOOerEtypeSampling(bool has_data) {
     NDArray::FromVector(
       std::vector<FloatType>({.0, .5, .0, .5, .0, .5, .5}));
   for (int k = 0; k < 10; ++k) {
-    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, 2, prob, true);
+    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, true);
     CheckSampledPerEtypeReplaceResult<Idx>(rst, rows, has_data);
     auto eset = ToEdgeSet<Idx>(rst);
     if (has_data) {
@@ -749,11 +749,11 @@ void _TestCOOerEtypeSamplingSorted(bool has_data, bool etype_sorted) {
       NDArray::FromVector(std::vector<int32_t>({0, 1, 0, 0, 2, 0, 3})) :
       NDArray::FromVector(std::vector<int32_t>({0, 0, 0, 3, 0, 1, 2}));
   for (int k = 0; k < 10; ++k) {
-    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, 2, prob, true, etype_sorted);
+    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, true, etype_sorted);
     CheckSampledPerEtypeReplaceResult<Idx>(rst, rows, has_data);
   }
   for (int k = 0; k < 10; ++k) {
-    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, 2, prob, false, etype_sorted);
+    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, false, etype_sorted);
     CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
     auto eset = ToEdgeSet<Idx>(rst);
     if (has_data) {
@@ -801,7 +801,7 @@ void _TestCOOerEtypeSamplingSorted(bool has_data, bool etype_sorted) {
     NDArray::FromVector(
       std::vector<FloatType>({.0, .5, .0, .5, .0, .5, .5}));
   for (int k = 0; k < 10; ++k) {
-    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, 2, prob, true, etype_sorted);
+    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, true, etype_sorted);
     CheckSampledPerEtypeReplaceResult<Idx>(rst, rows, has_data);
     auto eset = ToEdgeSet<Idx>(rst);
     if (has_data) {
@@ -850,12 +850,12 @@ void _TestCOOPerEtypeSamplingUniform(bool has_data) {
       NDArray::FromVector(std::vector<int32_t>({3, 1, 3, 3, 2, 3, 0})) :
       NDArray::FromVector(std::vector<int32_t>({3, 3, 3, 0, 3, 1, 2}));
   for (int k = 0; k < 10; ++k) {
-    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, 2, prob, true);
+    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, true);
     CheckSampledPerEtypeReplaceResult<Idx>(rst, rows, has_data);
   }
 
   for (int k = 0; k < 10; ++k) {
-    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, 2, prob, false);
+    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, false);
     CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
     auto eset = ToEdgeSet<Idx>(rst);
     if (has_data) {
@@ -907,12 +907,12 @@ void _TestCOOPerEtypeSamplingUniformSorted(bool has_data, bool etype_sorted) {
       NDArray::FromVector(std::vector<int32_t>({0, 1, 0, 0, 2, 0, 3})) :
       NDArray::FromVector(std::vector<int32_t>({0, 0, 0, 3, 0, 1, 2}));
   for (int k = 0; k < 10; ++k) {
-    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, 2, prob, true, etype_sorted);
+    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, true, etype_sorted);
     CheckSampledPerEtypeReplaceResult<Idx>(rst, rows, has_data);
   }
 
   for (int k = 0; k < 10; ++k) {
-    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, 2, prob, false, etype_sorted);
+    auto rst = COORowWisePerEtypeSampling(mat, rows, etypes, {2, 2, 2, 2}, prob, false, etype_sorted);
     CheckSampledPerEtypeResult<Idx>(rst, rows, has_data);
     auto eset = ToEdgeSet<Idx>(rst);
     if (has_data) {
@@ -1134,4 +1134,4 @@ TEST(RowwiseTest, TestCSRSamplingBiased) {
   _TestCSRSamplingBiased<int32_t, double>(false);
   _TestCSRSamplingBiased<int64_t, double>(true);
   _TestCSRSamplingBiased<int64_t, double>(false);
-}
\ No newline at end of file
+}
diff --git a/tests/distributed/test_mp_dataloader.py b/tests/distributed/test_mp_dataloader.py
--- a/tests/distributed/test_mp_dataloader.py
+++ b/tests/distributed/test_mp_dataloader.py
@@ -278,7 +278,10 @@ def start_node_dataloader(rank, tmpdir, num_server, num_workers, orig_nid, orig_
         part, _, _, _, _, _, _ = load_partition(tmpdir / 'test_sampling.json', i)
 
     # Create sampler
-    sampler = dgl.dataloading.MultiLayerNeighborSampler([5, 10])
+    sampler = dgl.dataloading.MultiLayerNeighborSampler([
+        # test dict for hetero
+        {etype: 5 for etype in dist_graph.etypes} if len(dist_graph.etypes) > 1 else 5,
+        10])        # test int for hetero
 
     # We need to test creating DistDataLoader multiple times.
     for i in range(2):