code: def test_get_diag(dtype, device):
      row, col = tensor(
          [[0, 0, 1, 2], [0, 1, 2, 2]], torch.long, device
      )
      value = tensor([[1, 1], [2, 2], [3, 3], [4, 4]], dtype, device)
      mat = SparseTensor(row=row, col=col, value=value)
      assert mat.get_diag().tolist() == [[1, 1], [0, 0], [4, 4]]
  
      row, col = tensor(
          [[0, 0, 1, 2], [0, 1, 2, 2]], torch.long, device
      )
      mat = SparseTensor(row=row, col=col)
      assert mat.get_diag().tolist() == [1, 0, 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 4 for calling a global function with example usage: test_get_diag(dtype, device) without return types