mct-nightly 2.1.0.20240617.451__py3-none-any.whl → 2.1.0.20240618.432__py3-none-any.whl

model_compression_toolkit/xquant/common/model_folding_utils.py

@@ -0,0 +1,104 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
+from model_compression_toolkit.core.common.framework_info import FrameworkInfo
+
+from model_compression_toolkit.core.common.model_builder_mode import ModelBuilderMode
+from model_compression_toolkit.core.common.quantization.quantization_config import DEFAULTCONFIG
+
+from model_compression_toolkit.core.graph_prep_runner import graph_preparation_runner
+from typing import Any, Callable
+
+from model_compression_toolkit.core.common import Graph
+from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
+
+
+class ModelFoldingUtils:
+    """
+    Utility class for handling model folding operations such as batch normalization (BN) folding,
+    residual collapsing, and other graph optimizations.
+    """
+
+    def __init__(self,
+                 fw_info: FrameworkInfo,
+                 fw_impl: FrameworkImplementation,
+                 fw_default_tpc: TargetPlatformCapabilities):
+        """
+        Initialize the ModelFoldingUtils class with framework-specific information, implementation details,
+        and default TPC.
+
+        Args:
+            fw_info: Framework-specific information.
+            fw_impl: Implementation functions for the framework.
+            fw_default_tpc: Default target platform capabilities for the handled framework.
+        """
+        self.fw_info = fw_info
+        self.fw_impl = fw_impl
+        self.fw_default_tpc = fw_default_tpc
+
+    def create_float_folded_model(self, float_model: Any, representative_dataset: Any = None) -> Any:
+        """
+        Create folded version of the model like MCT does (bn folding, residual collapsing, etc.).
+        This is needed since we need the models we compare to have the same architecture for
+        comparing tensors in different points of the models.
+
+        Args:
+            float_model: The floating-point model to be folded.
+            representative_dataset: A callable for generating representative data.
+
+        Returns:
+            The folded floating-point model.
+
+        """
+        float_graph = self.create_float_folded_graph(model=float_model,
+                                                     repr_dataset=representative_dataset)
+        float_folded_model, _ = self.fw_impl.model_builder(
+            float_graph,
+            mode=ModelBuilderMode.FLOAT,
+            append2output=None,
+            fw_info=self.fw_info
+        )
+        return float_folded_model
+
+    def create_float_folded_graph(self, model: Any, repr_dataset: Callable) -> Graph:
+        """
+        Create a folded graph for the float model. This process involves
+        graph optimizations similar to those applied during quantization (e.g., batch normalization folding,
+        residual collapsing).
+
+        Args:
+            model: The floating-point model to be folded into a graph.
+            repr_dataset: A callable that generates representative data.
+
+        Returns:
+            The folded graph.
+        """
+        # TODO:
+        # Consider simplifying graph_preparation_runner by extracting relevant parts to a separate method in MCT.
+        #
+        # Issues:
+        # 1. The quantization config affects how the optimized graph looks (e.g., collapsing).
+        # 2. The back2fw function requires quantization info even for float models.
+        #
+        # Future Considerations:
+        # - Remove quantization config parts related to graph optimizations.
+        # - Update back2fw to handle float models without needing quantization info.
+        graph = graph_preparation_runner(in_model=model,
+                                         representative_data_gen=repr_dataset,
+                                         fw_impl=self.fw_impl,
+                                         fw_info=self.fw_info,
+                                         quantization_config=DEFAULTCONFIG,
+                                         tpc=self.fw_default_tpc)
+        return graph

model_compression_toolkit/xquant/common/similarity_calculator.py

