class_docstr: An implementation of an in-memory (heterogeneous) neighbor sampler used
  by :class:`~torch_geometric.loader.NeighborLoader`.
class_docstr_tok: ["An","implementation","of","an","in-memory","(","heterogeneous",")","neighbor","sampler","used","by",":","class",":","`","~torch_geometric.loader.NeighborLoader","`","."]
code: def __init__(
      self,
      data: Union[Data, HeteroData, Tuple[FeatureStore, GraphStore]],
      num_neighbors: NumNeighborsType,
      subgraph_type: Union[SubgraphType, str] = "directional",
      replace: bool = False,
      disjoint: bool = False,
      temporal_strategy: str = "uniform",
      time_attr: Optional[str] = None,
      weight_attr: Optional[str] = None,
      is_sorted: bool = False,
      share_memory: bool = False,
      # Deprecated:
      directed: bool = True,
  ):
      if not directed:
          subgraph_type = SubgraphType.induced
          warnings.warn(
              f"The usage of the 'directed' argument in "
              f"'{self.__class__.__name__}' is deprecated. Use "
              f"`subgraph_type='induced'` instead."
          )
  
      if (
          not torch_geometric.typing.WITH_PYG_LIB
          and sys.platform == "linux"
          and subgraph_type != SubgraphType.induced
      ):
          warnings.warn(
              f"Using '{self.__class__.__name__}' without a "
              f"'pyg-lib' installation is deprecated and will be "
              f"removed soon. Please install 'pyg-lib' for "
              f"accelerated neighborhood sampling"
          )
  
      self.data_type = DataType.from_data(data)
  
      if self.data_type == DataType.homogeneous:
          self.num_nodes = data.num_nodes
  
          self.node_time: Optional[Tensor] = None
          self.edge_time: Optional[Tensor] = None
  
          if time_attr is not None:
              if data.is_node_attr(time_attr):
                  self.node_time = data[time_attr]
              elif data.is_edge_attr(time_attr):
                  self.edge_time = data[time_attr]
              else:
                  raise ValueError(
                      f"The time attribute '{time_attr}' is neither a "
                      f"node-level or edge-level attribute"
                  )
  
          # Convert the graph data into CSC format for sampling:
          self.colptr, self.row, self.perm = to_csc(
              data,
              device="cpu",
              share_memory=share_memory,
              is_sorted=is_sorted,
              src_node_time=self.node_time,
              edge_time=self.edge_time,
          )
  
          if self.edge_time is not None and self.perm is not None:
              self.edge_time = self.edge_time[self.perm]
  
          self.edge_weight: Optional[Tensor] = None
          if weight_attr is not None:
              self.edge_weight = data[weight_attr]
              if self.perm is not None:
                  self.edge_weight = self.edge_weight[self.perm]
  
      elif self.data_type == DataType.heterogeneous:
          self.node_types, self.edge_types = data.metadata()
  
          self.num_nodes = {
              k: data[k].num_nodes for k in self.node_types
          }
  
          self.node_time: Optional[Dict[NodeType, Tensor]] = None
          self.edge_time: Optional[Dict[EdgeType, Tensor]] = None
  
          if time_attr is not None:
              is_node_level_time = is_edge_level_time = False
  
              for store in data.node_stores:
                  if time_attr in store:
                      is_node_level_time = True
              for store in data.edge_stores:
                  if time_attr in store:
                      is_edge_level_time = True
  
              if is_node_level_time and is_edge_level_time:
                  raise ValueError(
                      f"The time attribute '{time_attr}' holds both "
                      f"node-level and edge-level information"
                  )
  
              if not is_node_level_time and not is_edge_level_time:
                  raise ValueError(
                      f"The time attribute '{time_attr}' is neither a "
                      f"node-level or edge-level attribute"
                  )
  
              if is_node_level_time:
                  self.node_time = data.collect(time_attr)
              else:
                  self.edge_time = data.collect(time_attr)
  
          # Conversion to/from C++ string type: Since C++ cannot take
          # dictionaries with tuples as key as input, edge type triplets need
          # to be converted into single strings.
          self.to_rel_type = {k: "__".join(k) for k in self.edge_types}
          self.to_edge_type = {
              v: k for k, v in self.to_rel_type.items()
          }
  
          # Convert the graph data into CSC format for sampling:
          colptr_dict, row_dict, self.perm = to_hetero_csc(
              data,
              device="cpu",
              share_memory=share_memory,
              is_sorted=is_sorted,
              node_time_dict=self.node_time,
              edge_time_dict=self.edge_time,
          )
  
          self.row_dict = remap_keys(row_dict, self.to_rel_type)
          self.colptr_dict = remap_keys(colptr_dict, self.to_rel_type)
  
          if self.edge_time is not None:
              for edge_type, edge_time in self.edge_time.items():
                  if self.perm.get(edge_type, None) is not None:
                      edge_time = edge_time[self.perm[edge_type]]
                      self.edge_time[edge_type] = edge_time
              self.edge_time = remap_keys(
                  self.edge_time, self.to_rel_type
              )
  
          self.edge_weight: Optional[Dict[EdgeType, Tensor]] = None
          if weight_attr is not None:
              self.edge_weight = data.collect(weight_attr)
              for edge_type, edge_weight in self.edge_weight.items():
                  if self.perm.get(edge_type, None) is not None:
                      edge_weight = edge_weight[self.perm[edge_type]]
                      self.edge_weight[edge_type] = edge_weight
              self.edge_weight = remap_keys(
                  self.edge_weight, self.to_rel_type
              )
  
      else:  # self.data_type == DataType.remote
          feature_store, graph_store = data
  
          # Obtain graph metadata:
          attrs = [
              attr for attr in feature_store.get_all_tensor_attrs()
          ]
  
          edge_attrs = graph_store.get_all_edge_attrs()
          self.edge_types = list(
              set(attr.edge_type for attr in edge_attrs)
          )
  
          if weight_attr is not None:
              raise NotImplementedError(
                  f"'weight_attr' argument not yet supported within "
                  f"'{self.__class__.__name__}' for "
                  f"'(FeatureStore, GraphStore)' inputs"
              )
  
          if time_attr is not None:
              # If the `time_attr` is present, we expect that `GraphStore`
              # holds all edges sorted by destination, and within local
              # neighborhoods, node indices should be sorted by time.
              # TODO (matthias, manan) Find an alternative way to ensure.
              for edge_attr in edge_attrs:
                  if edge_attr.layout == EdgeLayout.CSR:
                      raise ValueError(
                          "Temporal sampling requires that edges are stored "
                          "in either COO or CSC layout"
                      )
                  if not edge_attr.is_sorted:
                      raise ValueError(
                          "Temporal sampling requires that edges are "
                          "sorted by destination, and by source time "
                          "within local neighborhoods"
                      )
  
              # We obtain all features with `node_attr.name=time_attr`:
              time_attrs = [
                  copy.copy(attr)
                  for attr in attrs
                  if attr.attr_name == time_attr
              ]
  
          if not self.is_hetero:
              self.node_types = [None]
              self.num_nodes = max(edge_attrs[0].size)
              self.edge_weight: Optional[Tensor] = None
  
              self.node_time: Optional[Tensor] = None
              self.edge_time: Optional[Tensor] = None
  
              if time_attr is not None:
                  if len(time_attrs) != 1:
                      raise ValueError(
                          "Temporal sampling specified but did "
                          "not find any temporal data"
                      )
                  time_attrs[
                      0
                  ].index = None  # Reset index for full data.
                  time_tensor = feature_store.get_tensor(time_attrs[0])
                  # Currently, we determine whether to use node-level or
                  # edge-level temporal sampling based on the attribute name.
                  if time_attr == "time":
                      self.node_time = time_tensor
                  else:
                      self.edge_time = time_tensor
  
              self.row, self.colptr, self.perm = graph_store.csc()
  
          else:
              node_types = [
                  attr.group_name
                  for attr in attrs
                  if isinstance(attr.group_name, str)
              ]
              self.node_types = list(set(node_types))
              self.num_nodes = {
                  node_type: remote_backend_utils.size(*data, node_type)
                  for node_type in self.node_types
              }
              self.edge_weight: Optional[Dict[EdgeType, Tensor]] = None
  
              self.node_time: Optional[Dict[NodeType, Tensor]] = None
              self.edge_time: Optional[Dict[EdgeType, Tensor]] = None
  
              if time_attr is not None:
                  for attr in time_attrs:  # Reset index for full data.
                      attr.index = None
  
                  time_tensors = feature_store.multi_get_tensor(
                      time_attrs
                  )
                  time = {
                      attr.group_name: time_tensor
                      for attr, time_tensor in zip(
                          time_attrs, time_tensors
                      )
                  }
  
                  group_names = [attr.group_name for attr in time_attrs]
                  if all([isinstance(g, str) for g in group_names]):
                      self.node_time = time
                  elif all([isinstance(g, tuple) for g in group_names]):
                      self.edge_time = time
                  else:
                      raise ValueError(
                          f"Found time attribute '{time_attr}' for both "
                          f"node-level and edge-level types"
                      )
  
              # Conversion to/from C++ string type (see above):
              self.to_rel_type = {
                  k: "__".join(k) for k in self.edge_types
              }
              self.to_edge_type = {
                  v: k for k, v in self.to_rel_type.items()
              }
              # Convert the graph data into CSC format for sampling:
              row_dict, colptr_dict, self.perm = graph_store.csc()
              self.row_dict = remap_keys(row_dict, self.to_rel_type)
              self.colptr_dict = remap_keys(
                  colptr_dict, self.to_rel_type
              )
  
      if (
          self.edge_time is not None
          and not torch_geometric.typing.WITH_EDGE_TIME_NEIGHBOR_SAMPLE
      ):
          raise ImportError(
              "Edge-level temporal sampling requires a "
              "more recent 'pyg-lib' installation"
          )
  
      if (
          self.edge_weight is not None
          and not torch_geometric.typing.WITH_WEIGHTED_NEIGHBOR_SAMPLE
      ):
          raise ImportError(
              "Weighted neighbor sampling requires " "'pyg-lib>=0.3.0'"
          )
  
      self.num_neighbors = num_neighbors
      self.replace = replace
      self.subgraph_type = SubgraphType(subgraph_type)
      self.disjoint = disjoint
      self.temporal_strategy = temporal_strategy
  
imports: import copy
  import math
  import sys
  import warnings
  from typing import Callable, Dict, List, Optional, Tuple, Union
  import torch
  from torch import Tensor
  import torch_geometric.typing
  from torch_geometric.data import Data, FeatureStore, GraphStore, HeteroData, remote_backend_utils
  from torch_geometric.data.graph_store import EdgeLayout
  from torch_geometric.sampler import BaseSampler, EdgeSamplerInput, HeteroSamplerOutput, NegativeSampling, NodeSamplerInput, SamplerOutput
  from torch_geometric.sampler.base import DataType, NumNeighbors, SubgraphType
  from torch_geometric.sampler.utils import remap_keys, to_csc, to_hetero_csc
  from torch_geometric.typing import EdgeType, NodeType, OptTensor
label_desc: Use image node_id 1 to create a new NeighborSampler object from inherited base classes: BaseSampler with example: obj = NeighborSampler(data, num_neighbors, subgraph_type, replace, disjoint, temporal_strategy, time_attr, weight_attr, is_sorted, share_memory, directed)