class_docstr: This wrapper augments a keras layer so the weight tensor may be pruned.
  
  This wrapper implements magnitude-based pruning of the weight tensors.
  Magnitude-based pruning achieves a target sparsity (s% of zeros) for a given
  weight tensor by monitoring the distribution of the absolute values of the
  weight tensor and determining the weight value (referred to as threshold)
  below which s% of elements lie. For every weight tensor being pruned, the
  wrapper maintains an identically shaped tensor (referred to as mask) which
  stores 0 if the weight value lies below the threshold.
  The mask and thresholds are computed during the training based on the
  evolution of the weight values.
  
  Block sparse patterns:
  For certain SIMD hardware architectures, it may be beneficial to induce
  spatially correlated sparsity. To train models in which the weight tensors
  have block sparse structure, the pruning wrapper can be configured with
  the block_height and block_width configuration parameters set to the desired
  block configuration (2x2, 4x4, 4x1, 1x8, etc). This is applicable to
  rank-2 weight tensor only and the tensor partitioned into non-overlapping
  blocks of size [block_height, block_dim]. Either the average or max absolute
  value in this block is taken as a proxy for the entire block
  (set by block_pooling_function configuration parameter)
  while computing the distribution of the weight values and
  the threshold for pruning.
  
  Custom keras layers:
  The pruning wrapper can also be applied to a user-defined keras layer.
  Such a layer may contain one or more weight tensors that may be pruned.
  To apply pruning wrapper to such layers, the layer should be a `PrunableLayer`
  instance or, more directly, user should define a `get_prunable_weights` method
  for the layer (Check the pruning_wrapper_test.CustomLayerPrunable for more
  details about how to define a user-defined prunable layer).
  
  Sparsity function:
  The target sparsity for the weight tensors are set through the
  pruning_schedule parameter of the pruning wrapper. The user must create a
  python callable that returns a scalar tensorflow tensor and pass this
  callable to the sparsity_function parameter. This scalar tensor contains the
  target sparsity value for the weight tensors in the layer.
  The wrapper provides the following pre-built sparsity functions:
  
  ConstantSparsity
  GradualSparsity
  
  Eg.
  params = PruningParams(frequency=10,pruning_schedule=ConstantSparsity(0.9))
  pruned_model = keras.model.Sequential()
  pruned_model.add(
      Prune(keras.layers.Dense(256), input_shape=(256,)))
  pruned_model.add(Prune(keras.layers.Dense(1024), params=params))
