code: def test_neighbor_nonuniform(idtype, mode, use_ddp, use_mask):
      if mode != "cpu" and F.ctx() == F.cpu():
          pytest.skip("UVA and GPU sampling require a GPU.")
      if mode != "cpu" and use_mask:
          pytest.skip("Masked sampling only works on CPU.")
      if use_ddp:
          if os.name == "nt":
              pytest.skip(
                  "PyTorch 1.13.0+ has problems in Windows DDP..."
              )
          dist.init_process_group(
              "gloo" if F.ctx() == F.cpu() else "nccl",
              "tcp://127.0.0.1:12347",
              world_size=1,
              rank=0,
          )
      g = dgl.graph(
          ([1, 2, 3, 4, 5, 6, 7, 8], [0, 0, 0, 0, 1, 1, 1, 1])
      ).astype(idtype)
      g.edata["p"] = torch.FloatTensor([1, 1, 0, 0, 1, 1, 0, 0])
      g.edata["mask"] = g.edata["p"] != 0
      if mode in ("cpu", "uva_cpu_indices"):
          indices = F.copy_to(F.tensor([0, 1], idtype), F.cpu())
      else:
          indices = F.copy_to(F.tensor([0, 1], idtype), F.cuda())
      if mode == "pure_gpu":
          g = g.to(F.cuda())
      use_uva = mode.startswith("uva")
  
      if use_mask:
          prob, mask = None, "mask"
      else:
          prob, mask = "p", None
  
      sampler = dgl.dataloading.MultiLayerNeighborSampler(
          [2], prob=prob, mask=mask
      )
      for num_workers in [0, 1, 2] if mode == "cpu" else [0]:
          dataloader = dgl.dataloading.DataLoader(
              g,
              indices,
              sampler,
              batch_size=1,
              device=F.ctx(),
              num_workers=num_workers,
              use_uva=use_uva,
              use_ddp=use_ddp,
          )
          for input_nodes, output_nodes, blocks in dataloader:
              seed = output_nodes.item()
              neighbors = set(input_nodes[1:].cpu().numpy())
              if seed == 1:
                  assert neighbors == {5, 6}
              elif seed == 0:
                  assert neighbors == {1, 2}
  
      g = dgl.heterograph(
          {
              ("B", "BA", "A"): (
                  [1, 2, 3, 4, 5, 6, 7, 8],
                  [0, 0, 0, 0, 1, 1, 1, 1],
              ),
              ("C", "CA", "A"): (
                  [1, 2, 3, 4, 5, 6, 7, 8],
                  [0, 0, 0, 0, 1, 1, 1, 1],
              ),
          }
      ).astype(idtype)
      g.edges["BA"].data["p"] = torch.FloatTensor(
          [1, 1, 0, 0, 1, 1, 0, 0]
      )
      g.edges["BA"].data["mask"] = g.edges["BA"].data["p"] != 0
      g.edges["CA"].data["p"] = torch.FloatTensor(
          [0, 0, 1, 1, 0, 0, 1, 1]
      )
      g.edges["CA"].data["mask"] = g.edges["CA"].data["p"] != 0
      if mode == "pure_gpu":
          g = g.to(F.cuda())
      for num_workers in [0, 1, 2] if mode == "cpu" else [0]:
          dataloader = dgl.dataloading.DataLoader(
              g,
              {"A": indices},
              sampler,
              batch_size=1,
              device=F.ctx(),
              num_workers=num_workers,
              use_uva=use_uva,
              use_ddp=use_ddp,
          )
          for input_nodes, output_nodes, blocks in dataloader:
              seed = output_nodes["A"].item()
              # Seed and neighbors are of different node types so slicing is not necessary here.
              neighbors = set(input_nodes["B"].cpu().numpy())
              if seed == 1:
                  assert neighbors == {5, 6}
              elif seed == 0:
                  assert neighbors == {1, 2}
  
              neighbors = set(input_nodes["C"].cpu().numpy())
              if seed == 1:
                  assert neighbors == {7, 8}
              elif seed == 0:
                  assert neighbors == {3, 4}
  
      if use_ddp:
          dist.destroy_process_group()
  
imports: import os
  import unittest
  from collections.abc import Iterator, Mapping
  from functools import partial
  import backend
  import dgl
  import dgl.ops
  import numpy
  import pytest
  import torch
  import torch.distributed
  import torch.multiprocessing
  from utils import parametrize_idtype
label_desc: Use image node_id 5 for calling a global function with example usage: test_neighbor_nonuniform(idtype, mode, use_ddp, use_mask) without return types