1# SPDX-FileCopyrightText: Copyright 2024 Arm Limited and/or its affiliates <open-source-office@arm.com>
2#
3# SPDX-License-Identifier: Apache-2.0
4#
5# Licensed under the Apache License, Version 2.0 (the License); you may
6# not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an AS IS BASIS, WITHOUT
13# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17import Lib.op_utils
18import tensorflow as tf
19import math
20import numpy as np
21
22from tensorflow.lite.python.interpreter import Interpreter
23from tensorflow.lite.python.interpreter import OpResolverType
24import tf_keras as keras
25
26class Op_maximum_minimum(Lib.op_utils.Op_type):
27
28    def get_shapes(params):
29        shapes = {}
30        shapes["input_tensor_1"] = (params["batch_1"], params["height_1"], params["width_1"], params["channel_1"])
31        shapes["input_tensor_2"] = (params["batch_2"], params["height_2"], params["width_2"], params["channel_2"])
32        shapes["representational_dataset"] = (params["batch_1"], params["height_1"], params["width_1"], params["channel_1"])
33        shapes["representational_dataset2"] = (params["batch_2"], params["height_2"], params["width_2"], params["channel_2"])
34        shapes["different_in_shapes"]=True
35
36        return shapes
37
38    def generate_keras_model(shapes, params):
39        tf.keras.backend.clear_session()
40        layer_type = None
41        if params['layer_type'] == "minimum":
42            layer_type = tf.minimum
43        elif params['layer_type'] == "maximum":
44            layer_type = tf.maximum
45
46        input_1_shape = (params["batch_1"], params["height_1"], params["width_1"], params["channel_1"])
47        input_2_shape = (params["batch_2"], params["height_2"], params["width_2"], params["channel_2"])
48        input_1 = keras.layers.Input(batch_input_shape=input_1_shape)
49        input_2 = keras.layers.Input(batch_input_shape=input_2_shape)
50
51        layer = layer_type(input_1, input_2)
52        model = keras.Model([input_1, input_2], [layer])
53
54        return model
55
56    def generate_data_tflite(tflite_fname, params):
57        tensors = {}
58        effective_scales = {}
59        scales = {}
60        generated_params = {}
61        aliases = {}
62
63        # To be removed
64        aliases["output_multiplier"] = "output_mult"
65        aliases["output"] = "output_ref"
66
67        interpreter = Interpreter(str(tflite_fname), experimental_op_resolver_type=OpResolverType.BUILTIN_REF)
68        interpreter.allocate_tensors()
69        tensor_details = interpreter.get_tensor_details()
70
71        input_1 = tensor_details[0]
72        input_2 = tensor_details[1]
73
74        input_details = interpreter.get_input_details()
75        (scales["scale_1"], scales["zero_point_1"]) = input_details[0]['quantization']
76        (scales["scale_2"], scales["zero_point2"]) = input_details[1]['quantization']
77
78        output_details = interpreter.get_output_details()
79        (scales["output_scale"], scales["output_zero_point"]) = output_details[0]['quantization']
80
81        minval = Lib.op_utils.get_dtype_min(params["input_data_type"])
82        maxval = Lib.op_utils.get_dtype_max(params["input_data_type"])
83
84        n_output = output_details[0]['shape'][0]
85        h_output = output_details[0]['shape'][1]
86        w_output = output_details[0]['shape'][2]
87        c_output = output_details[0]['shape'][3]
88
89        generated_params["dst_size"] = n_output * h_output * w_output * c_output
90        generated_params["output_batch"] = n_output
91        generated_params["output_height"] = h_output
92        generated_params["output_width"] = w_output
93        generated_params["output_channel"] = c_output
94        generated_params["input_1_offset"] = -input_1['quantization_parameters']['zero_points'][0]
95        generated_params["input_2_offset"] = -input_2['quantization_parameters']['zero_points'][0]
96        generated_params["output_offset"] = output_details[0]['quantization'][1]
97
98        def quantize_scale(scales):
99            effective_output_scale = scales["scale_1"] * scales["scale_2"] / scales["output_scale"]
100
101            significand, shift = math.frexp(effective_output_scale)
102            significand_q31 = round(significand * (1 << 31))
103            return significand_q31, shift
104
105        mult, shift = quantize_scale(scales)
106        generated_params["output_multiplier"] = mult
107        generated_params["output_shift"] = shift
108
109        return Lib.op_utils.Generated_data(generated_params, tensors, scales, effective_scales, aliases)
110
111