code: def cluster_model(model, x_train, y_train, x_test, y_test):
      print("Clustering model")
  
      clustering_params = {
          "number_of_clusters": 8,
          "cluster_centroids_init": cluster_config.CentroidInitialization.DENSITY_BASED,
      }
  
      # Cluster model
      clustered_model = cluster.cluster_weights(
          model, **clustering_params
      )
  
      # Use smaller learning rate for fine-tuning
      # clustered model
      opt = tf.keras.optimizers.Adam(learning_rate=1e-5)
  
      clustered_model.compile(
          loss=tf.keras.losses.SparseCategoricalCrossentropy(
              from_logits=True
          ),
          optimizer=opt,
          metrics=["accuracy"],
      )
  
      # Add callback for tensorboard summaries
      log_dir = os.path.join(
          FLAGS.output_dir,
          datetime.datetime.now().strftime("%Y%m%d-%H%M%S-clustering"),
      )
      callbacks = [
          clustering_callbacks.ClusteringSummaries(
              log_dir,
              cluster_update_freq="epoch",
              update_freq="batch",
              histogram_freq=1,
          )
      ]
  
      # Fine-tune clustered model
      clustered_model.fit(
          x_train,
          y_train,
          batch_size=batch_size,
          epochs=epochs_fine_tuning,
          verbose=1,
          callbacks=callbacks,
          validation_split=0.1,
      )
  
      score = clustered_model.evaluate(x_test, y_test, verbose=0)
      print("Clustered model test loss:", score[0])
      print("Clustered model test accuracy:", score[1])
  
      return clustered_model
  
file_path: model-optimization/tensorflow_model_optimization/python/examples/clustering/keras/mnist/mnist_cnn.py
imports: from __future__ import print_function
  import datetime
  import os
  from absl import app
  from absl import flags
  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 import clustering_callbacks
label_desc: Use image node_id 4 for calling a global function with example usage: cluster_model(model, x_train, y_train, x_test, y_test) and returns: clustered_model