class_docstr: An abstract class for implementing basic GNN models.
  
  Args:
      in_channels (int or tuple): Size of each input sample, or :obj:`-1` to
          derive the size from the first input(s) to the forward method.
          A tuple corresponds to the sizes of source and target
          dimensionalities.
      hidden_channels (int): Size of each hidden sample.
      num_layers (int): Number of message passing layers.
      out_channels (int, optional): If not set to :obj:`None`, will apply a
          final linear transformation to convert hidden node embeddings to
          output size :obj:`out_channels`. (default: :obj:`None`)
      dropout (float, optional): Dropout probability. (default: :obj:`0.`)
      act (str or Callable, optional): The non-linear activation function to
          use. (default: :obj:`"relu"`)
      act_first (bool, optional): If set to :obj:`True`, activation is
          applied before normalization. (default: :obj:`False`)
      act_kwargs (Dict[str, Any], optional): Arguments passed to the
          respective activation function defined by :obj:`act`.
          (default: :obj:`None`)
      norm (str or Callable, optional): The normalization function to
          use. (default: :obj:`None`)
      norm_kwargs (Dict[str, Any], optional): Arguments passed to the
          respective normalization function defined by :obj:`norm`.
          (default: :obj:`None`)
      jk (str, optional): The Jumping Knowledge mode. If specified, the model
          will additionally apply a final linear transformation to transform
          node embeddings to the expected output feature dimensionality.
          (:obj:`None`, :obj:`"last"`, :obj:`"cat"`, :obj:`"max"`,
          :obj:`"lstm"`). (default: :obj:`None`)
      **kwargs (optional): Additional arguments of the underlying
          :class:`torch_geometric.nn.conv.MessagePassing` layers.
docstr: Performs layer-wise inference on large-graphs using a
  :class:`~torch_geometric.loader.NeighborLoader`, where
  :class:`~torch_geometric.loader.NeighborLoader` should sample the
  full neighborhood for only one layer.
  This is an efficient way to compute the output embeddings for all
  nodes in the graph.
  Only applicable in case :obj:`jk=None` or `jk='last'`.
  
  Args:
      loader (torch_geometric.loader.NeighborLoader): A neighbor loader
          object that generates full 1-hop subgraphs, *i.e.*,
          :obj:`loader.num_neighbors = [-1]`.
      device (torch.device, optional): The device to run the GNN on.
          (default: :obj:`None`)
      embedding_device (torch.device, optional): The device to store
          intermediate embeddings on. If intermediate embeddings fit on
          GPU, this option helps to avoid unnecessary device transfers.
          (default: :obj:`"cpu"`)
      progress_bar (bool, optional): If set to :obj:`True`, will print a
          progress bar during computation. (default: :obj:`False`)
      cache (bool, optional): If set to :obj:`True`, caches intermediate
          sampler outputs for usage in later epochs.
          This will avoid repeated sampling to accelerate inference.
          (default: :obj:`False`)
