code: def test_compact(create_func, dim, index):
      ctx = F.ctx()
      shape = (5, 5)
      ans_idx = []
      if index is not None:
          ans_idx = list(dict.fromkeys(index))
          index = torch.tensor(index).to(ctx)
  
      A = create_func(shape, 8, ctx)
  
      A_compact, ret_id = A.compact(dim, index)
      A_compact_dense = sparse_matrix_to_dense(A_compact)
  
      A_dense = sparse_matrix_to_dense(A)
  
      for i in range(shape[dim]):
          if dim == 0:
              row = list(A_dense[i, :].nonzero().reshape(-1))
          else:
              row = list(A_dense[:, i].nonzero().reshape(-1))
          if (i not in list(ans_idx)) and len(row) > 0:
              ans_idx.append(i)
      if len(ans_idx):
          ans_idx = torch.tensor(ans_idx).to(ctx)
      A_dense_select = sparse_matrix_to_dense(
          A.index_select(dim, ans_idx)
      )
  
      assert A_compact_dense.shape == A_dense_select.shape
      assert torch.allclose(A_compact_dense, A_dense_select)
      assert torch.allclose(ans_idx, ret_id)
  
imports: import backend
  import pytest
  import torch
  from .utils import rand_coo, rand_csc, rand_csr, rand_diag, sparse_matrix_to_dense
label_desc: Use image node_id 1 for calling a global function with example usage: test_compact(create_func, dim, index) without return types