code: def trim_adj(
      edge_index: Adj,
      layer: int,
      num_sampled_src_nodes_per_hop: List[int],
      num_sampled_dst_nodes_per_hop: List[int],
      num_sampled_edges_per_hop: List[int],
  ) -> Adj:
      if layer <= 0:
          return edge_index
  
      if isinstance(edge_index, Tensor):
          return edge_index.narrow(
              dim=1,
              start=0,
              length=edge_index.size(1)
              - num_sampled_edges_per_hop[-layer],
          )
  
      elif isinstance(edge_index, SparseTensor):
          size = (
              edge_index.size(0)
              - num_sampled_dst_nodes_per_hop[-layer],
              edge_index.size(1)
              - num_sampled_src_nodes_per_hop[-layer],
          )
  
          num_seed_nodes = (
              size[0] - num_sampled_dst_nodes_per_hop[-(layer + 1)]
          )
  
          return trim_sparse_tensor(edge_index, size, num_seed_nodes)
  
      raise ValueError(
          f"Unsupported 'edge_index' type '{type(edge_index)}'"
      )
  
file_path: pytorch_geometric/torch_geometric/utils/_trim_to_layer.py
imports: import copy
  from typing import Any, Dict, List, Optional, Tuple, Union, overload
  import torch
  from torch import Tensor
  from torch_geometric.typing import Adj, EdgeType, MaybeHeteroAdjTensor, MaybeHeteroEdgeTensor, MaybeHeteroNodeTensor, NodeType, SparseStorage, SparseTensor
label_desc: Use image node_id 6 for calling a global function with example usage: trim_adj(edge_index, layer, num_sampled_src_nodes_per_hop, num_sampled_dst_nodes_per_hop, num_sampled_edges_per_hop) and returns: edge_index, edge_index, trim_sparse_tensor