class_docstr_tok: ["This","wrapper","augments","a","keras","layer","so","the","weight","tensor","may","be","pruned",".","This","wrapper","implements","magnitude-based","pruning","of","the","weight","tensors",".","Magnitude-based","pruning","achieves","a","target","sparsity","(","s","%","of","zeros",")","for","a","given","weight","tensor","by","monitoring","the","distribution","of","the","absolute","values","of","the","weight","tensor","and","determining","the","weight","value","(","referred","to","as","threshold",")","below","which","s","%","of","elements","lie",".","For","every","weight","tensor","being","pruned",",","the","wrapper","maintains","an","identically","shaped","tensor","(","referred","to","as","mask",")","which","stores","0","if","the","weight","value","lies","below","the","threshold",".","The","mask","and","thresholds","are","computed","during","the","training","based","on","the","evolution","of","the","weight","values",".","Block","sparse","patterns",":","For","certain","SIMD","hardware","architectures",",","it","may","be","beneficial","to","induce","spatially","correlated","sparsity",".","To","train","models","in","which","the","weight","tensors","have","block","sparse","structure",",","the","pruning","wrapper","can","be","configured","with","the","block_height","and","block_width","configuration","parameters","set","to","the","desired","block","configuration","(","2x2",",","4x4",",","4x1",",","1x8",",","etc",")",".","This","is","applicable","to","rank-2","weight","tensor","only","and","the","tensor","partitioned","into","non-overlapping","blocks","of","size","[","block_height",",","block_dim","]",".","Either","the","average","or","max","absolute","value","in","this","block","is","taken","as","a","proxy","for","the","entire","block","(","set","by","block_pooling_function","configuration","parameter",")","while","computing","the","distribution","of","the","weight","values","and","the","threshold","for","pruning",".","Custom","keras","layers",":","The","pruning","wrapper","can","also","be","applied","to","a","user-defined","keras","layer",".","Such","a","layer","may","contain","one","or","more","weight","tensors","that","may","be","pruned",".","To","apply","pruning","wrapper","to","such","layers",",","the","layer","should","be","a","`","PrunableLayer","`","instance","or",",","more","directly",",","user","should","define","a","`","get_prunable_weights","`","method","for","the","layer","(","Check","the","pruning_wrapper_test.CustomLayerPrunable","for","more","details","about","how","to","define","a","user-defined","prunable","layer",")",".","Sparsity","function",":","The","target","sparsity","for","the","weight","tensors","are","set","through","the","pruning_schedule","parameter","of","the","pruning","wrapper",".","The","user","must","create","a","python","callable","that","returns","a","scalar","tensorflow","tensor","and","pass","this","callable","to","the","sparsity_function","parameter",".","This","scalar","tensor","contains","the","target","sparsity","value","for","the","weight","tensors","in","the","layer",".","The","wrapper","provides","the","following","pre-built","sparsity","functions",":","ConstantSparsity","GradualSparsity","Eg",".","params","=","PruningParams","(","frequency=10",",","pruning_schedule=ConstantSparsity","(","0.9",")",")","pruned_model","=","keras.model.Sequential","(",")","pruned_model.add","(","Prune","(","keras.layers.Dense","(","256",")",",","input_shape=","(","256",",",")",")",")","pruned_model.add","(","Prune","(","keras.layers.Dense","(","1024",")",",","params=params",")",")"]
docstr: Create a pruning wrapper for a keras layer.
  
  #TODO(pulkitb): Consider if begin_step should be 0 by default.
  
  Args:
    layer: The keras layer to be pruned.
    pruning_schedule: A `PruningSchedule` object that controls pruning rate
      throughout training.
    block_size: (optional) The dimensions (height, weight) for the block
      sparse pattern in rank-2 weight tensors.
    block_pooling_type: (optional) The function to use to pool weights in the
      block. Must be 'AVG' or 'MAX'.
    sparsity_m_by_n: default None, otherwise a tuple of 2 integers, indicates
      pruning with m_by_n sparsity, e.g., (2, 4): 2 zeros out of 4 consecutive
      elements. It check whether we can do pruning with m_by_n sparsity.
  
    **kwargs: Additional keyword arguments to be passed to the keras layer.
