code: def test_nm_pruner_mask_computation(weights_shape, dim):
      weights = torch.tensor(
          [[2, 3, 0, 4, 5, 9, 1, 1], [3, 6, 1, 0, 2, 3, 8, 9]]
      )
      if dim == 1:
          expected_mask = torch.tensor(
              [[0, 1, 0, 1, 1, 1, 0, 0], [1, 1, 0, 0, 0, 0, 1, 1]]
          )
          nm = (2, 4)
      else:
          expected_mask = torch.tensor(
              [[0, 0, 0, 1, 1, 1, 0, 0], [1, 1, 1, 0, 0, 0, 1, 1]]
          )
          nm = (1, 2)
      if weights_shape == [2, 8, 1, 1]:
          weights = weights.reshape([2, 8, 1, 1])
          expected_mask = expected_mask.reshape([2, 8, 1, 1])
      mask = n_m_mask(weights, nm, dim=dim)
      np.testing.assert_array_equal(mask, expected_mask)
  
imports: import copy
  from collections import OrderedDict
  import numpy
  import pytest
  import torch
  from coremltools.optimize.torch.pruning import MagnitudePruner, MagnitudePrunerConfig, ModuleMagnitudePrunerConfig
  from coremltools.optimize.torch.pruning._utils import n_m_mask
label_desc: Use image node_id 21 for calling a global function with example usage: test_nm_pruner_mask_computation(weights_shape, dim) without return types