class_docstr: Unit tests for clustering lookup algorithms.
class_docstr_tok: ["Unit","tests","for","clustering","lookup","algorithms","."]
code: def _check_gradients_clustered_weight(
      self,
      clustering_algo,
      weight,
      pulling_indices,
      expected_grad_centroids,
  ):
      pulling_indices = tf.convert_to_tensor(pulling_indices)
      cluster_centroids = clustering_algo.cluster_centroids
  
      with tf.GradientTape(persistent=True) as t:
          t.watch(weight)
          t.watch(cluster_centroids)
  
          out = clustering_algo.get_clustered_weight(
              pulling_indices, original_weight=weight
          )
  
      grad_original_weight = t.gradient(out, weight)
      grad_cluster_centroids = t.gradient(out, cluster_centroids)
      # Because of tf.gather, the grad will be of type tf.IndexedSlices
      # Convert it back to a tf.tensor
      grad_cluster_centroids = tf.convert_to_tensor(
          grad_cluster_centroids
      )
  
      # grad_original_weight with respect to out should always be 1s
      expected_grad_weight = tf.ones(shape=grad_original_weight.shape)
  
      self.assertAllEqual(grad_original_weight, expected_grad_weight)
      self.assertAllEqual(
          grad_cluster_centroids, expected_grad_centroids
      )
  
imports: from absl.testing import parameterized
  import tensorflow
  from tensorflow_model_optimization.python.core.clustering.keras import clusterable_layer
  from tensorflow_model_optimization.python.core.clustering.keras import clustering_registry
  from tensorflow_model_optimization.python.core.clustering.keras.cluster_config import GradientAggregation
label_desc: Use image node_id 2 for calling the ClusteringAlgorithmTest obj's underlying member method code with example usage: obj._check_gradients_clustered_weight(clustering_algo, weight, pulling_indices, expected_grad_centroids) without return types