docstr_tok: ["Create","a","pruning","wrapper","for","a","keras","layer",".","#","TODO","(","pulkitb",")",":","Consider","if","begin_step","should","be","0","by","default",".","Args",":","layer",":","The","keras","layer","to","be","pruned",".","pruning_schedule",":","A","`","PruningSchedule","`","object","that","controls","pruning","rate","throughout","training",".","block_size",":","(","optional",")","The","dimensions","(","height",",","weight",")","for","the","block","sparse","pattern","in","rank-2","weight","tensors",".","block_pooling_type",":","(","optional",")","The","function","to","use","to","pool","weights","in","the","block",".","Must","be","'AVG","'","or","'MAX","'",".","sparsity_m_by_n",":","default","None",",","otherwise","a","tuple","of","2","integers",",","indicates","pruning","with","m_by_n","sparsity",",","e.g.",",","(","2",",","4",")",":","2","zeros","out","of","4","consecutive","elements",".","It","check","whether","we","can","do","pruning","with","m_by_n","sparsity",".","*","*","kwargs",":","Additional","keyword","arguments","to","be","passed","to","the","keras","layer","."]
code: def __init__(
      self,
      layer,
      pruning_schedule=pruning_sched.ConstantSparsity(0.5, 0),
      block_size=(1, 1),
      block_pooling_type="AVG",
      sparsity_m_by_n=None,
      **kwargs,
  ):
      """Create a pruning wrapper for a keras layer.
  
      #TODO(pulkitb): Consider if begin_step should be 0 by default.
  
      Args:
        layer: The keras layer to be pruned.
        pruning_schedule: A `PruningSchedule` object that controls pruning rate
          throughout training.
        block_size: (optional) The dimensions (height, weight) for the block
          sparse pattern in rank-2 weight tensors.
        block_pooling_type: (optional) The function to use to pool weights in the
          block. Must be 'AVG' or 'MAX'.
        sparsity_m_by_n: default None, otherwise a tuple of 2 integers, indicates
          pruning with m_by_n sparsity, e.g., (2, 4): 2 zeros out of 4 consecutive
          elements. It check whether we can do pruning with m_by_n sparsity.
  
        **kwargs: Additional keyword arguments to be passed to the keras layer.
      """
      self.pruning_schedule = pruning_schedule
      self.block_size = block_size
      self.block_pooling_type = block_pooling_type
      self.sparsity_m_by_n = None
  
      if sparsity_m_by_n:
          self.sparsity_m_by_n = convert_to_tuple_of_two_int(
              sparsity_m_by_n, "sparsity_m_by_n"
          )
  
      # An instance of the Pruning class. This class contains the logic to prune
      # the weights of this layer.
      self.pruning_obj = None
  
      # A list of all (weight,mask,threshold) tuples for this layer
      self.pruning_vars = []
  
      if block_pooling_type not in ["AVG", "MAX"]:
          raise ValueError(
              "Unsupported pooling type '{}'. Should be 'AVG' or 'MAX'.".format(
                  block_pooling_type
              )
          )
  
      if not isinstance(layer, tf.keras.layers.Layer):
          raise ValueError(
              "Please initialize `Prune` layer with a "
              "`Layer` instance. You passed: {input}".format(
                  input=layer
              )
          )
  
      # TODO(pulkitb): This should be pushed up to the wrappers.py
      # Name the layer using the wrapper and underlying layer name.
      # Prune(Dense) becomes prune_dense_1
      kwargs.update(
          {
              "name": f"{_to_snake_case(self.__class__.__name__)}_{layer.name}"
          }
      )
  
      if isinstance(layer, prunable_layer.PrunableLayer) or hasattr(
          layer, "get_prunable_weights"
      ):
          # Custom layer in client code which supports pruning.
          super(PruneLowMagnitude, self).__init__(layer, **kwargs)
      elif prune_registry.PruneRegistry.supports(layer):
          # Built-in keras layers which support pruning.
          super(PruneLowMagnitude, self).__init__(
              prune_registry.PruneRegistry.make_prunable(layer),
              **kwargs,
          )
      else:
          raise ValueError(
              "Please initialize `Prune` with a supported layer. Layers should "
              "either be supported by the PruneRegistry (built-in keras layers) or "
              "should be a `PrunableLayer` instance, or should has a customer "
              "defined `get_prunable_weights` method. You passed: "
              "{input}".format(input=layer.__class__)
          )
  
      self._track_trackable(layer, name="layer")
  
      # TODO(yunluli): Work-around to handle the first layer of Sequential model
      # properly. Can remove this when it is implemented in the Wrapper base
      # class.
      #
      # Enables end-user to prune the first layer in Sequential models, while
      # passing the input shape to the original layer.
      #
      # tf.keras.Sequential(
      #   prune_low_magnitude(tf.keras.layers.Dense(2, input_shape=(3,)))
      # )
      #
      # as opposed to
      #
      # tf.keras.Sequential(
      #   prune_low_magnitude(tf.keras.layers.Dense(2), input_shape=(3,))
      # )
      #
      # Without this code, the pruning wrapper doesn't have an input
      # shape and being the first layer, this causes the model to not be
      # built. Being not built is confusing since the end-user has passed an
      # input shape.
      if not hasattr(self, "_batch_input_shape") and hasattr(
          layer, "_batch_input_shape"
      ):
          self._batch_input_shape = self.layer._batch_input_shape
      metrics.MonitorBoolGauge("prune_low_magnitude_wrapper_usage").set(
          layer.__class__.__name__
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import inspect
  import re
  import numpy
  import tensorflow
  from tensorflow_model_optimization.python.core.keras import compat
  from tensorflow_model_optimization.python.core.keras import metrics
  from tensorflow_model_optimization.python.core.keras import utils
  from tensorflow_model_optimization.python.core.sparsity.keras import prunable_layer
  from tensorflow_model_optimization.python.core.sparsity.keras import prune_registry
  from tensorflow_model_optimization.python.core.sparsity.keras import pruning_impl
  from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule
  from tensorflow_model_optimization.python.core.sparsity.keras.pruning_utils import convert_to_tuple_of_two_int
label_desc: Use image node_id 1 to create a new PruneLowMagnitude object from inherited base classes: Wrapper with example: obj = PruneLowMagnitude(layer, pruning_schedule, block_size, block_pooling_type, sparsity_m_by_n)