Add docstring for cluster_config.CentroidInitialization

Author: arovir01

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

@@ -108,17 +108,8 @@ def cluster_weights(to_cluster,
       number_of_clusters: the number of cluster centroids to form when
         clustering a layer/model. For example, if number_of_clusters=8 then only
         8 unique values will be used in each weight array.
-      cluster_centroids_init: enum value that determines how the cluster
-        centroids will be initialized.
-        Can have following values:
-          1. RANDOM : centroids are sampled using the uniform distribution
-          between the minimum and maximum weight values in a given layer
-          2. DENSITY_BASED : density-based sampling. First, cumulative
-          distribution function is built for weights, then y-axis is evenly
-          spaced into number_of_clusters regions. After this the corresponding x
-          values are obtained and used to initialize clusters centroids.
-          3. LINEAR : cluster centroids are evenly spaced between the minimum
-          and maximum values of a given weight
+      cluster_centroids_init: `cluster_config.CentroidInitialization` instance
+        that determines how the cluster centroids will be initialized.
       **kwargs: Additional keyword arguments to be passed to the keras layer.
         Ignored when to_cluster is not a keras layer.
 

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

@@ -18,6 +18,17 @@
 
 
 class CentroidInitialization(str, enum.Enum):
+  """Specifies how the cluster centroids should be initialized.
+  * `LINEAR`: Cluster centroids are evenly spaced between the minimum and
+      maximum values of a given weight tensor.
+  * `RANDOM`: Centroids are sampled using the uniform distribution between the
+      minimum and maximum weight values in a given layer.
+  * `DENSITY_BASED`: Density-based sampling obtained as follows: first a
+       cumulative distribution function is built for the weights, then the Y
+       axis is evenly spaced into as many regions as many clusters we want to
+       have. After this the corresponding X values are obtained and used to
+       initialize the clusters centroids.
+  """
   LINEAR = "LINEAR"
   RANDOM = "RANDOM"
   DENSITY_BASED = "DENSITY_BASED"