code: def set_diag(
      src: SparseTensor, values: Optional[Tensor] = None, k: int = 0
  ) -> SparseTensor:
      src = remove_diag(src, k=k)
      row, col, value = src.coo()
  
      mask = torch.ops.torch_sparse.non_diag_mask(
          row, col, src.size(0), src.size(1), k
      )
      inv_mask = ~mask
  
      start, num_diag = -k if k < 0 else 0, mask.numel() - row.numel()
      diag = torch.arange(start, start + num_diag, device=row.device)
  
      new_row = row.new_empty(mask.size(0))
      new_row[mask] = row
      new_row[inv_mask] = diag
  
      new_col = col.new_empty(mask.size(0))
      new_col[mask] = col
      new_col[inv_mask] = diag.add_(k)
  
      new_value: Optional[Tensor] = None
      if value is not None:
          new_value = value.new_empty(
              (mask.size(0),) + value.size()[1:]
          )
          new_value[mask] = value
          if values is not None:
              new_value[inv_mask] = values
          else:
              new_value[inv_mask] = torch.ones(
                  (num_diag,), dtype=value.dtype, device=value.device
              )
  
      rowcount = src.storage._rowcount
      if rowcount is not None:
          rowcount = rowcount.clone()
          rowcount[start : start + num_diag] += 1
  
      colcount = src.storage._colcount
      if colcount is not None:
          colcount = colcount.clone()
          colcount[start + k : start + num_diag + k] += 1
  
      storage = SparseStorage(
          row=new_row,
          rowptr=None,
          col=new_col,
          value=new_value,
          sparse_sizes=src.sparse_sizes(),
          rowcount=rowcount,
          colptr=None,
          colcount=colcount,
          csr2csc=None,
          csc2csr=None,
          is_sorted=True,
      )
      return src.from_storage(storage)
  
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 2 for calling a global function with example usage: set_diag(src, values, k) and returns: src