class_docstr: Integration tests for clustering.
docstr: Verifies that weights are updated during training a clustered model.
  
  Training a clustered model should update original_weights,
  clustered_centroids and bias.
docstr_tok: ["Verifies","that","weights","are","updated","during","training","a","clustered","model",".","Training","a","clustered","model","should","update","original_weights",",","clustered_centroids","and","bias","."]
code: def testWeightsAreLearningDuringClustering(self):
      """Verifies that weights are updated during training a clustered model.
  
      Training a clustered model should update original_weights,
      clustered_centroids and bias.
      """
      original_model = keras.Sequential(
          [layers.Dense(5, input_shape=(5,))]
      )
  
      clustered_model = cluster.cluster_weights(
          original_model, **self.params
      )
  
      clustered_model.compile(
          loss=keras.losses.categorical_crossentropy,
          optimizer="adam",
          metrics=["accuracy"],
      )
  
      class CheckWeightsCallback(keras.callbacks.Callback):
          def on_train_batch_begin(self, batch, logs=None):
              # Save weights before batch
              self.original_weight_kernel = (
                  self.model.layers[0]
                  .original_clusterable_weights["kernel"]
                  .numpy()
              )
              self.cluster_centroids_kernel = (
                  self.model.layers[0]
                  .cluster_centroids["kernel"]
                  .numpy()
              )
              self.bias = self.model.layers[0].layer.bias.numpy()
  
          def on_train_batch_end(self, batch, logs=None):
              # Skip the first batch since the update from the start of training is
              # too small.
              if not batch:
                  return
              # Check weights are different after batch
              assert not np.array_equal(
                  self.original_weight_kernel,
                  self.model.layers[0]
                  .original_clusterable_weights["kernel"]
                  .numpy(),
              )
              assert not np.array_equal(
                  self.cluster_centroids_kernel,
                  self.model.layers[0]
                  .cluster_centroids["kernel"]
                  .numpy(),
              )
              assert not np.array_equal(
                  self.bias, self.model.layers[0].layer.bias.numpy()
              )
  
      clustered_model.fit(
          x=self.dataset_generator(),
          steps_per_epoch=5,
          callbacks=[CheckWeightsCallback()],
      )
  
imports: import os
  import tempfile
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tensorflow_model_optimization.python.core.clustering.keras import cluster
  from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
  from tensorflow_model_optimization.python.core.clustering.keras.experimental import cluster
label_desc: Use image node_id 18 for calling the ClusterIntegrationTest obj's underlying member method code with example usage: obj.testWeightsAreLearningDuringClustering() without return types