code: def index_select_nnz(
      src: SparseTensor, idx: torch.Tensor, layout: Optional[str] = None
  ) -> SparseTensor:
      assert idx.dim() == 1
  
      if get_layout(layout) == "csc":
          idx = src.storage.csc2csr()[idx]
  
      row, col, value = src.coo()
      row, col = row[idx], col[idx]
  
      if value is not None:
          value = value[idx]
  
      return SparseTensor(
          row=row,
          rowptr=None,
          col=col,
          value=value,
          sparse_sizes=src.sparse_sizes(),
          is_sorted=True,
      )
  
imports: from typing import Optional
  import torch
  from torch_scatter import gather_csr
  from torch_sparse.storage import SparseStorage, get_layout
  from torch_sparse.tensor import SparseTensor
label_desc: Use image node_id 2 for calling a global function with example usage: index_select_nnz(src, idx, layout) and returns: SparseTensor