docstr_tok: ["Performs","layer-wise","inference","on","large-graphs","using","a",":","class",":","`","~torch_geometric.loader.NeighborLoader","`",",","where",":","class",":","`","~torch_geometric.loader.NeighborLoader","`","should","sample","the","full","neighborhood","for","only","one","layer",".","This","is","an","efficient","way","to","compute","the","output","embeddings","for","all","nodes","in","the","graph",".","Only","applicable","in","case",":","obj",":","`","jk=None","`","or","`","jk='last","'","`",".","Args",":","loader","(","torch_geometric.loader.NeighborLoader",")",":","A","neighbor","loader","object","that","generates","full","1-hop","subgraphs",",","*","i.e",".","*",",",":","obj",":","`","loader.num_neighbors","=","[","-1","]","`",".","device","(","torch.device",",","optional",")",":","The","device","to","run","the","GNN","on",".","(","default",":",":","obj",":","`","None","`",")","embedding_device","(","torch.device",",","optional",")",":","The","device","to","store","intermediate","embeddings","on",".","If","intermediate","embeddings","fit","on","GPU",",","this","option","helps","to","avoid","unnecessary","device","transfers",".","(","default",":",":","obj",":","`","``","cpu","''","`",")","progress_bar","(","bool",",","optional",")",":","If","set","to",":","obj",":","`","True","`",",","will","print","a","progress","bar","during","computation",".","(","default",":",":","obj",":","`","False","`",")","cache","(","bool",",","optional",")",":","If","set","to",":","obj",":","`","True","`",",","caches","intermediate","sampler","outputs","for","usage","in","later","epochs",".","This","will","avoid","repeated","sampling","to","accelerate","inference",".","(","default",":",":","obj",":","`","False","`",")"]
code: def inference(
      self,
      loader: NeighborLoader,
      device: Optional[Union[str, torch.device]] = None,
      embedding_device: Union[str, torch.device] = "cpu",
      progress_bar: bool = False,
      cache: bool = False,
  ) -> Tensor:
      r"""Performs layer-wise inference on large-graphs using a
      :class:`~torch_geometric.loader.NeighborLoader`, where
      :class:`~torch_geometric.loader.NeighborLoader` should sample the
      full neighborhood for only one layer.
      This is an efficient way to compute the output embeddings for all
      nodes in the graph.
      Only applicable in case :obj:`jk=None` or `jk='last'`.
  
      Args:
          loader (torch_geometric.loader.NeighborLoader): A neighbor loader
              object that generates full 1-hop subgraphs, *i.e.*,
              :obj:`loader.num_neighbors = [-1]`.
          device (torch.device, optional): The device to run the GNN on.
              (default: :obj:`None`)
          embedding_device (torch.device, optional): The device to store
              intermediate embeddings on. If intermediate embeddings fit on
              GPU, this option helps to avoid unnecessary device transfers.
              (default: :obj:`"cpu"`)
          progress_bar (bool, optional): If set to :obj:`True`, will print a
              progress bar during computation. (default: :obj:`False`)
          cache (bool, optional): If set to :obj:`True`, caches intermediate
              sampler outputs for usage in later epochs.
              This will avoid repeated sampling to accelerate inference.
              (default: :obj:`False`)
      """
      assert self.jk_mode is None or self.jk_mode == "last"
      assert isinstance(loader, NeighborLoader)
      assert len(loader.dataset) == loader.data.num_nodes
      assert len(loader.node_sampler.num_neighbors) == 1
      assert not self.training
      # assert not loader.shuffle  # TODO (matthias) does not work :(
      if progress_bar:
          pbar = tqdm(total=len(self.convs) * len(loader))
          pbar.set_description("Inference")
  
      x_all = loader.data.x.to(embedding_device)
  
      if cache:
          # Only cache necessary attributes:
          def transform(data: Data) -> Data:
              kwargs = dict(n_id=data.n_id, batch_size=data.batch_size)
              if hasattr(data, "adj_t"):
                  kwargs["adj_t"] = data.adj_t
              else:
                  kwargs["edge_index"] = data.edge_index
  
              return Data.from_dict(kwargs)
  
          loader = CachedLoader(
              loader, device=device, transform=transform
          )
  
      for i in range(self.num_layers):
          xs: List[Tensor] = []
          for batch in loader:
              x = x_all[batch.n_id].to(device)
              batch_size = batch.batch_size
              if hasattr(batch, "adj_t"):
                  edge_index = batch.adj_t.to(device)
              else:
                  edge_index = batch.edge_index.to(device)
  
              x = self.inference_per_layer(i, x, edge_index, batch_size)
              xs.append(x.to(embedding_device))
  
              if progress_bar:
                  pbar.update(1)
  
          x_all = torch.cat(xs, dim=0)
  
      if progress_bar:
          pbar.close()
  
      return x_all
  
imports: import copy
  import inspect
  from typing import Any, Callable, Dict, Final, List, Optional, Tuple, Union
  import torch
  from torch import Tensor
  from torch.nn import Linear, ModuleList
  from tqdm import tqdm
  from torch_geometric.data import Data
  from torch_geometric.loader import CachedLoader, NeighborLoader
  from torch_geometric.nn.conv import EdgeConv, GATConv, GATv2Conv, GCNConv, GINConv, MessagePassing, PNAConv, SAGEConv
  from torch_geometric.nn.models import MLP
  from torch_geometric.nn.models.jumping_knowledge import JumpingKnowledge
  from torch_geometric.nn.resolver import activation_resolver, normalization_resolver
  from torch_geometric.typing import Adj, OptTensor
  from torch_geometric.utils._trim_to_layer import TrimToLayer
label_desc: Use image node_id 6 for calling the BasicGNN obj's underlying member method code with example usage: obj.inference(loader, device, embedding_device, progress_bar, cache) and returns: x_all, Data