fix: Extend docstrings in the rewrite module

Rework doctrings in rewrite functions based on recent changes

Resolves MLIA-944

Signed-off-by: Nathan Bailey <nathan.bailey@arm.com>
Change-Id: I31a37e17a296f8a16d0db408d48c6de65c05300e

1 files changed, 3 insertions, 3 deletions

diff --git a/src/mlia/nn/rewrite/library/clustering.py b/src/mlia/nn/rewrite/library/clustering.py
index 6f06c48..81bfd90 100644
--- a/src/mlia/nn/rewrite/library/clustering.py
+++ b/src/mlia/nn/rewrite/library/clustering.py
@@ -1,6 +1,6 @@
 # SPDX-FileCopyrightText: Copyright 2024, Arm Limited and/or its affiliates.
 # SPDX-License-Identifier: Apache-2.0
-"""Example rewrite with one fully connected clustered layer."""
+"""Rewrite functions used to return layers ready for clustering."""
 from typing import Any
 
 import tensorflow_model_optimization as tfmot
@@ -10,7 +10,7 @@ from mlia.nn.rewrite.library.helper_functions import compute_conv2d_parameters
 
 
 def fc_clustering_rewrite(input_shape: Any, output_shape: Any) -> keras.Model:
-    """Generate TensorFlow Lite model for clustering rewrite."""
+    """Fully connected TensorFlow Lite model ready for clustering."""
     rewrite_params = {
         "number_of_clusters": 4,
         "cluster_centroids_init": tfmot.clustering.keras.CentroidInitialization.LINEAR,
@@ -29,7 +29,7 @@ def fc_clustering_rewrite(input_shape: Any, output_shape: Any) -> keras.Model:
 
 
 def conv2d_clustering_rewrite(input_shape: Any, output_shape: Any) -> keras.Model:
-    """Generate TensorFlow Lite model for clustering rewrite."""
+    """Conv2d TensorFlow Lite model ready for clustering."""
     rewrite_params = {
         "number_of_clusters": 4,
         "cluster_centroids_init": tfmot.clustering.keras.CentroidInitialization.LINEAR,