code: def test_set_diag(dtype, device):
      row, col = tensor(
          [[0, 0, 9, 9], [0, 1, 0, 1]], torch.long, device
      )
      value = tensor([1, 2, 3, 4], dtype, device)
      mat = SparseTensor(row=row, col=col, value=value)
  
      mat = mat.set_diag(tensor([-8, -8], dtype, device), k=-1)
      mat = mat.set_diag(tensor([-8], dtype, device), k=1)
  
file_path: pytorch_sparse/test/test_diag.py
imports: from itertools import product
  import pytest
  import torch
  from torch_sparse.tensor import SparseTensor
  from torch_sparse.testing import devices, dtypes, tensor
label_desc: Use image node_id 2 for calling a global function with example usage: test_set_diag(dtype, device) without return types