code: def trim_to_layer(
      layer: int,
      num_sampled_nodes_per_hop: Union[
          List[int], Dict[NodeType, List[int]]
      ],
      num_sampled_edges_per_hop: Union[
          List[int], Dict[EdgeType, List[int]]
      ],
      x: MaybeHeteroNodeTensor,
      edge_index: MaybeHeteroEdgeTensor,
      edge_attr: Optional[MaybeHeteroEdgeTensor] = None,
  ) -> Tuple[
      MaybeHeteroNodeTensor,
      MaybeHeteroAdjTensor,
      Optional[MaybeHeteroEdgeTensor],
  ]:
      r"""Trims the :obj:`edge_index` representation, node features :obj:`x` and
      edge features :obj:`edge_attr` to a minimal-sized representation for the
      current GNN layer :obj:`layer` in directed
      :class:`~torch_geometric.loader.NeighborLoader` scenarios.
  
      This ensures that no computation is performed for nodes and edges that are
      not included in the current GNN layer, thus avoiding unnecessary
      computation within the GNN when performing neighborhood sampling.
  
      Args:
          layer (int): The current GNN layer.
          num_sampled_nodes_per_hop (List[int] or Dict[NodeType, List[int]]): The
              number of sampled nodes per hop.
          num_sampled_edges_per_hop (List[int] or Dict[EdgeType, List[int]]): The
              number of sampled edges per hop.
          x (torch.Tensor or Dict[NodeType, torch.Tensor]): The homogeneous or
              heterogeneous (hidden) node features.
          edge_index (torch.Tensor or Dict[EdgeType, torch.Tensor]): The
              homogeneous or heterogeneous edge indices.
          edge_attr (torch.Tensor or Dict[EdgeType, torch.Tensor], optional): The
              homogeneous or heterogeneous (hidden) edge features.
      """
      if layer <= 0:
          return x, edge_index, edge_attr
  
      if isinstance(num_sampled_edges_per_hop, dict):
          assert isinstance(num_sampled_nodes_per_hop, dict)
  
          assert isinstance(x, dict)
          x = {
              k: trim_feat(v, layer, num_sampled_nodes_per_hop[k])
              for k, v in x.items()
          }
          x = filter_empty_entries(x)
  
          assert isinstance(edge_index, dict)
          edge_index = {
              k: trim_adj(
                  v,
                  layer,
                  num_sampled_nodes_per_hop[k[0]],
                  num_sampled_nodes_per_hop[k[-1]],
                  num_sampled_edges_per_hop[k],
              )
              for k, v in edge_index.items()
          }
          edge_index = filter_empty_entries(edge_index)
  
          if edge_attr is not None:
              assert isinstance(edge_attr, dict)
              edge_attr = {
                  k: trim_feat(v, layer, num_sampled_edges_per_hop[k])
                  for k, v in edge_attr.items()
              }
              edge_attr = filter_empty_entries(edge_attr)
  
          return x, edge_index, edge_attr
  
      assert isinstance(num_sampled_nodes_per_hop, list)
  
      assert isinstance(x, Tensor)
      x = trim_feat(x, layer, num_sampled_nodes_per_hop)
  
      assert isinstance(edge_index, (Tensor, SparseTensor))
      edge_index = trim_adj(
          edge_index,
          layer,
          num_sampled_nodes_per_hop,
          num_sampled_nodes_per_hop,
          num_sampled_edges_per_hop,
      )
  
      if edge_attr is not None:
          assert isinstance(edge_attr, Tensor)
          edge_attr = trim_feat(
              edge_attr, layer, num_sampled_edges_per_hop
          )
  
      return x, edge_index, edge_attr
  
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 4 for calling a global function with example usage: trim_to_layer(layer, num_sampled_nodes_per_hop, num_sampled_edges_per_hop, x, edge_index, edge_attr) and returns: x,  edge_index,  edge_attr, x,  edge_index,  edge_attr, x,  edge_index,  edge_attr