@@ -0,0 +1,194 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+from functools import partial
+
+from typing import Tuple, Any, Dict, Callable
+
+from model_compression_toolkit.xquant.common.constants import MODEL_OUTPUT_KEY
+from model_compression_toolkit.xquant.common.dataset_utils import DatasetUtils
+from model_compression_toolkit.xquant.common.model_analyzer import ModelAnalyzer
+from model_compression_toolkit.xquant.common.model_folding_utils import ModelFoldingUtils
+from model_compression_toolkit.xquant.common.similarity_functions import SimilarityFunctions
+from model_compression_toolkit.logger import Logger
+
+class SimilarityCalculator:
+    """
+    A class to calculate the similarity between two models (that are often referred as float
+    and quantized models). It utilizes various utility classes for dataset preparation, model folding,
+    similarity computation, and model analysis.
+    """
+
+    def __init__(self,
+                 dataset_utils: DatasetUtils,
+                 model_folding: ModelFoldingUtils,
+                 similarity_functions: SimilarityFunctions,
+                 model_analyzer_utils: ModelAnalyzer,
+                 device: str = None):
+        """
+        Initialize the SimilarityCalculator with required utilities.
+
+        Args:
+            dataset_utils (DatasetUtils): Utility class for dataset preparation.
+            model_folding (ModelFoldingUtils): Utility class for model folding operations.
+            similarity_functions (SimilarityFunctions): Class containing similarity functions.
+            model_analyzer_utils (ModelAnalyzer): Utility class for model analysis.
+            device (str, optional): Device to perform computations on (e.g., 'cpu', 'cuda'). Defaults to None.
+        """
+        self.dataset_utils = dataset_utils
+        self.model_folding = model_folding
+        self.similarity_functions = similarity_functions
+        self.model_analyzer_utils = model_analyzer_utils
+        self.device = device
+
+    @staticmethod
+    def compute_tensors_similarity(tensors_to_compare: Tuple[Any, Any],
+                                   similarity_metrics: Dict[str, Callable]) -> Dict[str, float]:
+        """
+        Compute the similarity between two tensors using provided similarity metrics.
+
+        Args:
+            tensors_to_compare (Tuple[Any, Any]): Tensors to compare by computing their similarity.
+            similarity_metrics (Dict[str, Callable]): A dictionary with similarity metric names and functions.
+
+        Returns:
+            Dict[str, float]: A dictionary of similarity metric names and their computed values.
+        """
+        x, y = tensors_to_compare
+        similarity_metrics = {k: v(x, y) for k, v in similarity_metrics.items()}
+        return similarity_metrics
+
+    def _get_float_to_quantized_compare_points(self,
+                                               quantized_model: Any,
+                                               float_model: Any) -> Dict[str, str]:
+        """
+        Map corresponding layers between the float and quantized models for comparison.
+
+        Args:
+            quantized_model (Any): The quantized model.
+            float_model (Any): The float model.
+
+        Returns:
+            Dict[str, str]: A dictionary mapping float model layer names to quantized model layer names.
+        """
+        # Identify the points in the quantized model to compare.
+        quant_points_names = self.model_analyzer_utils.identify_quantized_compare_points(quantized_model)
+
+        float_name2quant_name = {}
+
+        # Extract the names of the layers in the float model.
+        float_layers_names = self.model_analyzer_utils.extract_float_layer_names(float_model)
+
+        # Map each quantized layer to the corresponding float layer.
+        for quant_point in quant_points_names:
+            candidate_float_layer_name = self.model_analyzer_utils.find_corresponding_float_layer(
+                quant_compare_point=quant_point, quantized_model=quantized_model)
+
+            if candidate_float_layer_name in float_layers_names:
+                if candidate_float_layer_name not in float_name2quant_name:
+                    float_name2quant_name[candidate_float_layer_name] = quant_point
+                else:
+                    Logger.critical(f"Duplicate mapping found for layer: {candidate_float_layer_name}.")
+            else:
+                Logger.warning(
+                    f"Could not find a matching layer in the float model for layer with name {quant_point}, "
+                    f"skipping it in similarity metrics comparison points computation.")
+
+        return float_name2quant_name
+
+    def compute_similarity_metrics(self,
+                                   float_model: Any,
+                                   quantized_model: Any,
+                                   dataset: Callable,
+                                   custom_similarity_metrics: Dict[str, Callable] = None,
+                                   is_validation: bool = False) -> Tuple[Dict[str, float], Dict[str, Dict[str, float]]]:
+        """
+        Compute the similarity metrics between the two models (usually, float and quantized models).
+
+        Args:
+            float_model (Any): The float model.
+            quantized_model (Any): The quantized model.
+            dataset (Callable): A callable to provide the dataset.
+            custom_similarity_metrics (Dict[str, Callable], optional): Custom similarity metrics. Defaults to None.
+            is_validation (bool, optional): Flag to indicate if the dataset is for validation. Defaults to False.
+
+        Returns:
+            Tuple[Dict[str, float], Dict[str, Dict[str, float]]]: Aggregated output similarity metrics and
+            intermediate similarity metrics for each layer.
+        """
+        # Prepare the dataset such that the rest of operations are indistinguishable between the representative
+        # dataset and the validation dataset.
+        dataset = partial(self.dataset_utils.prepare_dataset,
+                          dataset=dataset,
+                          is_validation=is_validation,
+                          device=self.device)
+
+        # Create a folded version of the float model.
+        float_model = self.model_folding.create_float_folded_model(float_model=float_model,
+                                                                   representative_dataset=dataset)
+
+        # Gather similarity metrics to compute (default and custom).
+        similarity_metrics_to_compute = self.similarity_functions.get_default_similarity_metrics()
+        if custom_similarity_metrics:
+            if not isinstance(custom_similarity_metrics, dict):
+                Logger.critical(
+                    f"custom_similarity_metrics should be a dictionary but is of type "
+                    f"{type(custom_similarity_metrics)}.")
+            similarity_metrics_to_compute.update(custom_similarity_metrics)
+
+        # Map float model layers to quantized model layers for comparison.
+        float_name2quant_name = self._get_float_to_quantized_compare_points(float_model=float_model,
+                                                                            quantized_model=quantized_model)
+
+        # Initialize dictionaries to store similarity metrics.
+        output_similarity_metrics = {key: [] for key in similarity_metrics_to_compute.keys()}
+        intermediate_similarity_metrics = {layer: {key: [] for key in similarity_metrics_to_compute.keys()} for layer in
+                                           float_name2quant_name.values()}
+
+        # Iterate over the dataset and compute similarity metrics.
+        for x in dataset():
+            # Extract activations and predictions from both models.
+            float_activations, quant_activations = (
+                self.model_analyzer_utils.extract_model_activations(
+                float_model, quantized_model, float_name2quant_name, x))
+
+            float_predictions = float_activations[MODEL_OUTPUT_KEY]
+            quant_predictions = quant_activations[MODEL_OUTPUT_KEY]
+
+            # Compute similarity metrics for the output predictions.
+            output_results = self.compute_tensors_similarity((float_predictions, quant_predictions),
+                                                             similarity_metrics_to_compute)
+            for key in output_similarity_metrics:
+                output_similarity_metrics[key].append(output_results[key])
+
+            # Compute similarity metrics for each intermediate layer.
+            for float_layer, quant_layer in float_name2quant_name.items():
+                intermediate_results = self.compute_tensors_similarity(
+                    (float_activations[float_layer], quant_activations[quant_layer]),
+                    similarity_metrics_to_compute)
+                for key in intermediate_similarity_metrics[quant_layer]:
+                    intermediate_similarity_metrics[quant_layer][key].append(intermediate_results[key])
+
+        # Aggregate the output similarity metrics.
+        aggregated_output_similarity_metrics = {key: sum(value) / len(value) for key, value in
+                                                output_similarity_metrics.items()}
+
+        # Aggregate the intermediate similarity metrics for each layer.
+        for layer_name, layer_similarity_metrics in intermediate_similarity_metrics.items():
+            for similarity_name, similarity_values_list in layer_similarity_metrics.items():
+                if len(similarity_values_list) == 0:
+                    Logger.critical(f"Can not average similarities of an empty list.")
+                intermediate_similarity_metrics[layer_name][similarity_name] = sum(similarity_values_list) / len(similarity_values_list)
+
+        return aggregated_output_similarity_metrics, intermediate_similarity_metrics

