class_docstr: An abstract class for implementing scalable GNNs via historical
  embeddings.
  This class will take care of initializing :obj:`num_layers - 1` historical
  embeddings, and provides a convenient interface to push recent node
  embeddings to the history, and to pull previous embeddings from the
  history.
  In case historical embeddings are stored on the CPU, they will reside
  inside pinned memory, which allows for asynchronous memory transfers of
  historical embeddings.
  For this, this class maintains a :class:`AsyncIOPool` object that
  implements the underlying mechanisms of asynchronous memory transfers as
  described in our paper.
  
  Args:
      num_nodes (int): The number of nodes in the graph.
      hidden_channels (int): The number of hidden channels of the model.
          As a current restriction, all intermediate node embeddings need to
          utilize the same number of features.
      num_layers (int): The number of layers of the model.
      pool_size (int, optional): The number of pinned CPU buffers for pulling
          histories and transfering them to GPU.
          Needs to be set in order to make use of asynchronous memory
          transfers. (default: :obj:`None`)
      buffer_size (int, optional): The size of pinned CPU buffers, i.e. the
          maximum number of out-of-mini-batch nodes pulled at once.
          Needs to be set in order to make use of asynchronous memory
          transfers. (default: :obj:`None`)
class_docstr_tok: ["An","abstract","class","for","implementing","scalable","GNNs","via","historical","embeddings",".","This","class","will","take","care","of","initializing",":","obj",":","`","num_layers","-","1","`","historical","embeddings",",","and","provides","a","convenient","interface","to","push","recent","node","embeddings","to","the","history",",","and","to","pull","previous","embeddings","from","the","history",".","In","case","historical","embeddings","are","stored","on","the","CPU",",","they","will","reside","inside","pinned","memory",",","which","allows","for","asynchronous","memory","transfers","of","historical","embeddings",".","For","this",",","this","class","maintains","a",":","class",":","`","AsyncIOPool","`","object","that","implements","the","underlying","mechanisms","of","asynchronous","memory","transfers","as","described","in","our","paper",".","Args",":","num_nodes","(","int",")",":","The","number","of","nodes","in","the","graph",".","hidden_channels","(","int",")",":","The","number","of","hidden","channels","of","the","model",".","As","a","current","restriction",",","all","intermediate","node","embeddings","need","to","utilize","the","same","number","of","features",".","num_layers","(","int",")",":","The","number","of","layers","of","the","model",".","pool_size","(","int",",","optional",")",":","The","number","of","pinned","CPU","buffers","for","pulling","histories","and","transfering","them","to","GPU",".","Needs","to","be","set","in","order","to","make","use","of","asynchronous","memory","transfers",".","(","default",":",":","obj",":","`","None","`",")","buffer_size","(","int",",","optional",")",":","The","size","of","pinned","CPU","buffers",",","i.e",".","the","maximum","number","of","out-of-mini-batch","nodes","pulled","at","once",".","Needs","to","be","set","in","order","to","make","use","of","asynchronous","memory","transfers",".","(","default",":",":","obj",":","`","None","`",")"]
docstr: Enhances the call of forward propagation by immediately start
  pulling historical embeddings for all layers asynchronously.
  After forward propogation is completed, the push of node embeddings to
  the histories will be synchronized.
  
  For example, given a mini-batch with node indices
  :obj:`n_id = [0, 1, 5, 6, 7, 3, 4]`, where the first 5 nodes
  represent the mini-batched nodes, and nodes :obj:`3` and :obj:`4`
  denote out-of-mini-batched nodes (i.e. the 1-hop neighbors of the
  mini-batch that are not included in the current mini-batch), then
  other input arguments should be given as:
  
  .. code-block:: python
  
      batch_size = 5
      offset = [0, 2, 5]
      count = [2, 3]
  
  Args:
      x (Tensor, optional): Node feature matrix. (default: :obj:`None`)
      adj_t (SparseTensor, optional) The sparse adjacency matrix.
          (default: :obj:`None`)
      batch_size (int, optional): The in-mini-batch size of nodes.
          (default: :obj:`None`)
      n_id (Tensor, optional): The global indices of mini-batched and
          out-of-mini-batched nodes. (default: :obj:`None`)
      offset (Tensor, optional): The offset of mini-batched nodes inside
          a utilize a contiguous memory layout. (default: :obj:`None`)
      count (Tensor, optional): The number of mini-batched nodes inside a
          contiguous memory layout. (default: :obj:`None`)
      loader (EvalSubgraphLoader, optional): A subgraph loader used for
          evaluating the given GNN in a layer-wise fashsion.
