code: def test_custom_neighbor_loader():
      # Initialize feature store, graph store, and reference:
      feature_store = MyFeatureStore()
      graph_store = MyGraphStore()
  
      # Set up node features:
      x = torch.arange(100, 300)
      feature_store.put_tensor(
          x, group_name=None, attr_name="x", index=None
      )
  
      y = torch.arange(100, 300)
      feature_store.put_tensor(
          y, group_name=None, attr_name="y", index=None
      )
  
      # COO:
      edge_index = get_random_edge_index(100, 100, 500, coalesce=True)
      edge_index = edge_index[:, torch.randperm(edge_index.size(1))]
      coo = (edge_index[0], edge_index[1])
      graph_store.put_edge_index(
          edge_index=coo, edge_type=None, layout="coo", size=(100, 100)
      )
  
      data = Data(x=x, edge_index=edge_index, y=y, num_nodes=200)
  
      # Construct neighbor loaders:
      loader1 = NeighborLoader(
          data,
          batch_size=20,
          input_nodes=torch.arange(100),
          num_neighbors=[-1] * 2,
      )
  
      loader2 = NeighborLoader(
          (feature_store, graph_store),
          batch_size=20,
          input_nodes=torch.arange(100),
          num_neighbors=[-1] * 2,
      )
  
      assert str(loader1) == str(loader2)
      assert len(loader1) == len(loader2)
  
      for batch1, batch2 in zip(loader1, loader2):
          assert len(batch1) == len(batch2)
          assert batch1.num_nodes == batch2.num_nodes
          assert batch1.num_edges == batch2.num_edges
          assert batch1.batch_size == batch2.batch_size
  
          # Mapped indices of neighbors may be differently sorted ...
          assert torch.allclose(batch1.x.sort()[0], batch2.x.sort()[0])
          assert torch.allclose(batch1.y.sort()[0], batch2.y.sort()[0])
  
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 7 for calling a global function with example usage: test_custom_neighbor_loader() without return types