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","`","."]
docstr: Implements neighbor sampling by calling either :obj:`pyg-lib` (if
  installed) or :obj:`torch-sparse` (if installed) sampling routines.
docstr_tok: ["Implements","neighbor","sampling","by","calling","either",":","obj",":","`","pyg-lib","`","(","if","installed",")","or",":","obj",":","`","torch-sparse","`","(","if","installed",")","sampling","routines","."]
code: def _sample(
      self,
      seed: Union[Tensor, Dict[NodeType, Tensor]],
      seed_time: Optional[Union[Tensor, Dict[NodeType, Tensor]]] = None,
      **kwargs,
  ) -> Union[SamplerOutput, HeteroSamplerOutput]:
      r"""Implements neighbor sampling by calling either :obj:`pyg-lib` (if
      installed) or :obj:`torch-sparse` (if installed) sampling routines.
      """
      if isinstance(seed, dict):  # Heterogeneous sampling:
          # TODO Support induced subgraph sampling in `pyg-lib`.
          if (
              torch_geometric.typing.WITH_PYG_LIB
              and self.subgraph_type != SubgraphType.induced
          ):
              # TODO (matthias) Ideally, `seed` inherits dtype from `colptr`
              colptrs = list(self.colptr_dict.values())
              dtype = (
                  colptrs[0].dtype if len(colptrs) > 0 else torch.int64
              )
              seed = {k: v.to(dtype) for k, v in seed.items()}
  
              args = (
                  self.node_types,
                  self.edge_types,
                  self.colptr_dict,
                  self.row_dict,
                  seed,
                  self.num_neighbors.get_mapped_values(self.edge_types),
                  self.node_time,
              )
              if torch_geometric.typing.WITH_EDGE_TIME_NEIGHBOR_SAMPLE:
                  args += (self.edge_time,)
              args += (seed_time,)
              if torch_geometric.typing.WITH_WEIGHTED_NEIGHBOR_SAMPLE:
                  args += (self.edge_weight,)
              args += (
                  True,  # csc
                  self.replace,
                  self.subgraph_type != SubgraphType.induced,
                  self.disjoint,
                  self.temporal_strategy,
                  # TODO (matthias) `return_edge_id` if edge features present
                  True,  # return_edge_id
              )
  
              out = torch.ops.pyg.hetero_neighbor_sample(*args)
              row, col, node, edge, batch = out[:4] + (None,)
  
              # `pyg-lib>0.1.0` returns sampled number of nodes/edges:
              num_sampled_nodes = num_sampled_edges = None
              if len(out) >= 6:
                  num_sampled_nodes, num_sampled_edges = out[4:6]
  
              if self.disjoint:
                  node = {
                      k: v.t().contiguous() for k, v in node.items()
                  }
                  batch = {k: v[0] for k, v in node.items()}
                  node = {k: v[1] for k, v in node.items()}
  
          elif torch_geometric.typing.WITH_TORCH_SPARSE:
              if self.disjoint:
                  raise ValueError(
                      "'disjoint' sampling not supported for "
                      "neighbor sampling via 'torch-sparse'. "
                      "Please install 'pyg-lib' for improved "
                      "and optimized sampling routines."
                  )
  
              out = torch.ops.torch_sparse.hetero_neighbor_sample(
                  self.node_types,
                  self.edge_types,
                  self.colptr_dict,
                  self.row_dict,
                  seed,  # seed_dict
                  self.num_neighbors.get_mapped_values(self.edge_types),
                  self.num_neighbors.num_hops,
                  self.replace,
                  self.subgraph_type != SubgraphType.induced,
              )
              node, row, col, edge, batch = out + (None,)
              num_sampled_nodes = num_sampled_edges = None
  
          else:
              raise ImportError(
                  f"'{self.__class__.__name__}' requires "
                  f"either 'pyg-lib' or 'torch-sparse'"
              )
  
          if num_sampled_edges is not None:
              num_sampled_edges = remap_keys(
                  num_sampled_edges,
                  self.to_edge_type,
              )
  
          return HeteroSamplerOutput(
              node=node,
              row=remap_keys(row, self.to_edge_type),
              col=remap_keys(col, self.to_edge_type),
              edge=remap_keys(edge, self.to_edge_type),
              batch=batch,
              num_sampled_nodes=num_sampled_nodes,
              num_sampled_edges=num_sampled_edges,
          )
  
      else:  # Homogeneous sampling:
          # TODO Support induced subgraph sampling in `pyg-lib`.
          if (
              torch_geometric.typing.WITH_PYG_LIB
              and self.subgraph_type != SubgraphType.induced
          ):
              args = (
                  self.colptr,
                  self.row,
                  # TODO (matthias) `seed` should inherit dtype from `colptr`
                  seed.to(self.colptr.dtype),
                  self.num_neighbors.get_mapped_values(),
                  self.node_time,
              )
              if torch_geometric.typing.WITH_EDGE_TIME_NEIGHBOR_SAMPLE:
                  args += (self.edge_time,)
              args += (seed_time,)
              if torch_geometric.typing.WITH_WEIGHTED_NEIGHBOR_SAMPLE:
                  args += (self.edge_weight,)
              args += (
                  True,  # csc
                  self.replace,
                  self.subgraph_type != SubgraphType.induced,
                  self.disjoint,
                  self.temporal_strategy,
                  # TODO (matthias) `return_edge_id` if edge features present
                  True,  # return_edge_id
              )
  
              out = torch.ops.pyg.neighbor_sample(*args)
              row, col, node, edge, batch = out[:4] + (None,)
  
              # `pyg-lib>0.1.0` returns sampled number of nodes/edges:
              num_sampled_nodes = num_sampled_edges = None
              if len(out) >= 6:
                  num_sampled_nodes, num_sampled_edges = out[4:6]
  
              if self.disjoint:
                  batch, node = node.t().contiguous()
  
          elif torch_geometric.typing.WITH_TORCH_SPARSE:
              if self.disjoint:
                  raise ValueError(
                      "'disjoint' sampling not supported for "
                      "neighbor sampling via 'torch-sparse'. "
                      "Please install 'pyg-lib' for improved "
                      "and optimized sampling routines."
                  )
  
              out = torch.ops.torch_sparse.neighbor_sample(
                  self.colptr,
                  self.row,
                  seed,  # seed
                  self.num_neighbors.get_mapped_values(),
                  self.replace,
                  self.subgraph_type != SubgraphType.induced,
              )
              node, row, col, edge, batch = out + (None,)
              num_sampled_nodes = num_sampled_edges = None
  
          else:
              raise ImportError(
                  f"'{self.__class__.__name__}' requires "
                  f"either 'pyg-lib' or 'torch-sparse'"
              )
  
          return SamplerOutput(
              node=node,
              row=row,
              col=col,
              edge=edge,
              batch=batch,
              num_sampled_nodes=num_sampled_nodes,
              num_sampled_edges=num_sampled_edges,
          )
  
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 11 for calling the NeighborSampler obj's underlying member method code with example usage: obj._sample(seed, seed_time) and returns: HeteroSamplerOutput, SamplerOutput