code: def test_non_uniform_sample():
      shape = (5, 5)
      prob = mx.nd.array([0.9, 0.8, 0.2, 0.4, 0.1], dtype=np.float32)
      data_np = np.array(
          [
              1,
              2,
              3,
              4,
              5,
              6,
              7,
              8,
              9,
              10,
              11,
              12,
              13,
              14,
              15,
              16,
              17,
              18,
              19,
              20,
          ],
          dtype=np.int64,
      )
      indices_np = np.array(
          [1, 2, 3, 4, 0, 2, 3, 4, 0, 1, 3, 4, 0, 1, 2, 4, 0, 1, 2, 3],
          dtype=np.int64,
      )
      indptr_np = np.array([0, 4, 8, 12, 16, 20], dtype=np.int64)
      a = mx.nd.sparse.csr_matrix(
          (data_np, indices_np, indptr_np), shape=shape
      )
  
      seed = mx.nd.array([0, 1, 2, 3, 4], dtype=np.int64)
      out = mx.nd.contrib.dgl_csr_neighbor_non_uniform_sample(
          a,
          prob,
          seed,
          num_args=3,
          num_hops=1,
          num_neighbor=2,
          max_num_vertices=5,
      )
      assert len(out) == 4
      check_non_uniform(out, num_hops=1, max_num_vertices=5)
  
      seed = mx.nd.array([0], dtype=np.int64)
      out = mx.nd.contrib.dgl_csr_neighbor_non_uniform_sample(
          a,
          prob,
          seed,
          num_args=3,
          num_hops=1,
          num_neighbor=1,
          max_num_vertices=4,
      )
      assert len(out) == 4
      check_non_uniform(out, num_hops=1, max_num_vertices=4)
  
      seed = mx.nd.array([0], dtype=np.int64)
      out = mx.nd.contrib.dgl_csr_neighbor_non_uniform_sample(
          a,
          prob,
          seed,
          num_args=3,
          num_hops=2,
          num_neighbor=1,
          max_num_vertices=4,
      )
      assert len(out) == 4
      check_non_uniform(out, num_hops=2, max_num_vertices=4)
  
      seed = mx.nd.array([0, 2, 4], dtype=np.int64)
      out = mx.nd.contrib.dgl_csr_neighbor_non_uniform_sample(
          a,
          prob,
          seed,
          num_args=3,
          num_hops=1,
          num_neighbor=2,
          max_num_vertices=5,
      )
      assert len(out) == 4
      check_non_uniform(out, num_hops=1, max_num_vertices=5)
  
      seed = mx.nd.array([0, 4], dtype=np.int64)
      out = mx.nd.contrib.dgl_csr_neighbor_non_uniform_sample(
          a,
          prob,
          seed,
          num_args=3,
          num_hops=1,
          num_neighbor=2,
          max_num_vertices=5,
      )
      assert len(out) == 4
      check_non_uniform(out, num_hops=1, max_num_vertices=5)
  
      seed = mx.nd.array([0, 4], dtype=np.int64)
      out = mx.nd.contrib.dgl_csr_neighbor_non_uniform_sample(
          a,
          prob,
          seed,
          num_args=3,
          num_hops=2,
          num_neighbor=2,
          max_num_vertices=5,
      )
      assert len(out) == 4
      check_non_uniform(out, num_hops=2, max_num_vertices=5)
  
      seed = mx.nd.array([0, 4], dtype=np.int64)
      out = mx.nd.contrib.dgl_csr_neighbor_non_uniform_sample(
          a,
          prob,
          seed,
          num_args=3,
          num_hops=1,
          num_neighbor=2,
          max_num_vertices=5,
      )
      assert len(out) == 4
      check_non_uniform(out, num_hops=1, max_num_vertices=5)
  
file_path: mxnet/tests/python/unittest/test_dgl_graph.py
imports: from __future__ import print_function
  import numpy
  import scipy
  import mxnet
  import random
  import itertools
  from numpy.testing import assert_allclose, assert_array_equal
  from mxnet.test_utils import *
  import unittest
label_desc: Use image node_id 5 for calling a global function with example usage: test_non_uniform_sample() without return types