class_docstr: Class to implement highly efficient vectorised look-ups.
  
  We do not utilise looping for that purpose, instead we `smartly` reshape and
  tile arrays. The trade-off is that we are potentially using way more memory
  than we would have if looping is used.
  
  Each class that inherits from this class is supposed to implement a
  particular lookup function for a certain shape.
  
  For example, look-ups for 2D table will be different in the case of 3D.
class_docstr_tok: ["Class","to","implement","highly","efficient","vectorised","look-ups",".","We","do","not","utilise","looping","for","that","purpose",",","instead","we","`","smartly","`","reshape","and","tile","arrays",".","The","trade-off","is","that","we","are","potentially","using","way","more","memory","than","we","would","have","if","looping","is","used",".","Each","class","that","inherits","from","this","class","is","supposed","to","implement","a","particular","lookup","function","for","a","certain","shape",".","For","example",",","look-ups","for","2D","table","will","be","different","in","the","case","of","3D","."]
docstr: Returns clustered weights with custom gradients.
  
  Take indices (pulling_indices) as input and then form a new array
  by gathering cluster centroids based on the given pulling indices.
  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 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","(","pulling_indices",")","as","input","and","then","form","a","new","array","by","gathering","cluster","centroids","based","on","the","given","pulling","indices",".","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","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 (pulling_indices) as input and then form a new array
      by gathering cluster centroids based on the given pulling indices.
      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 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.
      """
  
      if self.cluster_gradient_aggregation == GradientAggregation.SUM:
          cluster_centroids = self.cluster_centroids
      elif self.cluster_gradient_aggregation == GradientAggregation.AVG:
          cluster_centroids = self.cluster_centroids
          # Compute the size of each cluster
          # (number of weights belonging to each cluster)
          cluster_sizes = tf.math.bincount(
              arr=tf.cast(pulling_indices, dtype=tf.int32),
              minlength=tf.size(cluster_centroids),
              dtype=cluster_centroids.dtype,
          )
          # Modify the gradient of cluster_centroids to be averaged by cluster sizes
          cluster_centroids = (
              self.average_centroids_gradient_by_cluster_size(
                  cluster_centroids,
                  tf.stop_gradient(cluster_sizes),
              )
          )
      else:
          raise ValueError(
              f"self.cluster_gradient_aggregation="
              f"{self.cluster_gradient_aggregation} not implemented."
          )
  
      # Gather the clustered weights based on cluster centroids and
      # pulling indices.
      clustered_weight = tf.gather(cluster_centroids, pulling_indices)
  
      # Add an estimated gradient to the original weight
      clustered_weight = self.add_gradient_to_original_weight(
          clustered_weight,
          # Fix the bug with MirroredVariable and tf.custom_gradient:
          # tf.identity will transform a MirroredVariable into a Variable
          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 5 for calling the ClusteringAlgorithm obj's underlying member method code with example usage: obj.get_clustered_weight(pulling_indices, original_weight) and returns: clustered_weight