def METHOD_NAME(matrix, m, n):
  mat = reshape_tensor(matrix, m)
  topk, indices = mat.abs().topk(k=n, dim=1, sorted=False)
  mask_matrix = torch.zeros_like(mat)
  mask_min = torch.min(topk, dim=-1).values
  mask_min = mask_min.unsqueeze(-1).repeat(1, m)
  mask = torch.where(
      torch.ge(mat.abs(), mask_min) & torch.ne(mat.abs(), mask_matrix), mat,
      mask_matrix)
  return mask