class_docstr: Class for Per-channel clustering of Conv2D layers.
class_docstr_tok: ["Class","for","Per-channel","clustering","of","Conv2D","layers","."]
docstr: Returns clustered weights with custom gradients.
  
  Take indices the per-channel pulling_indices as input and retrieve
  the corresponding clustered weights by using the gather operation
  for each of the channels.
  The original gradients will also be modified in two ways:
  - By averaging the gradient of cluster_centroids based on the size of
    each cluster.
  - By adding an estimated gradient onto the non-differentiable
    original weight.
  Args:
    pulling_indices: a tensor of per-channel indices used for lookup of the
      same size as original_weight.
    original_weight: the original weights of the wrapped layer.
  
  Returns:
    array with the same shape as `pulling_indices`. Each array element
    is a member of self.cluster_centroids. The backward pass is modified by
    adding custom gradients.
docstr_tok: ["Returns","clustered","weights","with","custom","gradients",".","Take","indices","the","per-channel","pulling_indices","as","input","and","retrieve","the","corresponding","clustered","weights","by","using","the","gather","operation","for","each","of","the","channels",".","The","original","gradients","will","also","be","modified","in","two","ways",":","-","By","averaging","the","gradient","of","cluster_centroids","based","on","the","size","of","each","cluster",".","-","By","adding","an","estimated","gradient","onto","the","non-differentiable","original","weight",".","Args",":","pulling_indices",":","a","tensor","of","per-channel","indices","used","for","lookup","of","the","same","size","as","original_weight",".","original_weight",":","the","original","weights","of","the","wrapped","layer",".","Returns",":","array","with","the","same","shape","as","`","pulling_indices","`",".","Each","array","element","is","a","member","of","self.cluster_centroids",".","The","backward","pass","is","modified","by","adding","custom","gradients","."]
code: def get_clustered_weight(self, pulling_indices, original_weight):
      """Returns clustered weights with custom gradients.
  
      Take indices the per-channel pulling_indices as input and retrieve
      the corresponding clustered weights by using the gather operation
      for each of the channels.
      The original gradients will also be modified in two ways:
      - By averaging the gradient of cluster_centroids based on the size of
        each cluster.
      - By adding an estimated gradient onto the non-differentiable
        original weight.
      Args:
        pulling_indices: a tensor of per-channel indices used for lookup of the
          same size as original_weight.
        original_weight: the original weights of the wrapped layer.
  
      Returns:
        array with the same shape as `pulling_indices`. Each array element
        is a member of self.cluster_centroids. The backward pass is modified by
        adding custom gradients.
      """
      num_channels = (
          original_weight.shape[1]
          if self.data_format == "channels_first"
          else original_weight.shape[-1]
      )
  
      # In case of channels_last, we have NHWC.
      # In case of channels_first, we have NCHW.
      # We need to transpose the tensor, so C is the first dimension
      # and then we could loop over channels
      pulling_indices = tf.transpose(
          pulling_indices,
          perm=(1, 0, 2, 3)
          if self.data_format == "channels_first"
          else (3, 0, 1, 2),
      )
  
      if self.cluster_gradient_aggregation == GradientAggregation.SUM:
          cluster_centroids = self.cluster_centroids
      elif self.cluster_gradient_aggregation == GradientAggregation.AVG:
          cluster_sizes = []
          for i in range(num_channels):
              # Compute the size of each cluster for each channel
              # (number of weights belonging to each cluster)
              cluster_sizes.append(
                  tf.math.bincount(
                      arr=tf.cast(pulling_indices[i], dtype=tf.int32),
                      minlength=tf.size(self.cluster_centroids[i]),
                      dtype=self.cluster_centroids.dtype,
                  )
              )
  
          cluster_sizes = tf.convert_to_tensor(cluster_sizes)
  
          # Modify the gradient of cluster_centroids to be averaged by cluster sizes
          cluster_centroids = (
              self.average_centroids_gradient_by_cluster_size(
                  self.cluster_centroids,
                  tf.stop_gradient(cluster_sizes),
              )
          )
      else:
          raise ValueError(
              f"self.cluster_gradient_aggregation="
              f"{self.cluster_gradient_aggregation} not implemented."
          )
  
      clustered_weights = []
  
      for i in range(num_channels):
          clustered_weights.append(
              tf.gather(cluster_centroids[i], pulling_indices[i])
          )
  
      clustered_weight = tf.convert_to_tensor(clustered_weights)
  
      # Permute weights to ensure the channels are first or last, as expected
      # based on the data_format attribute
      clustered_weight = (
          tf.transpose(clustered_weights, perm=[1, 0, 2, 3])
          if self.data_format == "channels_first"
          else tf.transpose(clustered_weights, perm=[1, 2, 3, 0])
      )
  
      # Add an estimated gradient to the original weight
      clustered_weight = self.add_gradient_to_original_weight(
          clustered_weight,
          tf.identity(original_weight),
      )
  
      return clustered_weight
  
file_path: model-optimization/tensorflow_model_optimization/python/core/clustering/keras/clustering_algorithm.py
imports: import abc
  import six
  import tensorflow
  from tensorflow_model_optimization.python.core.clustering.keras.cluster_config import GradientAggregation
label_desc: Use image node_id 2 for calling the ClusteringAlgorithmPerChannel obj's underlying member method code with example usage: obj.get_clustered_weight(pulling_indices, original_weight) and returns: clustered_weight