code: def _blockMasking(
      self, block_size, block_pooling_type, weight, expected_mask
  ):
      mask = tf.Variable(
          tf.ones(weight.get_shape(), dtype=weight.dtype),
          name="mask",
          dtype=weight.dtype,
      )
      threshold = tf.Variable(
          tf.zeros([], dtype=weight.dtype),
          name="threshold",
          dtype=weight.dtype,
      )
      self.initialize()
  
      # Set up pruning
      p = pruning_impl.Pruning(
          pruning_vars=[(weight, mask, threshold)],
          training_step_fn=self.training_step_fn,
          pruning_schedule=self.constant_sparsity,
          block_size=block_size,
          block_pooling_type=block_pooling_type,
      )
  
      _, new_mask = p._maybe_update_block_mask(weight)
      # Check if the mask is the same size as the weights
      self.assertAllEqual(new_mask.get_shape(), weight.get_shape())
      mask_after_pruning = K.get_value(new_mask)
      self.assertAllEqual(mask_after_pruning, expected_mask)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tensorflow_model_optimization.python.core.keras import compat
  from tensorflow_model_optimization.python.core.sparsity.keras import pruning_impl
  from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule
label_desc: Use image node_id 6 for calling the PruningTest obj's underlying member method code with example usage: obj._blockMasking(block_size, block_pooling_type, weight, expected_mask) without return types