ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more


Strips clustering wrappers from the model.

Used in the notebooks

Used in the guide

Once a model has been clustered, this method can be used to restore the original model with the clustered weights.

Only sequential and functional models are supported for now.

model A tf.keras.Model instance with clustered layers.

A keras model with clustering wrappers removed.

ValueError if the model is not a tf.keras.Model instance.
NotImplementedError if the model is a subclass model.


orig_model = tf.keras.Model(inputs, outputs)
clustered_model = cluster_weights(orig_model)
exported_model = strip_clustering(clustered_model)

The exported_model and the orig_model have the same structure.