code: def mn_2d_greedy(matrix, m, n):
      # Convert to numpy
      mat = matrix.cpu().detach().numpy()
      mask = np.ones(mat.shape, dtype=int)
  
      rowCount = int(mat.shape[0] / m) * m
      colCount = int(mat.shape[1] / m) * m
      for rowStartIdx in range(0, rowCount, m):
          rowEndIdx = rowStartIdx + m
          for colStartIdx in range(0, colCount, m):
              colEndIdx = colStartIdx + m
              matrixSub = np.absolute(
                  np.squeeze(
                      mat[rowStartIdx:rowEndIdx, colStartIdx:colEndIdx]
                  )
              )
              maskSub = np.squeeze(
                  mask[rowStartIdx:rowEndIdx, colStartIdx:colEndIdx]
              )
              maskSub.fill(0.0)
              matrixVecView = matrixSub.reshape(-1)
              maskVecView = maskSub.reshape(-1)
              linearIdx = np.argsort(matrixVecView)
              matrixIdx = [(int(x / m), x % m) for x in linearIdx]
              rowCounter = collections.Counter()
              colCounter = collections.Counter()
              for currIdx in range(len(linearIdx) - 1, -1, -1):
                  currMatrixEntry = matrixIdx[currIdx]
                  if (rowCounter[currMatrixEntry[0]] == n) or (
                      colCounter[currMatrixEntry[1]] == n
                  ):
                      continue
                  # end if
                  maskSub[currMatrixEntry[0], currMatrixEntry[1]] = 1.0
                  rowCounter[currMatrixEntry[0]] += 1
                  colCounter[currMatrixEntry[1]] += 1
  
      return torch.tensor(mask.cuda())
  
file_path: apex/apex/contrib/sparsity/sparse_masklib.py
imports: import sys
  import torch
  import numpy
  import collections
  from itertools import permutations
label_desc: Use image node_id 6 for calling a global function with example usage: mn_2d_greedy(matrix, m, n) and returns: torch