code: def test_custom_hetero_neighbor_loader():
      # Initialize feature store, graph store, and reference:
      feature_store = MyFeatureStore()
      graph_store = MyGraphStore()
      data = HeteroData()
  
      # Set up node features:
      x = torch.arange(100)
      data["paper"].x = x
      feature_store.put_tensor(
          x, group_name="paper", attr_name="x", index=None
      )
  
      x = torch.arange(100, 300)
      data["author"].x = x
      feature_store.put_tensor(
          x, group_name="author", attr_name="x", index=None
      )
  
      # COO:
      edge_index = get_random_edge_index(100, 100, 500, coalesce=True)
      edge_index = edge_index[:, torch.randperm(edge_index.size(1))]
      data["paper", "to", "paper"].edge_index = edge_index
      coo = (edge_index[0], edge_index[1])
      graph_store.put_edge_index(
          edge_index=coo,
          edge_type=("paper", "to", "paper"),
          layout="coo",
          size=(100, 100),
      )
  
      # CSR:
      edge_index = get_random_edge_index(100, 200, 1000, coalesce=True)
      data["paper", "to", "author"].edge_index = edge_index
      adj = to_torch_csr_tensor(edge_index, size=(100, 200))
      csr = (adj.crow_indices(), adj.col_indices())
      graph_store.put_edge_index(
          edge_index=csr,
          edge_type=("paper", "to", "author"),
          layout="csr",
          size=(100, 200),
      )
  
      # CSC:
      edge_index = get_random_edge_index(200, 100, 1000, coalesce=True)
      data["author", "to", "paper"].edge_index = edge_index
      adj = to_torch_csr_tensor(edge_index.flip([0]), size=(100, 200))
      csc = (adj.col_indices(), adj.crow_indices())
      graph_store.put_edge_index(
          edge_index=csc,
          edge_type=("author", "to", "paper"),
          layout="csc",
          size=(200, 100),
      )
  
      # COO (sorted):
      edge_index = get_random_edge_index(200, 200, 100, coalesce=True)
      edge_index = edge_index[:, edge_index[1].argsort()]
      data["author", "to", "author"].edge_index = edge_index
      coo = (edge_index[0], edge_index[1])
      graph_store.put_edge_index(
          edge_index=coo,
          edge_type=("author", "to", "author"),
          layout="coo",
          size=(200, 200),
          is_sorted=True,
      )
  
      # Construct neighbor loaders:
      loader1 = NeighborLoader(
          data,
          batch_size=20,
          input_nodes=("paper", range(100)),
          num_neighbors=[-1] * 2,
      )
  
      loader2 = NeighborLoader(
          (feature_store, graph_store),
          batch_size=20,
          input_nodes=("paper", range(100)),
          num_neighbors=[-1] * 2,
      )
  
      assert str(loader1) == str(loader2)
      assert len(loader1) == len(loader2)
  
      for batch1, batch2 in zip(loader1, loader2):
          # `loader2` explicitly adds `num_nodes` to the batch:
          assert len(batch1) + 1 == len(batch2)
          assert (
              batch1["paper"].batch_size == batch2["paper"].batch_size
          )
  
          # Mapped indices of neighbors may be differently sorted ...
          for node_type in data.node_types:
              assert torch.allclose(
                  batch1[node_type].x.sort()[0],
                  batch2[node_type].x.sort()[0],
              )
  
          # ... but should sample the exact same number of edges:
          for edge_type in data.edge_types:
              assert (
                  batch1[edge_type].num_edges
                  == batch2[edge_type].num_edges
              )
  
imports: import os.path
  import numpy
  import pytest
  import torch
  import torch_geometric.typing
  from torch_geometric import EdgeIndex
  from torch_geometric.data import Data, HeteroData
  from torch_geometric.loader import NeighborLoader
  from torch_geometric.nn import GraphConv, to_hetero
  from torch_geometric.sampler.base import SubgraphType
  from torch_geometric.testing import MyFeatureStore, MyGraphStore, get_random_edge_index, get_random_tensor_frame, onlyLinux, onlyNeighborSampler, onlyOnline, withCUDA, withPackage
  from torch_geometric.typing import WITH_EDGE_TIME_NEIGHBOR_SAMPLE, WITH_PYG_LIB, WITH_TORCH_SPARSE, WITH_WEIGHTED_NEIGHBOR_SAMPLE, TensorFrame
  from torch_geometric.utils import is_undirected, sort_edge_index, to_torch_csr_tensor, to_undirected
label_desc: Use image node_id 8 for calling a global function with example usage: test_custom_hetero_neighbor_loader() without return types