code: 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, edata=["p", "mask"])
      h_g_etype = F.asnumpy(h_g.edata[dgl.ETYPE])
      h_g_offset = np.cumsum(
          np.insert(np.bincount(h_g_etype), 0, 0)
      ).tolist()
      sg = g.edge_subgraph(g.edata["mask"], relabel_nodes=False)
      h_sg = h_g.edge_subgraph(h_g.edata["mask"], relabel_nodes=False)
      h_sg_etype = F.asnumpy(h_sg.edata[dgl.ETYPE])
      h_sg_offset = np.cumsum(
          np.insert(np.bincount(h_sg_etype), 0, 0)
      ).tolist()
  
      seed_ntype = g.get_ntype_id("u")
      seeds = F.nonzero_1d(h_g.ndata[dgl.NTYPE] == seed_ntype)
      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()
          all_etype_array = F.asnumpy(h_g.edata[dgl.ETYPE])
          num_etypes = all_etype_array.max() + 1
          etype_array = F.asnumpy(subg.edata[dgl.ETYPE])
          src = F.asnumpy(src)
          dst = F.asnumpy(dst)
          fanouts = F.asnumpy(fanouts)
  
          all_src = F.asnumpy(all_src)
          all_dst = F.asnumpy(all_dst)
  
          src_per_etype = []
          dst_per_etype = []
          all_src_per_etype = []
          all_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])
              all_src_per_etype.append(
                  all_src[all_etype_array == etype]
              )
              all_dst_per_etype.append(
                  all_dst[all_etype_array == etype]
              )
  
          if replace:
              if direction == "in":
                  in_degree_per_etype = [
                      np.bincount(d) for d in dst_per_etype
                  ]
                  for etype in range(len(fanouts)):
                      in_degree = in_degree_per_etype[etype]
                      fanout = fanouts[etype]
                      ans = np.zeros_like(in_degree)
                      if len(in_degree) > 0:
                          ans[all_dst_per_etype[etype]] = fanout
                      assert np.all(in_degree == ans)
              else:
                  out_degree_per_etype = [
                      np.bincount(s) for s in src_per_etype
                  ]
                  for etype in range(len(fanouts)):
                      out_degree = out_degree_per_etype[etype]
                      fanout = fanouts[etype]
                      ans = np.zeros_like(out_degree)
                      if len(out_degree) > 0:
                          ans[all_src_per_etype[etype]] = fanout
                      assert np.all(out_degree == ans)
          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:
                  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()
      all_sub_src, all_sub_dst = h_sg.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(
              h_g,
              seeds,
              h_g_offset,
              fanouts,
              replace=replace,
              edge_dir=direction,
          )
          check_num(
              h_g, all_src, all_dst, subg, replace, fanouts, direction
          )
  
          p = [g.edges[etype].data["p"] for etype in g.etypes]
          subg = dgl.sampling.sample_etype_neighbors(
              h_g,
              seeds,
              h_g_offset,
              fanouts,
              replace=replace,
              edge_dir=direction,
              prob=p,
          )
          check_num(
              h_sg,
              all_sub_src,
              all_sub_dst,
              subg,
              replace,
              fanouts,
              direction,
          )
  
          p = [g.edges[etype].data["mask"] for etype in g.etypes]
          subg = dgl.sampling.sample_etype_neighbors(
              h_g,
              seeds,
              h_g_offset,
              fanouts,
              replace=replace,
              edge_dir=direction,
              prob=p,
          )
          check_num(
              h_sg,
              all_sub_src,
              all_sub_dst,
              subg,
              replace,
              fanouts,
              direction,
          )
  
imports: import unittest
  import warnings
  import backend
  import dgl
  import numpy
  import pytest
label_desc: Use image node_id 27 for calling a global function with example usage: test_sample_neighbors_etype_homogeneous(format_, direction, replace) without return types