code: def m_by_n_sparsity_mask_prepare(mask, weights_shape):
      """Reshape and permute sparsity mask, so that it match original weights data format.
  
      This is a m_by_n sparsity helper function.
  
      Args:
        mask: A 2-D tensor. Must be rank 2 or 4.
        weights_shape: shape of weights
  
      Returns:
        A sparsity mask that matches weights data format.
  
      Raises:
        ValueError:
          if the input tensor is not rank 2 or 4.
        InvalidArgumentError:
          if number of elements mismatch between mask and weights,
          if shape of prepared mask mismatch shape of weights.
      """
      tf.debugging.assert_equal(
          tf.size(mask),
          tf.reduce_prod(weights_shape),
          message="number of elements mismatch between mask and weights.",
      )
  
      if mask.shape.rank != 2:
          raise ValueError(
              f"rank of mask(rank:{mask.shape.rank}) should be 2."
          )
  
      if weights_shape.rank == 2:
          prepared_mask = tf.transpose(mask)
      elif weights_shape.rank == 4:
          reshaped_mask = tf.reshape(
              mask,
              [
                  weights_shape[-1],
                  weights_shape[0],
                  weights_shape[1],
                  weights_shape[2],
              ],
          )
          prepared_mask = tf.transpose(reshaped_mask, perm=[1, 2, 3, 0])
      else:
          raise ValueError(
              f"weight tensor with shape: {weights_shape} is not supported."
          )
  
      tf.debugging.assert_equal(
          prepared_mask.shape,
          weights_shape,
          message="shape of prepared mask mismatch shape of weights.",
      )
  
      return prepared_mask
  
file_path: model-optimization/tensorflow_model_optimization/python/core/sparsity/keras/pruning_utils.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  import tensorflow
label_desc: Use image node_id 6 for calling a global function with example usage: m_by_n_sparsity_mask_prepare(mask, weights_shape) and returns: prepared_mask