code: def sample(
      src: SparseTensor,
      num_neighbors: int,
      subset: Optional[torch.Tensor] = None,
  ) -> torch.Tensor:
      rowptr, col, _ = src.csr()
      rowcount = src.storage.rowcount()
  
      if subset is not None:
          rowcount = rowcount[subset]
          rowptr = rowptr[subset]
      else:
          rowptr = rowptr[:-1]
  
      rand = torch.rand(
          (rowcount.size(0), num_neighbors), device=col.device
      )
      rand.mul_(rowcount.to(rand.dtype).view(-1, 1))
      rand = rand.to(torch.long)
      rand.add_(rowptr.view(-1, 1))
  
      return col[rand]
  
imports: from typing import Optional, Tuple
  import torch
  from torch_sparse.tensor import SparseTensor
label_desc: Use image node_id 1 for calling a global function with example usage: sample(src, num_neighbors, subset) and returns: col