aboutsummaryrefslogtreecommitdiff
path: root/src/mlia/nn/rewrite/library/fc_clustering_layer.py
blob: 07c07aca2b04cb5f9ea1aa58e00f24fabc0552d1 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
# SPDX-FileCopyrightText: Copyright 2024, Arm Limited and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
"""Example rewrite with one fully connected clustered layer."""
from typing import Any

from keras.api._v2 import keras  # Temporary workaround for now: MLIA-1107

from mlia.nn.rewrite.library.fc_layer import get_keras_model


def get_keras_model_clus(input_shape: Any, output_shape: Any) -> keras.Model:
    """Generate TensorFlow Lite model for clustering rewrite."""
    return get_keras_model(input_shape, output_shape)