model_compression_toolkit/xquant/common/similarity_functions.py

@@ -0,0 +1,81 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+
+from typing import Any, Dict, Callable
+
+from model_compression_toolkit.xquant.common.constants import CS_SIMILARITY_METRIC_NAME, SQNR_SIMILARITY_METRIC_NAME, MSE_SIMILARITY_METRIC_NAME
+
+DEFAULT_SIMILARITY_METRICS_NAMES = [CS_SIMILARITY_METRIC_NAME, MSE_SIMILARITY_METRIC_NAME, SQNR_SIMILARITY_METRIC_NAME]
+
+class SimilarityFunctions:
+    """
+    A class that provides various static methods to compute similarity metrics between tensors.
+    """
+
+    @staticmethod
+    def compute_mse(x: Any, y: Any) -> float:
+        """
+        Compute the Mean Squared Error (MSE) between two tensors (usually, the float and quantized predictions).
+
+        Args:
+            x (Any): First tensor to compare.
+            y (Any): Second tensor to compare.
+
+        Returns:
+            float: The computed MSE value.
+        """
+        raise NotImplemented  # pragma: no cover
+
+    @staticmethod
+    def compute_cs(x: Any, y: Any) -> float:
+        """
+        Compute the Cosine Similarity (CS) between two tensors (usually, the float and quantized predictions).
+
+        Args:
+            x (Any): First tensor to compare.
+            y (Any): Second tensor to compare.
+
+        Returns:
+            float: The computed CS value.
+        """
+        raise NotImplemented  # pragma: no cover
+
+    @staticmethod
+    def compute_sqnr(x: Any, y: Any) -> float:
+        """
+        Compute the Signal-to-Quantization-Noise Ratio (SQNR) between two tensors (usually, the float and quantized predictions).
+
+        Args:
+            x (Any): First tensor to compare.
+            y (Any): Second tensor to compare.
+
+        Returns:
+            float: The computed SQNR value.
+        """
+        raise NotImplemented  # pragma: no cover
+
+    def get_default_similarity_metrics(self) -> Dict[str, Callable]:
+        """
+        Get the default similarity metrics to compute.
+
+        Returns:
+            Dict[str, Callable]: A dictionary where the keys are similarity metric names and the values are the corresponding functions.
+        """
+        return {
+            MSE_SIMILARITY_METRIC_NAME: self.compute_mse,
+            CS_SIMILARITY_METRIC_NAME: self.compute_cs,
+            SQNR_SIMILARITY_METRIC_NAME: self.compute_sqnr
+        }
+