docstr_tok: ["Enhances","the","call","of","forward","propagation","by","immediately","start","pulling","historical","embeddings","for","all","layers","asynchronously",".","After","forward","propogation","is","completed",",","the","push","of","node","embeddings","to","the","histories","will","be","synchronized",".","For","example",",","given","a","mini-batch","with","node","indices",":","obj",":","`","n_id","=","[","0",",","1",",","5",",","6",",","7",",","3",",","4","]","`",",","where","the","first","5","nodes","represent","the","mini-batched","nodes",",","and","nodes",":","obj",":","`","3","`","and",":","obj",":","`","4","`","denote","out-of-mini-batched","nodes","(","i.e",".","the","1-hop","neighbors","of","the","mini-batch","that","are","not","included","in","the","current","mini-batch",")",",","then","other","input","arguments","should","be","given","as",":","..","code-block",":",":","python","batch_size","=","5","offset","=","[","0",",","2",",","5","]","count","=","[","2",",","3","]","Args",":","x","(","Tensor",",","optional",")",":","Node","feature","matrix",".","(","default",":",":","obj",":","`","None","`",")","adj_t","(","SparseTensor",",","optional",")","The","sparse","adjacency","matrix",".","(","default",":",":","obj",":","`","None","`",")","batch_size","(","int",",","optional",")",":","The","in-mini-batch","size","of","nodes",".","(","default",":",":","obj",":","`","None","`",")","n_id","(","Tensor",",","optional",")",":","The","global","indices","of","mini-batched","and","out-of-mini-batched","nodes",".","(","default",":",":","obj",":","`","None","`",")","offset","(","Tensor",",","optional",")",":","The","offset","of","mini-batched","nodes","inside","a","utilize","a","contiguous","memory","layout",".","(","default",":",":","obj",":","`","None","`",")","count","(","Tensor",",","optional",")",":","The","number","of","mini-batched","nodes","inside","a","contiguous","memory","layout",".","(","default",":",":","obj",":","`","None","`",")","loader","(","EvalSubgraphLoader",",","optional",")",":","A","subgraph","loader","used","for","evaluating","the","given","GNN","in","a","layer-wise","fashsion","."]
code: def __call__(
      self,
      x: Optional[Tensor] = None,
      adj_t: Optional[SparseTensor] = None,
      batch_size: Optional[int] = None,
      n_id: Optional[Tensor] = None,
      offset: Optional[Tensor] = None,
      count: Optional[Tensor] = None,
      loader: EvalSubgraphLoader = None,
      **kwargs,
  ) -> Tensor:
      r"""Enhances the call of forward propagation by immediately start
      pulling historical embeddings for all layers asynchronously.
      After forward propogation is completed, the push of node embeddings to
      the histories will be synchronized.
  
      For example, given a mini-batch with node indices
      :obj:`n_id = [0, 1, 5, 6, 7, 3, 4]`, where the first 5 nodes
      represent the mini-batched nodes, and nodes :obj:`3` and :obj:`4`
      denote out-of-mini-batched nodes (i.e. the 1-hop neighbors of the
      mini-batch that are not included in the current mini-batch), then
      other input arguments should be given as:
  
      .. code-block:: python
  
          batch_size = 5
          offset = [0, 2, 5]
          count = [2, 3]
  
      Args:
          x (Tensor, optional): Node feature matrix. (default: :obj:`None`)
          adj_t (SparseTensor, optional) The sparse adjacency matrix.
              (default: :obj:`None`)
          batch_size (int, optional): The in-mini-batch size of nodes.
              (default: :obj:`None`)
          n_id (Tensor, optional): The global indices of mini-batched and
              out-of-mini-batched nodes. (default: :obj:`None`)
          offset (Tensor, optional): The offset of mini-batched nodes inside
              a utilize a contiguous memory layout. (default: :obj:`None`)
          count (Tensor, optional): The number of mini-batched nodes inside a
              contiguous memory layout. (default: :obj:`None`)
          loader (EvalSubgraphLoader, optional): A subgraph loader used for
              evaluating the given GNN in a layer-wise fashsion.
      """
  
      if loader is not None:
          return self.mini_inference(loader)
  
      # We only perform asynchronous history transfer in case the following
      # conditions are met:
      self._async = (
          self.pool is not None
          and batch_size is not None
          and n_id is not None
          and offset is not None
          and count is not None
      )
  
      if self._async:
          for hist in self.histories:
              self.pool.async_pull(
                  hist.emb, None, None, n_id[batch_size:]
              )
  
      out = self.forward(
          x, adj_t, batch_size, n_id, offset, count, **kwargs
      )
  
      if self._async:
          for hist in self.histories:
              self.pool.synchronize_push()
  
      self._async = False
  
      return out
  
file_path: DIG/dig/lsgraph/method/GraphFMOB/models/base.py
imports: from typing import Optional, Callable, Dict, Any
  import torch
  from torch import Tensor
  from torch_sparse import SparseTensor
  from dig.lsgraph.method.FM import FeatureMomentum
  from ..pool import AsyncIOPool
  from ..loader import SubgraphLoader, EvalSubgraphLoader
label_desc: Use image node_id 6 for calling the ScalableGNN obj's underlying member method code with example usage: obj.__call__(x, adj_t, batch_size, n_id, offset, count, loader) and returns: out, self