code: def get_diag(src: SparseTensor) -> Tensor:
      row, col, value = src.coo()
  
      if value is None:
          value = torch.ones(row.size(0), device=row.device)
  
      sizes = list(value.size())
      sizes[0] = min(src.size(0), src.size(1))
  
      out = value.new_zeros(sizes)
  
      mask = row == col
      out[row[mask]] = value[mask]
      return out
  
file_path: pytorch_sparse/torch_sparse/diag.py
imports: from typing import Optional
  import torch
  from torch import Tensor
  from torch_sparse.storage import SparseStorage
  from torch_sparse.tensor import SparseTensor
label_desc: Use image node_id 4 for calling a global function with example usage: get_diag(src) and returns: out