model_compression_toolkit/xquant/common/tensorboard_utils.py

@@ -0,0 +1,101 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+
+from model_compression_toolkit.core.common import Graph
+from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
+from model_compression_toolkit.core.common.framework_info import FrameworkInfo
+
+
+from model_compression_toolkit.core.common.visualization.tensorboard_writer import TensorboardWriter
+from model_compression_toolkit.xquant.common.constants import TENSORBOARD_DEFAULT_TAG
+from model_compression_toolkit.logger import Logger
+
+
+from typing import Any, Dict, Callable
+
+
+class TensorboardUtils:
+    """
+    Utility class for handling Tensorboard operations like adding graph to display
+    and histograms on the float model.
+    """
+
+    def __init__(self,
+                 report_dir: str,
+                 fw_info: FrameworkInfo,
+                 fw_impl: FrameworkImplementation):
+        """
+        Initialize the TensorboardUtils.
+
+        Args:
+            report_dir (str): Directory where Tensorboard logs will be stored.
+            fw_info (FrameworkInfo): Framework-specific information.
+            fw_impl (FrameworkImplementation): Framework-specific implementation.
+        """
+        self.fw_impl = fw_impl
+        self.fw_info = fw_info
+        self.tb_writer = TensorboardWriter(report_dir, fw_info)
+        Logger.info(f"Please run: tensorboard --logdir {self.tb_writer.dir_path}")
+
+    def get_graph_for_tensorboard_display(self,
+                                          quantized_model: Any,
+                                          similarity_metrics: Dict[str, Any],
+                                          repr_dataset: Callable) -> Graph:
+        """
+        Get the graph for Tensorboard display. The framework-specific implementations
+        (like KerasTensorboardUtils and PytorchTensorboardUtils) should implement this
+        as it differs between them when combining the similarity metrics into the graph.
+
+        Args:
+            quantized_model (Any): The quantized model.
+            similarity_metrics (Dict[str, Any]): Metrics for model similarity.
+            repr_dataset (Callable): Representative dataset function.
+
+        Returns:
+            Graph: The generated graph for Tensorboard display.
+        """
+        Logger.critical("This method should be implemented by the framework-specific TensorboardUtils.") # pragma: no cover
+
+    def add_histograms_to_tensorboard(self,
+                                      graph: Graph):
+        """
+        Add histograms to Tensorboard from a graph that holds these statistics.
+
+        Args:
+            graph (Graph): Graph with histograms to add to the tensorboard.
+        """
+        self.tb_writer.add_histograms(graph, TENSORBOARD_DEFAULT_TAG)
+
+    def add_graph_to_tensorboard(self,
+                                 quantized_model: Any,
+                                 similarity_metrics: Dict[str, Any],
+                                 repr_dataset: Callable):
+        """
+        Add a graph to Tensorboard. The graph represents the quantized graph
+        with the similarity metrics that were measured in different nodes.
+
+        Args:
+            quantized_model (Any): The quantized model.
+            similarity_metrics (Dict[str, Any]): The similarity metrics that were collected.
+            repr_dataset (Callable): Representative dataset to use (if needed, like in pytorch case).
+        """
+        # Generate the quantized graph with similarity metrics.
+        tb_graph = self.get_graph_for_tensorboard_display(quantized_model=quantized_model,
+                                                          similarity_metrics=similarity_metrics,
+                                                          repr_dataset=repr_dataset)
+
+        self.tb_writer.add_graph(tb_graph, TENSORBOARD_DEFAULT_TAG)
+
+

model_compression_toolkit/xquant/common/xquant_config.py

