code: def _cluster_model(
      model,
      number_of_clusters,
      preserve_sparsity=False,
      cluster_per_channel=False,
  ):
      (x_train, y_train), _ = _get_dataset()
  
      clustering_params = {
          "number_of_clusters": number_of_clusters,
          "cluster_centroids_init": cluster_config.CentroidInitialization.KMEANS_PLUS_PLUS,
          "cluster_per_channel": cluster_per_channel,
          "preserve_sparsity": preserve_sparsity,
      }
  
      # 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"],
      )
  
      # Fine-tune clustered model
      clustered_model.fit(x_train, y_train, epochs=EPOCHS_FINE_TUNING)
  
      stripped_model = cluster.strip_clustering(clustered_model)
      stripped_model.compile(
          loss=tf.keras.losses.SparseCategoricalCrossentropy(
              from_logits=True
          ),
          optimizer=opt,
          metrics=["accuracy"],
      )
  
      return stripped_model
  
file_path: model-optimization/tensorflow_model_optimization/python/core/clustering/keras/mnist_clustering_test.py
imports: from absl.testing import parameterized
  import tensorflow
  from tensorflow_model_optimization.python.core.clustering.keras import cluster
  from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
label_desc: Use image node_id 4 for calling a global function with example usage: _cluster_model(model, number_of_clusters, preserve_sparsity, cluster_per_channel) and returns: stripped_model