code: def test_hgt_loader():
      torch.manual_seed(12345)
  
      data = HeteroData()
  
      data["paper"].x = torch.arange(100)
      data["author"].x = torch.arange(100, 300)
  
      data["paper", "paper"].edge_index = get_random_edge_index(
          100, 100, 500
      )
      data["paper", "paper"].edge_attr = torch.arange(500)
      data["paper", "author"].edge_index = get_random_edge_index(
          100, 200, 1000
      )
      data["paper", "author"].edge_attr = torch.arange(500, 1500)
      data["author", "paper"].edge_index = get_random_edge_index(
          200, 100, 1000
      )
      data["author", "paper"].edge_attr = torch.arange(1500, 2500)
  
      r1, c1 = data["paper", "paper"].edge_index
      r2, c2 = data["paper", "author"].edge_index + torch.tensor(
          [[0], [100]]
      )
      r3, c3 = data["author", "paper"].edge_index + torch.tensor(
          [[100], [0]]
      )
      full_adj = SparseTensor(
          row=torch.cat([r1, r2, r3]),
          col=torch.cat([c1, c2, c3]),
          value=torch.arange(2500),
      )
  
      batch_size = 20
      loader = HGTLoader(
          data,
          num_samples=[5] * 4,
          batch_size=batch_size,
          input_nodes="paper",
      )
      assert str(loader) == "HGTLoader()"
      assert len(loader) == (100 + batch_size - 1) // batch_size
  
      for batch in loader:
          assert isinstance(batch, HeteroData)
  
          # Test node and types:
          assert set(batch.node_types) == {"paper", "author"}
          assert set(batch.edge_types) == set(data.edge_types)
  
          assert len(batch["paper"]) == 4
          assert batch["paper"].n_id.size() == (
              batch["paper"].num_nodes,
          )
          assert batch["paper"].x.size() == (40,)  # 20 + 4 * 5
          assert batch["paper"].input_id.numel() == batch_size
          assert batch["paper"].batch_size == batch_size
          assert (
              batch["paper"].x.min() >= 0
              and batch["paper"].x.max() < 100
          )
  
          assert len(batch["author"]) == 2
          assert batch["author"].n_id.size() == (
              batch["author"].num_nodes,
          )
          assert batch["author"].x.size() == (20,)  # 4 * 5
          assert (
              batch["author"].x.min() >= 100
              and batch["author"].x.max() < 300
          )
  
          # Test edge type selection:
          assert set(batch.edge_types) == {
              ("paper", "to", "paper"),
              ("paper", "to", "author"),
              ("author", "to", "paper"),
          }
  
          assert len(batch["paper", "paper"]) == 3
          num_edges = batch["paper", "paper"].num_edges
          assert batch["paper", "paper"].e_id.size() == (num_edges,)
          row, col = batch["paper", "paper"].edge_index
          value = batch["paper", "paper"].edge_attr
          adj = full_adj[batch["paper"].x, batch["paper"].x]
          assert row.min() >= 0 and row.max() < 40
          assert col.min() >= 0 and col.max() < 40
          assert value.min() >= 0 and value.max() < 500
          assert adj.nnz() == row.size(0)
          assert torch.allclose(
              row.unique(), adj.storage.row().unique()
          )
          assert torch.allclose(
              col.unique(), adj.storage.col().unique()
          )
          assert torch.allclose(
              value.unique(), adj.storage.value().unique()
          )
  
          assert is_subset(
              batch["paper", "paper"].edge_index,
              data["paper", "paper"].edge_index,
              batch["paper"].x,
              batch["paper"].x,
          )
  
          assert len(batch["paper", "author"]) == 3
          num_edges = batch["paper", "author"].num_edges
          assert batch["paper", "author"].e_id.size() == (num_edges,)
          row, col = batch["paper", "author"].edge_index
          value = batch["paper", "author"].edge_attr
          adj = full_adj[batch["paper"].x, batch["author"].x]
          assert row.min() >= 0 and row.max() < 40
          assert col.min() >= 0 and col.max() < 20
          assert value.min() >= 500 and value.max() < 1500
          assert adj.nnz() == row.size(0)
          assert torch.allclose(
              row.unique(), adj.storage.row().unique()
          )
          assert torch.allclose(
              col.unique(), adj.storage.col().unique()
          )
          assert torch.allclose(
              value.unique(), adj.storage.value().unique()
          )
  
          assert is_subset(
              batch["paper", "author"].edge_index,
              data["paper", "author"].edge_index,
              batch["paper"].x,
              batch["author"].x - 100,
          )
  
          assert len(batch["author", "paper"]) == 3
          num_edges = batch["author", "paper"].num_edges
          assert batch["author", "paper"].e_id.size() == (num_edges,)
          row, col = batch["author", "paper"].edge_index
          value = batch["author", "paper"].edge_attr
          adj = full_adj[batch["author"].x, batch["paper"].x]
          assert row.min() >= 0 and row.max() < 20
          assert col.min() >= 0 and col.max() < 40
          assert value.min() >= 1500 and value.max() < 2500
          assert adj.nnz() == row.size(0)
          assert torch.allclose(
              row.unique(), adj.storage.row().unique()
          )
          assert torch.allclose(
              col.unique(), adj.storage.col().unique()
          )
          assert torch.allclose(
              value.unique(), adj.storage.value().unique()
          )
  
          assert is_subset(
              batch["author", "paper"].edge_index,
              data["author", "paper"].edge_index,
              batch["author"].x - 100,
              batch["paper"].x,
          )
  
          # Test for isolated nodes (there shouldn't exist any):
          n_id = torch.cat([batch["paper"].x, batch["author"].x])
          row, col, _ = full_adj[n_id, n_id].coo()
          assert torch.cat([row, col]).unique().numel() >= 59
  
file_path: pytorch_geometric/test/loader/test_hgt_loader.py
imports: import numpy
  import torch
  from torch_geometric.data import HeteroData
  from torch_geometric.loader import HGTLoader
  from torch_geometric.nn import GraphConv, to_hetero
  from torch_geometric.testing import get_random_edge_index, onlyOnline, withPackage
  from torch_geometric.typing import SparseTensor
  from torch_geometric.utils import k_hop_subgraph
label_desc: Use image node_id 2 for calling a global function with example usage: test_hgt_loader() without return types