@@ -0,0 +1,39 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+
+from typing import Dict, Callable
+
+
+class XQuantConfig:
+    """
+    Configuration for generating the report.
+    It allows to set the log dir that the report will be saved in and to add similarity metrics
+    to measure between tensors of the two models.
+    """
+
+    def __init__(self,
+                 report_dir: str,
+                 custom_similarity_metrics: Dict[str, Callable] = None):
+        """
+        Initializes the configuration for explainable quantization.
+
+        Args:
+            report_dir (str): Directory where the reports will be saved.
+            custom_similarity_metrics (Dict[str, Callable]): Custom similarity metrics to be computed between tensors
+            of the two models. The dictionary keys are similarity metric names and the values are callables that implement the
+            similarity metric computation.
+        """
+        self.report_dir = report_dir
+        self.custom_similarity_metrics = custom_similarity_metrics

model_compression_toolkit/xquant/keras/__init__.py

@@ -0,0 +1,15 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+

model_compression_toolkit/xquant/keras/dataset_utils.py

@@ -0,0 +1,57 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+
+from typing import Any
+
+from model_compression_toolkit.xquant.common.dataset_utils import DatasetUtils
+
+
+class KerasDatasetUtils(DatasetUtils):
+    """
+    Class with helpful methods for handling different kinds of Keras datasets from the user.
+    """
+
+    @staticmethod
+    def prepare_dataset(dataset: Any, is_validation: bool, device: str = None):
+        """
+        Prepare the dataset so calling it will return only inputs for the model (like in the case
+        of the representative dataset). For example, when the validation dataset is used, the labels
+        should be removed.
+
+        Args:
+            dataset: Dataset to prepare.
+            is_validation: Whether it's validation dataset or not.
+            device: Device to transfer the data to.
+
+        Returns:
+            Generator to use for retrieving the dataset inputs.
+        """
+        def process_data(x: Any, is_validation: bool):
+            """
+            Processes individual data samples to transfer them to the device and convert to torch tensors if needed.
+
+            Args:
+                data: The data sample to process.
+                is_validation: A flag indicating if this is a validation dataset.
+                device: The device to transfer the data to.
+
+            Returns:
+                The data as torch tensors on the desired device.
+            """
+            return x[0] if is_validation else x  # Assume data[0] contains the inputs and data[1] the labels
+
+        for x in dataset():
+            yield process_data(x, is_validation)
+

model_compression_toolkit/xquant/keras/facade_xquant_report.py

@@ -0,0 +1,63 @@
+#  Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  ==============================================================================
+
+from typing import Callable, Dict, Any
+
+from model_compression_toolkit.constants import FOUND_TF
+from model_compression_toolkit.xquant.common.core_report_generator import core_report_generator
+from model_compression_toolkit.xquant import XQuantConfig
+from model_compression_toolkit.logger import Logger
+
+if FOUND_TF:
+    import keras
+    from model_compression_toolkit.xquant.keras.keras_report_utils import KerasReportUtils
+
+    def xquant_report_keras_experimental(float_model: keras.Model,
+                                         quantized_model: keras.Model,
+                                         repr_dataset: Callable,
+                                         validation_dataset: Callable,
+                                         xquant_config: XQuantConfig) -> Dict[str, Any]:
+        """
+        Generate an explainable quantization report for a quantized Keras model.
+
+        Args:
+            float_model (keras.Model): The original floating-point Keras model.
+            quantized_model (keras.Model): The quantized Keras model.
+            repr_dataset (Callable): The representative dataset used during quantization for similarity metrics computation.
+            validation_dataset (Callable): The validation dataset used for evaluation for similarity metrics computation.
+            xquant_config (XQuantConfig): Configuration settings for explainable quantization.
+
+        Returns:
+            Dict[str, Any]: A dictionary containing the collected similarity metrics and report data.
+        """
+        # Initialize the logger with the report directory.
+        Logger.set_log_file(log_folder=xquant_config.report_dir)
+
+        # Initialize a utility class for handling Keras-specific reporting tasks.
+        keras_report_utils = KerasReportUtils(xquant_config.report_dir)
+
+        # Create the report after collecting useful data like histograms and similarity metrics.
+        _collected_data = core_report_generator(float_model=float_model,
+                                                quantized_model=quantized_model,
+                                                repr_dataset=repr_dataset,
+                                                validation_dataset=validation_dataset,
+                                                fw_report_utils=keras_report_utils,
+                                                xquant_config=xquant_config)
+
+        return _collected_data
+else:
+    def xquant_report_keras_experimental(*args, **kwargs):
+        Logger.critical("Tensorflow must be installed to use xquant_report_keras_experimental. "
+                                     "The 'tensorflow' package is missing.")  # pragma: no cover