code: def edge_sample(
      inputs: EdgeSamplerInput,
      sample_fn: Callable,
      num_nodes: Union[int, Dict[NodeType, int]],
      disjoint: bool,
      node_time: Optional[Union[Tensor, Dict[str, Tensor]]] = None,
      neg_sampling: Optional[NegativeSampling] = None,
  ) -> Union[SamplerOutput, HeteroSamplerOutput]:
      r"""Performs sampling from an edge sampler input, leveraging a sampling
      function of the same signature as `node_sample`.
      """
      input_id = inputs.input_id
      src = inputs.row
      dst = inputs.col
      edge_label = inputs.label
      edge_label_time = inputs.time
      input_type = inputs.input_type
  
      src_time = dst_time = edge_label_time
      assert edge_label_time is None or disjoint
  
      assert isinstance(num_nodes, (dict, int))
      if not isinstance(num_nodes, dict):
          num_src_nodes = num_dst_nodes = num_nodes
      else:
          num_src_nodes = num_nodes[input_type[0]]
          num_dst_nodes = num_nodes[input_type[-1]]
  
      num_pos = src.numel()
      num_neg = 0
  
      # Negative Sampling #######################################################
  
      if neg_sampling is not None:
          # When we are doing negative sampling, we append negative information
          # of nodes/edges to `src`, `dst`, `src_time`, `dst_time`.
          # Later on, we can easily reconstruct what belongs to positive and
          # negative examples by slicing via `num_pos`.
          num_neg = math.ceil(num_pos * neg_sampling.amount)
  
          if neg_sampling.is_binary():
              # In the "binary" case, we randomly sample negative pairs of nodes.
              if isinstance(node_time, dict):
                  src_node_time = node_time.get(input_type[0])
              else:
                  src_node_time = node_time
  
              src_neg = neg_sample(
                  src,
                  neg_sampling,
                  num_src_nodes,
                  src_time,
                  src_node_time,
              )
              src = torch.cat([src, src_neg], dim=0)
  
              if isinstance(node_time, dict):
                  dst_node_time = node_time.get(input_type[-1])
              else:
                  dst_node_time = node_time
  
              dst_neg = neg_sample(
                  dst,
                  neg_sampling,
                  num_dst_nodes,
                  dst_time,
                  dst_node_time,
              )
              dst = torch.cat([dst, dst_neg], dim=0)
  
              if edge_label is None:
                  edge_label = torch.ones(num_pos)
              size = (num_neg,) + edge_label.size()[1:]
              edge_neg_label = edge_label.new_zeros(size)
              edge_label = torch.cat([edge_label, edge_neg_label])
  
              if edge_label_time is not None:
                  src_time = dst_time = edge_label_time.repeat(
                      1 + math.ceil(neg_sampling.amount)
                  )[: num_pos + num_neg]
  
          elif neg_sampling.is_triplet():
              # In the "triplet" case, we randomly sample negative destinations.
              if isinstance(node_time, dict):
                  dst_node_time = node_time.get(input_type[-1])
              else:
                  dst_node_time = node_time
  
              dst_neg = neg_sample(
                  dst,
                  neg_sampling,
                  num_dst_nodes,
                  dst_time,
                  dst_node_time,
              )
              dst = torch.cat([dst, dst_neg], dim=0)
  
              assert edge_label is None
  
              if edge_label_time is not None:
                  dst_time = edge_label_time.repeat(
                      1 + neg_sampling.amount
                  )
  
      # Heterogeneus Neighborhood Sampling ######################################
  
      if input_type is not None:
          seed_time_dict = None
          if (
              input_type[0] != input_type[-1]
          ):  # Two distinct node types:
              if not disjoint:
                  src, inverse_src = src.unique(return_inverse=True)
                  dst, inverse_dst = dst.unique(return_inverse=True)
  
              seed_dict = {input_type[0]: src, input_type[-1]: dst}
  
              if edge_label_time is not None:  # Always disjoint.
                  seed_time_dict = {
                      input_type[0]: src_time,
                      input_type[-1]: dst_time,
                  }
  
          else:  # Only a single node type: Merge both source and destination.
              seed = torch.cat([src, dst], dim=0)
  
              if not disjoint:
                  seed, inverse_seed = seed.unique(return_inverse=True)
  
              seed_dict = {input_type[0]: seed}
  
              if edge_label_time is not None:  # Always disjoint.
                  seed_time_dict = {
                      input_type[0]: torch.cat(
                          [src_time, dst_time], dim=0
                      ),
                  }
  
          out = sample_fn(seed_dict, seed_time_dict)
  
          # Enhance `out` by label information ##################################
          if disjoint:
              for key, batch in out.batch.items():
                  out.batch[key] = batch % num_pos
  
          if neg_sampling is None or neg_sampling.is_binary():
              if disjoint:
                  if input_type[0] != input_type[-1]:
                      edge_label_index = torch.arange(num_pos + num_neg)
                      edge_label_index = edge_label_index.repeat(
                          2
                      ).view(2, -1)
                  else:
                      edge_label_index = torch.arange(
                          2 * (num_pos + num_neg)
                      )
                      edge_label_index = edge_label_index.view(2, -1)
              else:
                  if input_type[0] != input_type[-1]:
                      edge_label_index = torch.stack(
                          [
                              inverse_src,
                              inverse_dst,
                          ],
                          dim=0,
                      )
                  else:
                      edge_label_index = inverse_seed.view(2, -1)
  
              out.metadata = (
                  input_id,
                  edge_label_index,
                  edge_label,
                  src_time,
              )
  
          elif neg_sampling.is_triplet():
              if disjoint:
                  src_index = torch.arange(num_pos)
                  if input_type[0] != input_type[-1]:
                      dst_pos_index = torch.arange(num_pos)
                      # `dst_neg_index` needs to be offset such that indices with
                      # offset `num_pos` belong to the same triplet:
                      dst_neg_index = torch.arange(
                          num_pos, seed_dict[input_type[-1]].numel()
                      )
                      dst_neg_index = dst_neg_index.view(
                          -1, num_pos
                      ).t()
                  else:
                      dst_pos_index = torch.arange(num_pos, 2 * num_pos)
                      dst_neg_index = torch.arange(
                          2 * num_pos, seed_dict[input_type[-1]].numel()
                      )
                      dst_neg_index = dst_neg_index.view(
                          -1, num_pos
                      ).t()
              else:
                  if input_type[0] != input_type[-1]:
                      src_index = inverse_src
                      dst_pos_index = inverse_dst[:num_pos]
                      dst_neg_index = inverse_dst[num_pos:]
                  else:
                      src_index = inverse_seed[:num_pos]
                      dst_pos_index = inverse_seed[
                          num_pos : 2 * num_pos
                      ]
                      dst_neg_index = inverse_seed[2 * num_pos :]
  
              dst_neg_index = dst_neg_index.view(num_pos, -1).squeeze(
                  -1
              )
  
              out.metadata = (
                  input_id,
                  src_index,
                  dst_pos_index,
                  dst_neg_index,
                  src_time,
              )
  
      # Homogeneus Neighborhood Sampling ########################################
  
      else:
          seed = torch.cat([src, dst], dim=0)
          seed_time = None
  
          if not disjoint:
              seed, inverse_seed = seed.unique(return_inverse=True)
  
          if edge_label_time is not None:  # Always disjoint.
              seed_time = torch.cat([src_time, dst_time])
  
          out = sample_fn(seed, seed_time)
  
          # Enhance `out` by label information ##################################
          if neg_sampling is None or neg_sampling.is_binary():
              if disjoint:
                  out.batch = out.batch % num_pos
                  edge_label_index = torch.arange(seed.numel()).view(
                      2, -1
                  )
              else:
                  edge_label_index = inverse_seed.view(2, -1)
  
              out.metadata = (
                  input_id,
                  edge_label_index,
                  edge_label,
                  src_time,
              )
  
          elif neg_sampling.is_triplet():
              if disjoint:
                  out.batch = out.batch % num_pos
                  src_index = torch.arange(num_pos)
                  dst_pos_index = torch.arange(num_pos, 2 * num_pos)
                  # `dst_neg_index` needs to be offset such that indices with
                  # offset `num_pos` belong to the same triplet:
                  dst_neg_index = torch.arange(
                      2 * num_pos, seed.numel()
                  )
                  dst_neg_index = dst_neg_index.view(-1, num_pos).t()
              else:
                  src_index = inverse_seed[:num_pos]
                  dst_pos_index = inverse_seed[num_pos : 2 * num_pos]
                  dst_neg_index = inverse_seed[2 * num_pos :]
              dst_neg_index = dst_neg_index.view(num_pos, -1).squeeze(
                  -1
              )
  
              out.metadata = (
                  input_id,
                  src_index,
                  dst_pos_index,
                  dst_neg_index,
                  src_time,
              )
  
      return out
  
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 2 for calling a global function with example usage: edge_sample(inputs, sample_fn, num_nodes, disjoint, node_time, neg_sampling) and returns: out