code: def test_CopyToWithMiniBatches(task):
      N = 16
      B = 2
      if task == "node_classification" or task == "extra_attrs":
          itemset = gb.ItemSet(
              (torch.arange(N), torch.arange(N)),
              names=("seed_nodes", "labels"),
          )
      elif task == "node_inference":
          itemset = gb.ItemSet(torch.arange(N), names="seed_nodes")
      elif task == "link_prediction":
          itemset = gb.ItemSet(
              (
                  torch.arange(2 * N).reshape(-1, 2),
                  torch.arange(3 * N).reshape(-1, 3),
              ),
              names=("node_pairs", "negative_dsts"),
          )
      elif task == "edge_classification":
          itemset = gb.ItemSet(
              (torch.arange(2 * N).reshape(-1, 2), torch.arange(N)),
              names=("node_pairs", "labels"),
          )
      else:
          itemset = gb.ItemSet(
              (torch.arange(2 * N).reshape(-1, 2), torch.arange(N)),
              names=("node_pairs", "seed_nodes"),
          )
      graph = gb_test_utils.rand_csc_graph(
          100, 0.15, bidirection_edge=True
      )
  
      features = {}
      keys = [("node", None, "a"), ("node", None, "b")]
      features[keys[0]] = gb.TorchBasedFeature(torch.randn(200, 4))
      features[keys[1]] = gb.TorchBasedFeature(torch.randn(200, 4))
      feature_store = gb.BasicFeatureStore(features)
  
      datapipe = gb.ItemSampler(itemset, batch_size=B)
      datapipe = gb.NeighborSampler(
          datapipe,
          graph,
          fanouts=[torch.LongTensor([2]) for _ in range(2)],
      )
      if task != "node_inference":
          datapipe = gb.FeatureFetcher(
              datapipe,
              feature_store,
              ["a"],
          )
  
      if task == "node_classification":
          copied_attrs = [
              "node_features",
              "edge_features",
              "sampled_subgraphs",
              "labels",
              "blocks",
          ]
      elif task == "node_inference":
          copied_attrs = [
              "seed_nodes",
              "sampled_subgraphs",
              "blocks",
              "labels",
          ]
      elif task == "link_prediction":
          copied_attrs = [
              "compacted_node_pairs",
              "node_features",
              "edge_features",
              "sampled_subgraphs",
              "compacted_negative_srcs",
              "compacted_negative_dsts",
              "blocks",
              "positive_node_pairs",
              "negative_node_pairs",
              "node_pairs_with_labels",
          ]
      elif task == "edge_classification":
          copied_attrs = [
              "compacted_node_pairs",
              "node_features",
              "edge_features",
              "sampled_subgraphs",
              "labels",
              "blocks",
              "positive_node_pairs",
              "negative_node_pairs",
              "node_pairs_with_labels",
          ]
      elif task == "extra_attrs":
          copied_attrs = [
              "node_features",
              "edge_features",
              "sampled_subgraphs",
              "labels",
              "blocks",
              "seed_nodes",
          ]
  
      def test_data_device(datapipe):
          for data in datapipe:
              for attr in dir(data):
                  var = getattr(data, attr)
                  if isinstance(var, Mapping):
                      var = var[next(iter(var))]
                  elif isinstance(var, Iterable):
                      var = next(iter(var))
                  if (
                      not callable(var)
                      and not attr.startswith("__")
                      and hasattr(var, "device")
                      and var is not None
                  ):
                      if task == "other":
                          assert var.device.type == "cuda"
                      else:
                          if attr in copied_attrs:
                              assert var.device.type == "cuda"
                          else:
                              assert var.device.type == "cpu"
  
      if task == "extra_attrs":
          extra_attrs = ["seed_nodes"]
      else:
          extra_attrs = None
  
      # Invoke CopyTo via class constructor.
      test_data_device(gb.CopyTo(datapipe, "cuda", extra_attrs))
  
      # Invoke CopyTo via functional form.
      test_data_device(datapipe.copy_to("cuda", extra_attrs))
  
file_path: dgl/tests/python/pytorch/graphbolt/test_base.py
imports: import re
  import unittest
  from collections.abc import Iterable, Mapping
  import backend
  import dgl.graphbolt
  import pytest
  import torch
  from .None import gb_test_utils
label_desc: Use image node_id 2 for calling a global function with example usage: test_CopyToWithMiniBatches(task) without return types