code: def masked_select_nnz(
      src: SparseTensor,
      mask: torch.Tensor,
      layout: Optional[str] = None,
  ) -> SparseTensor:
      assert mask.dim() == 1
  
      if get_layout(layout) == "csc":
          mask = mask[src.storage.csc2csr()]
  
      row, col, value = src.coo()
      row, col = row[mask], col[mask]
  
      if value is not None:
          value = value[mask]
  
      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_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: masked_select_nnz(src, mask, layout) and returns: SparseTensor