code: def test_uniform_neighbor_sample_batched(
      dask_client, dataset, input_df, max_batches
  ):
      num_workers = len(dask_client.scheduler_info()["workers"])
  
      df = dataset.get_edgelist()
      df["eid"] = cupy.arange(len(df), dtype=df["src"].dtype)
      df["etp"] = cupy.zeros_like(df["eid"].to_cupy())
      ddf = dask_cudf.from_cudf(df, npartitions=num_workers)
  
      G = cugraph.Graph(directed=True)
      G.from_dask_cudf_edgelist(
          ddf,
          source="src",
          destination="dst",
          edge_attr=["wgt", "eid", "etp"],
      )
  
      input_vertices = (
          dask_cudf.concat([df.src, df.dst]).unique().compute()
      )
      assert isinstance(input_vertices, cudf.Series)
  
      input_vertices.name = "start"
      input_vertices.index = cupy.random.permutation(
          len(input_vertices)
      )
      input_vertices = input_vertices.to_frame().reset_index(drop=True)
  
      input_vertices["batch"] = cudf.Series(
          cupy.random.randint(0, max_batches, len(input_vertices)),
          dtype="int32",
      )
  
      if input_df == dask_cudf.DataFrame:
          input_vertices = dask_cudf.from_cudf(
              input_vertices, npartitions=num_workers
          )
  
      sampling_results = cugraph.dask.uniform_neighbor_sample(
          G,
          start_list=input_vertices,
          fanout_vals=[5, 5],
          with_replacement=False,
          with_edge_properties=True,
          with_batch_ids=True,
      )
  
      for batch_id in range(max_batches):
          output_starts_per_batch = (
              sampling_results[
                  (sampling_results.batch_id == batch_id)
                  & (sampling_results.hop_id == 0)
              ]
              .sources.nunique()
              .compute()
          )
  
          input_starts_per_batch = len(
              input_vertices[input_vertices.batch == batch_id]
          )
  
          # Should be <= to account for starts without outgoing edges
          assert output_starts_per_batch <= input_starts_per_batch
  
imports: import gc
  import random
  import os
  import pytest
  import pandas
  import cupy
  import cudf
  import cugraph
  import dask_cudf
  import cugraph.dask
  from cugraph.testing import UNDIRECTED_DATASETS
  from cugraph.dask import uniform_neighbor_sample
  from cugraph.dask.common.mg_utils import is_single_gpu
  from cugraph.datasets import email_Eu_core, small_tree
  from pylibcugraph.testing.utils import gen_fixture_params_product
label_desc: Use image node_id 13 for calling a global function with example usage: test_uniform_neighbor_sample_batched(dask_client, dataset, input_df, max_batches) without return types