Merge pull request #567 from wwwind:clustering_tidy_up_test

PiperOrigin-RevId: 337365407

Author: tensorflower-gardener

tensorflow_model_optimization/python/core/clustering/keras/cluster_wrapper_test.py

@@ -160,11 +160,7 @@ def testIfLayerHasBatchShapeClusterWeightsMustHaveIt(self):
   @parameterized.parameters(
       *itertools.product(
           range(2, 16, 4),
-          (
-              CentroidInitialization.LINEAR,
-              CentroidInitialization.RANDOM,
-              CentroidInitialization.DENSITY_BASED
-          )
+          [type_centroid for type_centroid in CentroidInitialization]
       )
   )
   def testValuesAreClusteredAfterStripping(self,
@@ -178,6 +174,9 @@ def testValuesAreClusteredAfterStripping(self,
     original_model = tf.keras.Sequential([
         layers.Dense(32, input_shape=(10,)),
     ])
+    self.assertGreater(
+      len(set(original_model.get_weights()[0].reshape(-1,).tolist())),
+      number_of_clusters)
     clustered_model = cluster.cluster_weights(
         original_model,
         number_of_clusters=number_of_clusters,