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: Overrides gradients in the backprop stage.
  
  This function overrides gradients in the backprop stage: the Jacobian
  matrix of multiplication is replaced with the identity matrix, which
  effectively changes multiplication into add in the backprop. Since
  the gradient of tf.sign is 0, overwriting it with identity follows
  the design of straight-through-estimator, which accepts all upstream
  gradients and uses them to update original non-clustered weights of
  the layer. Here, we assume the gradient updates on individual elements
  inside a cluster will be different so that there is no point in mapping
  the gradient updates back to original non-clustered weights using the LUT.
  
  Args:
    clustered_weight: clustered weights
    original_weight: original weights
  
  Returns:
    result and custom gradient, as expected by @tf.custom_gradient
docstr_tok: ["Overrides","gradients","in","the","backprop","stage",".","This","function","overrides","gradients","in","the","backprop","stage",":","the","Jacobian","matrix","of","multiplication","is","replaced","with","the","identity","matrix",",","which","effectively","changes","multiplication","into","add","in","the","backprop",".","Since","the","gradient","of","tf.sign","is","0",",","overwriting","it","with","identity","follows","the","design","of","straight-through-estimator",",","which","accepts","all","upstream","gradients","and","uses","them","to","update","original","non-clustered","weights","of","the","layer",".","Here",",","we","assume","the","gradient","updates","on","individual","elements","inside","a","cluster","will","be","different","so","that","there","is","no","point","in","mapping","the","gradient","updates","back","to","original","non-clustered","weights","using","the","LUT",".","Args",":","clustered_weight",":","clustered","weights","original_weight",":","original","weights","Returns",":","result","and","custom","gradient",",","as","expected","by","@","tf.custom_gradient"]
code: def add_gradient_to_original_weight(
      self, clustered_weight, original_weight
  ):
      """Overrides gradients in the backprop stage.
  
      This function overrides gradients in the backprop stage: the Jacobian
      matrix of multiplication is replaced with the identity matrix, which
      effectively changes multiplication into add in the backprop. Since
      the gradient of tf.sign is 0, overwriting it with identity follows
      the design of straight-through-estimator, which accepts all upstream
      gradients and uses them to update original non-clustered weights of
      the layer. Here, we assume the gradient updates on individual elements
      inside a cluster will be different so that there is no point in mapping
      the gradient updates back to original non-clustered weights using the LUT.
  
      Args:
        clustered_weight: clustered weights
        original_weight: original weights
  
      Returns:
        result and custom gradient, as expected by @tf.custom_gradient
      """
      override_weights = tf.sign(original_weight + 1e6)
      override_clustered_weight = clustered_weight * override_weights
  
      def grad(d_override_clustered_weight):
          return (
              d_override_clustered_weight,
              d_override_clustered_weight,
          )
  
      return override_clustered_weight, grad
  
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 3 for calling the ClusteringAlgorithm obj's underlying member method code with example usage: obj.add_gradient_to_original_weight(clustered_weight, original_weight) and returns: override_clustered_weight,  grad, d_override_clustered_weight,  d_override_clustered_weight