code: def mn_2d_best(matrix, m, n):
      # Find all possible patterns.
      patterns = compute_valid_2d_patterns(m, n).cuda()
  
      # Find the best m:n pattern (sum of non-masked weights).
      mask = torch.cuda.IntTensor(matrix.shape).fill_(1)
      mat = reshape_2d(matrix, m, m).abs()
      pmax = torch.argmax(
          torch.matmul(
              mat, patterns.view(patterns.shape[0], m * m).t()
          ),
          dim=2,
      )
  
      # Copy best m:n patterns into mask.
      mat = mat.view(mat.shape[0] * mat.shape[1], -1)
      pmax = (
          pmax.view(pmax.shape[0] * pmax.shape[1])
          .unsqueeze(1)
          .expand(-1, mat.shape[1])
      )
      patterns = patterns.view(
          patterns.shape[0], patterns.shape[1] * patterns.shape[2]
      )
      mat = torch.gather(patterns, 0, pmax)
      mat = reshape_2d_inv(
          mat.view(matrix.shape[0] // m, matrix.shape[1] // m, m, m)
      )
      mask.copy_(mat.type(mask.type()))
      return mask
  
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 9 for calling a global function with example usage: mn_2d_best(matrix, m, n) and returns: mask