blob: e12267f05e09480ab5ffbd70b4530175b90fe617 [file] [log] [blame]
erik.andersson@arm.com460c6892021-02-24 14:38:09 +01001# Copyright (C) 2020-2021 Arm Limited or its affiliates. All rights reserved.
Tim Hall79d07d22020-04-27 18:20:16 +01002#
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.
Tim Hall79d07d22020-04-27 18:20:16 +010016# Description:
17# TensorFlow Lite mapping functions used by both reader and writer.
18# Contains a mapping from the various TensorFlow Lite enums and options structs, generated by the FlatBuffer code
19# generator, to Vela's internal format.
Tim Hall79d07d22020-04-27 18:20:16 +010020import struct
21
Diego Russoea6111a2020-04-14 18:41:58 +010022import numpy as np
23
Tim Hall79d07d22020-04-27 18:20:16 +010024from .data_type import DataType
Louis Verhaardaee5d752020-09-30 09:01:52 +020025from .operation import CustomType
26from .operation import Op
Michael McGeagh16895482020-12-14 15:51:20 +000027from .operation import Padding as opPad
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +020028from .operation import TensorIndices
Diego Russoe8a10452020-04-21 17:39:10 +010029from .tflite import AbsOptions
30from .tflite import AddNOptions
Tim Hall79d07d22020-04-27 18:20:16 +010031from .tflite import AddOptions
Tim Hall79d07d22020-04-27 18:20:16 +010032from .tflite import ArgMaxOptions
Tim Hall79d07d22020-04-27 18:20:16 +010033from .tflite import ArgMinOptions
Rickard Bolin2de898a2021-12-20 08:35:23 +000034from .tflite import AssignVariableOptions
Jacob Bohlin8daf6b72020-09-15 16:28:35 +020035from .tflite import BatchMatMulOptions
Diego Russoe8a10452020-04-21 17:39:10 +010036from .tflite import BatchToSpaceNDOptions
Tim Hall79d07d22020-04-27 18:20:16 +010037from .tflite import BidirectionalSequenceLSTMOptions
38from .tflite import BidirectionalSequenceRNNOptions
Dwight Lidman8a12da12021-07-19 13:43:05 +020039from .tflite import BroadcastToOptions
40from .tflite import CallOnceOptions
Diego Russoe8a10452020-04-21 17:39:10 +010041from .tflite import CallOptions
42from .tflite import CastOptions
43from .tflite import ConcatEmbeddingsOptions
44from .tflite import ConcatenationOptions
45from .tflite import Conv2DOptions
Dwight Lidman8a12da12021-07-19 13:43:05 +020046from .tflite import Conv3DOptions
Tim Hall79d07d22020-04-27 18:20:16 +010047from .tflite import CosOptions
Tim Hall42abec12021-02-04 21:31:57 +000048from .tflite import CumsumOptions
Tim Hall79d07d22020-04-27 18:20:16 +010049from .tflite import DensifyOptions
50from .tflite import DepthToSpaceOptions
Diego Russoe8a10452020-04-21 17:39:10 +010051from .tflite import DepthwiseConv2DOptions
52from .tflite import DequantizeOptions
53from .tflite import DivOptions
54from .tflite import EmbeddingLookupSparseOptions
55from .tflite import EqualOptions
56from .tflite import ExpandDimsOptions
57from .tflite import ExpOptions
58from .tflite import FakeQuantOptions
59from .tflite import FillOptions
60from .tflite import FloorDivOptions
61from .tflite import FloorModOptions
62from .tflite import FullyConnectedOptions
63from .tflite import GatherNdOptions
64from .tflite import GatherOptions
65from .tflite import GreaterEqualOptions
66from .tflite import GreaterOptions
Louis Verhaard678645b2020-06-15 15:22:47 +020067from .tflite import HardSwishOptions
Dwight Lidman8a12da12021-07-19 13:43:05 +020068from .tflite import HashtableFindOptions
69from .tflite import HashtableImportOptions
70from .tflite import HashtableOptions
71from .tflite import HashtableSizeOptions
Tim Hall79d07d22020-04-27 18:20:16 +010072from .tflite import IfOptions
Diego Russoe8a10452020-04-21 17:39:10 +010073from .tflite import L2NormOptions
74from .tflite import LeakyReluOptions
75from .tflite import LessEqualOptions
76from .tflite import LessOptions
77from .tflite import LocalResponseNormalizationOptions
78from .tflite import LogicalAndOptions
79from .tflite import LogicalNotOptions
80from .tflite import LogicalOrOptions
81from .tflite import LogSoftmaxOptions
82from .tflite import LSHProjectionOptions
83from .tflite import LSTMOptions
84from .tflite import MatrixDiagOptions
85from .tflite import MatrixSetDiagOptions
86from .tflite import MaximumMinimumOptions
87from .tflite import MirrorPadOptions
88from .tflite import MulOptions
89from .tflite import NegOptions
Tim Hall79d07d22020-04-27 18:20:16 +010090from .tflite import NonMaxSuppressionV4Options
91from .tflite import NonMaxSuppressionV5Options
Diego Russoe8a10452020-04-21 17:39:10 +010092from .tflite import NotEqualOptions
93from .tflite import OneHotOptions
94from .tflite import PackOptions
95from .tflite import PadOptions
96from .tflite import PadV2Options
97from .tflite import Pool2DOptions
98from .tflite import PowOptions
99from .tflite import QuantizeOptions
Rickard Bolin2de898a2021-12-20 08:35:23 +0000100from .tflite import RandomOptions
Diego Russoe8a10452020-04-21 17:39:10 +0100101from .tflite import RangeOptions
102from .tflite import RankOptions
Rickard Bolin2de898a2021-12-20 08:35:23 +0000103from .tflite import ReadVariableOptions
Diego Russoe8a10452020-04-21 17:39:10 +0100104from .tflite import ReducerOptions
105from .tflite import ReshapeOptions
106from .tflite import ResizeBilinearOptions
107from .tflite import ResizeNearestNeighborOptions
108from .tflite import ReverseSequenceOptions
109from .tflite import ReverseV2Options
Dwight Lidman8a12da12021-07-19 13:43:05 +0200110from .tflite import Rfft2dOptions
Diego Russoe8a10452020-04-21 17:39:10 +0100111from .tflite import RNNOptions
Tim Hall79d07d22020-04-27 18:20:16 +0100112from .tflite import ScatterNdOptions
113from .tflite import SegmentSumOptions
Diego Russoe8a10452020-04-21 17:39:10 +0100114from .tflite import SelectOptions
Tim Hall79d07d22020-04-27 18:20:16 +0100115from .tflite import SelectV2Options
Diego Russoe8a10452020-04-21 17:39:10 +0100116from .tflite import SequenceRNNOptions
117from .tflite import ShapeOptions
118from .tflite import SkipGramOptions
119from .tflite import SliceOptions
120from .tflite import SoftmaxOptions
121from .tflite import SpaceToBatchNDOptions
122from .tflite import SpaceToDepthOptions
123from .tflite import SparseToDenseOptions
124from .tflite import SplitOptions
125from .tflite import SplitVOptions
126from .tflite import SquaredDifferenceOptions
127from .tflite import SquareOptions
128from .tflite import SqueezeOptions
129from .tflite import StridedSliceOptions
130from .tflite import SubOptions
131from .tflite import SVDFOptions
132from .tflite import TileOptions
133from .tflite import TopKV2Options
134from .tflite import TransposeConvOptions
135from .tflite import TransposeOptions
136from .tflite import UnidirectionalSequenceLSTMOptions
137from .tflite import UniqueOptions
138from .tflite import UnpackOptions
Rickard Bolin2de898a2021-12-20 08:35:23 +0000139from .tflite import VarHandleOptions
Diego Russoe8a10452020-04-21 17:39:10 +0100140from .tflite import WhereOptions
Tim Hall79d07d22020-04-27 18:20:16 +0100141from .tflite import WhileOptions
Diego Russoe8a10452020-04-21 17:39:10 +0100142from .tflite import ZerosLikeOptions
143from .tflite.ActivationFunctionType import ActivationFunctionType
Diego Russoea6111a2020-04-14 18:41:58 +0100144from .tflite.BuiltinOperator import BuiltinOperator
145from .tflite.BuiltinOptions import BuiltinOptions
146from .tflite.Padding import Padding
Diego Russoe8a10452020-04-21 17:39:10 +0100147from .tflite.TensorType import TensorType
Tim Hall79d07d22020-04-27 18:20:16 +0100148
149
150def inverse_map(map):
151 return {v: k for k, v in map.items()}
152
153
154datatype_map = {
155 TensorType.UINT8: DataType.uint8,
156 TensorType.INT8: DataType.int8,
157 TensorType.INT16: DataType.int16,
158 TensorType.INT32: DataType.int32,
159 TensorType.INT64: DataType.int64,
160 TensorType.FLOAT16: DataType.float16,
161 TensorType.FLOAT32: DataType.float32,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200162 TensorType.FLOAT64: DataType.float64,
Tim Hall79d07d22020-04-27 18:20:16 +0100163 TensorType.STRING: DataType.string,
164 TensorType.BOOL: DataType.bool,
Jacob Bohlinf767b932020-08-13 15:32:45 +0200165 TensorType.COMPLEX64: DataType.complex64,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200166 TensorType.COMPLEX128: DataType.complex128,
Dwight Lidmanf53b3322021-10-08 11:58:08 +0200167 TensorType.UINT64: DataType.uint64,
168 TensorType.RESOURCE: DataType.resource,
169 TensorType.VARIANT: DataType.variant,
170 TensorType.UINT32: DataType.uint32,
Tim Hall79d07d22020-04-27 18:20:16 +0100171}
172
173datatype_inv_map = inverse_map(datatype_map)
174datatype_inv_map[DataType.quint8] = TensorType.UINT8
175
176datatype_inv_map[DataType.qint8] = TensorType.INT8
177datatype_inv_map[DataType.qint16] = TensorType.INT16
178datatype_inv_map[DataType.qint32] = TensorType.INT32
179
180
181datatype_map_numpy = {
182 TensorType.UINT8: np.uint8,
Dwight Lidmanf53b3322021-10-08 11:58:08 +0200183 TensorType.UINT32: np.uint32,
184 TensorType.UINT64: np.uint64,
Tim Hall79d07d22020-04-27 18:20:16 +0100185 TensorType.INT8: np.int8,
186 TensorType.INT16: np.int16,
187 TensorType.INT32: np.int32,
188 TensorType.INT64: np.int64,
189 TensorType.FLOAT16: np.float16,
190 TensorType.FLOAT32: np.float32,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200191 TensorType.FLOAT64: np.float64,
Tim Hall79d07d22020-04-27 18:20:16 +0100192 TensorType.BOOL: np.bool,
Jacob Bohlinf767b932020-08-13 15:32:45 +0200193 TensorType.COMPLEX64: np.complex64,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200194 TensorType.COMPLEX128: np.complex128,
Louis Verhaardf4e12be2020-12-18 14:23:06 +0100195 TensorType.STRING: np.uint8,
Dwight Lidmanf53b3322021-10-08 11:58:08 +0200196 TensorType.RESOURCE: None, # experimental type
197 TensorType.VARIANT: None, # experimental type
Tim Hall79d07d22020-04-27 18:20:16 +0100198}
199
200
201builtin_options_map = {
202 BuiltinOptions.Conv2DOptions: Conv2DOptions.Conv2DOptions,
203 BuiltinOptions.DepthwiseConv2DOptions: DepthwiseConv2DOptions.DepthwiseConv2DOptions,
204 BuiltinOptions.ConcatEmbeddingsOptions: ConcatEmbeddingsOptions.ConcatEmbeddingsOptions,
205 BuiltinOptions.LSHProjectionOptions: LSHProjectionOptions.LSHProjectionOptions,
206 BuiltinOptions.Pool2DOptions: Pool2DOptions.Pool2DOptions,
207 BuiltinOptions.SVDFOptions: SVDFOptions.SVDFOptions,
208 BuiltinOptions.RNNOptions: RNNOptions.RNNOptions,
209 BuiltinOptions.FullyConnectedOptions: FullyConnectedOptions.FullyConnectedOptions,
210 BuiltinOptions.SoftmaxOptions: SoftmaxOptions.SoftmaxOptions,
211 BuiltinOptions.ConcatenationOptions: ConcatenationOptions.ConcatenationOptions,
212 BuiltinOptions.AddOptions: AddOptions.AddOptions,
213 BuiltinOptions.L2NormOptions: L2NormOptions.L2NormOptions,
214 BuiltinOptions.LocalResponseNormalizationOptions: LocalResponseNormalizationOptions.LocalResponseNormalizationOptions, # noqa: E501
215 BuiltinOptions.LSTMOptions: LSTMOptions.LSTMOptions,
216 BuiltinOptions.ResizeBilinearOptions: ResizeBilinearOptions.ResizeBilinearOptions,
217 BuiltinOptions.CallOptions: CallOptions.CallOptions,
218 BuiltinOptions.ReshapeOptions: ReshapeOptions.ReshapeOptions,
219 BuiltinOptions.SkipGramOptions: SkipGramOptions.SkipGramOptions,
220 BuiltinOptions.SpaceToDepthOptions: SpaceToDepthOptions.SpaceToDepthOptions,
221 BuiltinOptions.EmbeddingLookupSparseOptions: EmbeddingLookupSparseOptions.EmbeddingLookupSparseOptions,
222 BuiltinOptions.MulOptions: MulOptions.MulOptions,
223 BuiltinOptions.PadOptions: PadOptions.PadOptions,
224 BuiltinOptions.GatherOptions: GatherOptions.GatherOptions,
225 BuiltinOptions.BatchToSpaceNDOptions: BatchToSpaceNDOptions.BatchToSpaceNDOptions,
226 BuiltinOptions.SpaceToBatchNDOptions: SpaceToBatchNDOptions.SpaceToBatchNDOptions,
227 BuiltinOptions.TransposeOptions: TransposeOptions.TransposeOptions,
228 BuiltinOptions.ReducerOptions: ReducerOptions.ReducerOptions,
229 BuiltinOptions.SubOptions: SubOptions.SubOptions,
230 BuiltinOptions.DivOptions: DivOptions.DivOptions,
231 BuiltinOptions.SqueezeOptions: SqueezeOptions.SqueezeOptions,
232 BuiltinOptions.SequenceRNNOptions: SequenceRNNOptions.SequenceRNNOptions,
233 BuiltinOptions.StridedSliceOptions: StridedSliceOptions.StridedSliceOptions,
234 BuiltinOptions.ExpOptions: ExpOptions.ExpOptions,
235 BuiltinOptions.TopKV2Options: TopKV2Options.TopKV2Options,
236 BuiltinOptions.SplitOptions: SplitOptions.SplitOptions,
237 BuiltinOptions.LogSoftmaxOptions: LogSoftmaxOptions.LogSoftmaxOptions,
238 BuiltinOptions.CastOptions: CastOptions.CastOptions,
239 BuiltinOptions.DequantizeOptions: DequantizeOptions.DequantizeOptions,
240 BuiltinOptions.MaximumMinimumOptions: MaximumMinimumOptions.MaximumMinimumOptions,
241 BuiltinOptions.ArgMaxOptions: ArgMaxOptions.ArgMaxOptions,
242 BuiltinOptions.LessOptions: LessOptions.LessOptions,
243 BuiltinOptions.NegOptions: NegOptions.NegOptions,
244 BuiltinOptions.PadV2Options: PadV2Options.PadV2Options,
245 BuiltinOptions.GreaterOptions: GreaterOptions.GreaterOptions,
246 BuiltinOptions.GreaterEqualOptions: GreaterEqualOptions.GreaterEqualOptions,
247 BuiltinOptions.LessEqualOptions: LessEqualOptions.LessEqualOptions,
248 BuiltinOptions.SelectOptions: SelectOptions.SelectOptions,
249 BuiltinOptions.SliceOptions: SliceOptions.SliceOptions,
250 BuiltinOptions.TransposeConvOptions: TransposeConvOptions.TransposeConvOptions,
251 BuiltinOptions.SparseToDenseOptions: SparseToDenseOptions.SparseToDenseOptions,
252 BuiltinOptions.TileOptions: TileOptions.TileOptions,
253 BuiltinOptions.ExpandDimsOptions: ExpandDimsOptions.ExpandDimsOptions,
254 BuiltinOptions.EqualOptions: EqualOptions.EqualOptions,
255 BuiltinOptions.NotEqualOptions: NotEqualOptions.NotEqualOptions,
256 BuiltinOptions.ShapeOptions: ShapeOptions.ShapeOptions,
257 BuiltinOptions.PowOptions: PowOptions.PowOptions,
258 BuiltinOptions.ArgMinOptions: ArgMinOptions.ArgMinOptions,
259 BuiltinOptions.FakeQuantOptions: FakeQuantOptions.FakeQuantOptions,
260 BuiltinOptions.PackOptions: PackOptions.PackOptions,
261 BuiltinOptions.LogicalOrOptions: LogicalOrOptions.LogicalOrOptions,
262 BuiltinOptions.OneHotOptions: OneHotOptions.OneHotOptions,
263 BuiltinOptions.LogicalAndOptions: LogicalAndOptions.LogicalAndOptions,
264 BuiltinOptions.LogicalNotOptions: LogicalNotOptions.LogicalNotOptions,
265 BuiltinOptions.UnpackOptions: UnpackOptions.UnpackOptions,
266 BuiltinOptions.FloorDivOptions: FloorDivOptions.FloorDivOptions,
267 BuiltinOptions.SquareOptions: SquareOptions.SquareOptions,
268 BuiltinOptions.ZerosLikeOptions: ZerosLikeOptions.ZerosLikeOptions,
269 BuiltinOptions.FillOptions: FillOptions.FillOptions,
270 BuiltinOptions.BidirectionalSequenceLSTMOptions: BidirectionalSequenceLSTMOptions.BidirectionalSequenceLSTMOptions,
271 BuiltinOptions.BidirectionalSequenceRNNOptions: BidirectionalSequenceRNNOptions.BidirectionalSequenceRNNOptions,
272 BuiltinOptions.UnidirectionalSequenceLSTMOptions: UnidirectionalSequenceLSTMOptions.UnidirectionalSequenceLSTMOptions, # noqa: E501
273 BuiltinOptions.FloorModOptions: FloorModOptions.FloorModOptions,
274 BuiltinOptions.RangeOptions: RangeOptions.RangeOptions,
275 BuiltinOptions.ResizeNearestNeighborOptions: ResizeNearestNeighborOptions.ResizeNearestNeighborOptions,
276 BuiltinOptions.LeakyReluOptions: LeakyReluOptions.LeakyReluOptions,
277 BuiltinOptions.SquaredDifferenceOptions: SquaredDifferenceOptions.SquaredDifferenceOptions,
278 BuiltinOptions.MirrorPadOptions: MirrorPadOptions.MirrorPadOptions,
279 BuiltinOptions.AbsOptions: AbsOptions.AbsOptions,
280 BuiltinOptions.SplitVOptions: SplitVOptions.SplitVOptions,
281 BuiltinOptions.UniqueOptions: UniqueOptions.UniqueOptions,
282 BuiltinOptions.ReverseV2Options: ReverseV2Options.ReverseV2Options,
283 BuiltinOptions.AddNOptions: AddNOptions.AddNOptions,
284 BuiltinOptions.GatherNdOptions: GatherNdOptions.GatherNdOptions,
285 BuiltinOptions.CosOptions: CosOptions.CosOptions,
286 BuiltinOptions.WhereOptions: WhereOptions.WhereOptions,
287 BuiltinOptions.RankOptions: RankOptions.RankOptions,
288 BuiltinOptions.ReverseSequenceOptions: ReverseSequenceOptions.ReverseSequenceOptions,
289 BuiltinOptions.MatrixDiagOptions: MatrixDiagOptions.MatrixDiagOptions,
290 BuiltinOptions.QuantizeOptions: QuantizeOptions.QuantizeOptions,
291 BuiltinOptions.MatrixSetDiagOptions: MatrixSetDiagOptions.MatrixSetDiagOptions,
292 BuiltinOptions.DensifyOptions: DensifyOptions.DensifyOptions,
293 BuiltinOptions.DepthToSpaceOptions: DepthToSpaceOptions.DepthToSpaceOptions,
Louis Verhaard678645b2020-06-15 15:22:47 +0200294 BuiltinOptions.HardSwishOptions: HardSwishOptions.HardSwishOptions,
Tim Hall79d07d22020-04-27 18:20:16 +0100295 BuiltinOptions.IfOptions: IfOptions.IfOptions,
296 BuiltinOptions.NonMaxSuppressionV4Options: NonMaxSuppressionV4Options.NonMaxSuppressionV4Options,
297 BuiltinOptions.NonMaxSuppressionV5Options: NonMaxSuppressionV5Options.NonMaxSuppressionV5Options,
298 BuiltinOptions.ScatterNdOptions: ScatterNdOptions.ScatterNdOptions,
299 BuiltinOptions.SegmentSumOptions: SegmentSumOptions.SegmentSumOptions,
300 BuiltinOptions.SelectV2Options: SelectV2Options.SelectV2Options,
301 BuiltinOptions.WhileOptions: WhileOptions.WhileOptions,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200302 BuiltinOptions.BatchMatMulOptions: BatchMatMulOptions.BatchMatMulOptions,
Tim Hall42abec12021-02-04 21:31:57 +0000303 BuiltinOptions.CumsumOptions: CumsumOptions.CumsumOptions,
Dwight Lidman8a12da12021-07-19 13:43:05 +0200304 BuiltinOptions.CallOnceOptions: CallOnceOptions.CallOnceOptions,
305 BuiltinOptions.BroadcastToOptions: BroadcastToOptions.BroadcastToOptions,
306 BuiltinOptions.Rfft2dOptions: Rfft2dOptions.Rfft2dOptions,
307 BuiltinOptions.Conv3DOptions: Conv3DOptions.Conv3DOptions,
308 BuiltinOptions.HashtableOptions: HashtableOptions.HashtableOptions,
309 BuiltinOptions.HashtableFindOptions: HashtableFindOptions.HashtableFindOptions,
310 BuiltinOptions.HashtableImportOptions: HashtableImportOptions.HashtableImportOptions,
311 BuiltinOptions.HashtableSizeOptions: HashtableSizeOptions.HashtableSizeOptions,
Rickard Bolin2de898a2021-12-20 08:35:23 +0000312 BuiltinOptions.VarHandleOptions: VarHandleOptions.VarHandleOptions,
313 BuiltinOptions.ReadVariableOptions: ReadVariableOptions.ReadVariableOptions,
314 BuiltinOptions.AssignVariableOptions: AssignVariableOptions.AssignVariableOptions,
315 BuiltinOptions.RandomOptions: RandomOptions.RandomOptions,
Tim Hall79d07d22020-04-27 18:20:16 +0100316}
317
Rickard Bolin2de898a2021-12-20 08:35:23 +0000318
Tim Hall79d07d22020-04-27 18:20:16 +0100319builtin_options_inv_map = inverse_map(builtin_options_map)
320
321
322def underscore_to_camel_case(s):
323 return "".join(x.title() for x in s.split("_"))
324
325
326def padding_deserialize(x):
327 return padding_map[x]
328
329
330def padding_serialize(builder, x):
331 return padding_inv_map[x]
332
333
334def activation_deserialize(x):
335 return activation_function_map[x]
336
337
338def activation_serialize(builder, x):
339 return activation_function_inv_map[x]
340
341
342def datatype_deserialize(x):
343 return datatype_map[x]
344
345
346def datatype_serialize(builder, x):
347 return datatype_inv_map[x]
348
349
350def identity(x):
351 return x
352
353
354def identity_serialize(builder, x):
355 return x
356
357
358def write_byte_vector(builder, v):
359 builder.StartVector(1, len(v), 1)
360 for e in v[::-1]:
361 builder.PrependByte(e)
362 return builder.EndVector(len(v))
363
364
365def write_int_vector(builder, v):
366 builder.StartVector(4, len(v), 4)
367 for e in v[::-1]:
368 builder.PrependInt32(e)
369 return builder.EndVector(len(v))
370
371
372class OptionsSerializer:
Michael McGeagh6f725262020-12-03 15:21:36 +0000373 def __init__(self, name, members=None):
Tim Hall79d07d22020-04-27 18:20:16 +0100374 self.name = name
375 self.module = globals()[self.name]
376 self.cls = getattr(self.module, self.name)
377 self.builtin_opt_type = builtin_options_inv_map[self.cls]
Tim Hall79d07d22020-04-27 18:20:16 +0100378 self.members = []
Michael McGeagh6f725262020-12-03 15:21:36 +0000379 if members is not None:
380 for mem in members:
381 deserialize = identity
382 serialize = identity_serialize
383 is_vector = False
384 if isinstance(mem, tuple):
385 if len(mem) == 3:
386 mem, deserialize, serialize = mem
387 elif len(mem) == 2:
388 mem, is_vector = mem
389 deserialize = tuple
390 serialize = write_int_vector
391 else:
392 assert 0
393 underscore_mem = mem
394 camelcase_mem = underscore_to_camel_case(mem)
395 self.members.append((underscore_mem, camelcase_mem, deserialize, serialize, is_vector))
Tim Hall79d07d22020-04-27 18:20:16 +0100396
Tim Hallc8310b12020-06-17 14:53:11 +0100397 def deserialize(self, op_data):
398 builtin_options = op_data.BuiltinOptions()
Tim Hall79d07d22020-04-27 18:20:16 +0100399 attrs = {}
Tim Hallc8310b12020-06-17 14:53:11 +0100400 if builtin_options:
Tim Hall79d07d22020-04-27 18:20:16 +0100401 tfattrs = self.cls()
Tim Hallc8310b12020-06-17 14:53:11 +0100402 tfattrs.Init(builtin_options.Bytes, builtin_options.Pos)
Tim Hall79d07d22020-04-27 18:20:16 +0100403 for underscore_mem, camelcase_mem, deserialize, serialize, is_vector in self.members:
404 fun = camelcase_mem
405 if is_vector:
406 fun += "AsNumpy"
407
Fredrik Svedberg4bb989b2020-11-16 15:32:25 +0100408 attr = getattr(tfattrs, fun)()
409 try:
410 attrs[underscore_mem] = deserialize(attr)
411 except TypeError:
412 print("Warning: {0} could not read attribute '{1}'.".format(self.name, underscore_mem))
413
Tim Hall79d07d22020-04-27 18:20:16 +0100414 return attrs
415
416 def serialize(self, builder, attrs):
417 ser_attrs = []
418 for underscore_mem, camelcase_mem, deserialize, serialize, is_vector in self.members:
419 a = serialize(builder, attrs[underscore_mem])
420 ser_attrs.append((camelcase_mem, a))
421
422 getattr(self.module, self.name + "Start")(builder)
423
424 for camelcase_mem, a in ser_attrs:
425 getattr(self.module, self.name + "Add" + camelcase_mem)(builder, a)
426
427 return getattr(self.module, self.name + "End")(builder), None
428
429
430class CustomOptionsSerializer:
Tim Hallc8310b12020-06-17 14:53:11 +0100431 CUSTOM_OPTIONS_NPU_OP = [0x01, 0x04, 0x01] # NpuOp=1, FlexbufferFormat.UINT8=4, byte length=1
432 CUSTOM_OPTIONS_FORMAT_DEFAULT = 0
433
Tim Hall79d07d22020-04-27 18:20:16 +0100434 def __init__(self):
Tim Hall79d07d22020-04-27 18:20:16 +0100435 self.custom_opt_format = 0
436
Tim Hallc8310b12020-06-17 14:53:11 +0100437 def deserialize(self, op_data):
Tim Hall79d07d22020-04-27 18:20:16 +0100438 attrs = {}
Tim Hallc8310b12020-06-17 14:53:11 +0100439 custom_options = op_data.CustomOptionsAsNumpy()
440 attrs["custom_options"] = custom_options
441 attrs["custom_options_format"] = op_data.CustomOptionsFormat()
442
443 if np.array_equal(custom_options, self.CUSTOM_OPTIONS_NPU_OP):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200444 attrs["custom_type"] = CustomType.ExistingNpuOp
Tim Hallc8310b12020-06-17 14:53:11 +0100445
Tim Hall79d07d22020-04-27 18:20:16 +0100446 return attrs
447
448 def serialize(self, builder, attrs):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200449 custom_type = attrs.get("custom_type", CustomType.ThirdPartyOp)
Tim Hallc8310b12020-06-17 14:53:11 +0100450 self.custom_opt_format = attrs.get("custom_options_format", self.CUSTOM_OPTIONS_FORMAT_DEFAULT)
Tim Hall79d07d22020-04-27 18:20:16 +0100451
452 # Set NPU op custom options for the TensorFlow Lite custom operator
Louis Verhaardaee5d752020-09-30 09:01:52 +0200453 if custom_type == CustomType.NpuOp:
Tim Hallc8310b12020-06-17 14:53:11 +0100454 custom_options = self.CUSTOM_OPTIONS_NPU_OP
455 else:
456 custom_options = attrs.get("custom_options", [])
Tim Hall79d07d22020-04-27 18:20:16 +0100457
Tim Hallc8310b12020-06-17 14:53:11 +0100458 custom_options_bytes = struct.pack("<{0}B".format(len(custom_options)), *custom_options)
459 custom_offset = write_byte_vector(builder, custom_options_bytes)
Tim Hall79d07d22020-04-27 18:20:16 +0100460
461 return None, custom_offset
462
463
464padding_map = {
Michael McGeagh16895482020-12-14 15:51:20 +0000465 Padding.SAME: opPad.SAME,
466 Padding.VALID: opPad.VALID,
Tim Hall79d07d22020-04-27 18:20:16 +0100467}
468
469padding_inv_map = inverse_map(padding_map)
470
471
472activation_function_map = {
473 ActivationFunctionType.NONE: None,
Louis Verhaardaee5d752020-09-30 09:01:52 +0200474 ActivationFunctionType.RELU: Op.Relu,
475 ActivationFunctionType.RELU_N1_TO_1: Op.ReluN1To1,
476 ActivationFunctionType.RELU6: Op.Relu6,
477 ActivationFunctionType.TANH: Op.Tanh,
478 ActivationFunctionType.SIGN_BIT: Op.SignBit,
Tim Hall79d07d22020-04-27 18:20:16 +0100479}
480
481activation_function_inv_map = inverse_map(activation_function_map)
482
483fused_act = ("fused_activation_function", activation_deserialize, activation_serialize)
484padding = ("padding", padding_deserialize, padding_serialize)
485
Tim Hall79d07d22020-04-27 18:20:16 +0100486is_int_vec = True
487
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200488TFLITE_NO_INDICES = TensorIndices([], [], [])
489TFLITE_IFM_INDICES = TensorIndices([0], [], [])
490TFLITE_IFM_WEIGHTS_INDICES = TensorIndices([0], [1], [])
491TFLITE_IFM_WEIGHTS_BIAS_INDICES = TensorIndices([0], [1], [2])
492TFLITE_IFM_IFM2_INDICES = TensorIndices([0, 1], [], [])
493TFLITE_CONV2D_BACKPROP_INDICES = TensorIndices([2], [1], [3])
494TFLITE_TRANSPOSE_CONV_INDICES = TensorIndices([0], [1], [3])
495TFLITE_CONCAT_INDICES = TensorIndices([1, 2], [], [])
496TFLITE_SPLIT_IFM_INDICES = TensorIndices([1], [], [])
497TFLITE_BLOCK_LSTM_INDICES = TensorIndices([3], [4], [])
498
Tim Hall79d07d22020-04-27 18:20:16 +0100499builtin_operator_map = {
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200500 BuiltinOperator.ADD: (
501 Op.Add,
502 OptionsSerializer("AddOptions", (fused_act, "pot_scale_int16")),
503 TFLITE_IFM_IFM2_INDICES,
504 ),
Dwight Lidman8a12da12021-07-19 13:43:05 +0200505 BuiltinOperator.AVERAGE_POOL_2D: (
506 Op.AvgPool,
507 OptionsSerializer(
508 "Pool2DOptions", ("filter_height", "filter_width", fused_act, padding, "stride_h", "stride_w")
509 ),
510 TFLITE_IFM_INDICES,
511 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200512 BuiltinOperator.CONCATENATION: (
513 Op.ConcatTFLite,
514 OptionsSerializer("ConcatenationOptions", ("axis", fused_act)),
515 TFLITE_CONCAT_INDICES,
516 ),
Dwight Lidman8a12da12021-07-19 13:43:05 +0200517 BuiltinOperator.CONV_2D: (
518 Op.Conv2DBias,
519 OptionsSerializer(
520 "Conv2DOptions", ("dilation_h_factor", "dilation_w_factor", fused_act, padding, "stride_h", "stride_w")
521 ),
522 TFLITE_IFM_WEIGHTS_BIAS_INDICES,
523 ),
524 BuiltinOperator.DEPTHWISE_CONV_2D: (
525 Op.DepthwiseConv2DBias,
526 OptionsSerializer(
527 "DepthwiseConv2DOptions",
528 ("depth_multiplier", "dilation_h_factor", "dilation_w_factor", fused_act, padding, "stride_h", "stride_w"),
529 ),
530 TFLITE_IFM_WEIGHTS_BIAS_INDICES,
531 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200532 BuiltinOperator.DEPTH_TO_SPACE: (
533 Op.DepthToSpace,
534 OptionsSerializer("DepthToSpaceOptions", ("block_size",)),
535 TFLITE_NO_INDICES,
536 ),
537 BuiltinOperator.DEQUANTIZE: (Op.Dequantize, OptionsSerializer("DequantizeOptions"), TFLITE_IFM_INDICES),
538 BuiltinOperator.EMBEDDING_LOOKUP: (Op.EmbeddingLookup, None, TFLITE_NO_INDICES),
539 BuiltinOperator.FLOOR: (Op.Floor, None, TFLITE_NO_INDICES),
Tim Hall79d07d22020-04-27 18:20:16 +0100540 BuiltinOperator.FULLY_CONNECTED: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200541 Op.FullyConnected,
erik.andersson@arm.com0cbb1662021-02-22 15:47:07 +0100542 OptionsSerializer(
Dwight Lidman8a12da12021-07-19 13:43:05 +0200543 "FullyConnectedOptions", ("asymmetric_quantize_inputs", fused_act, "keep_num_dims", "weights_format")
erik.andersson@arm.com0cbb1662021-02-22 15:47:07 +0100544 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200545 TFLITE_IFM_WEIGHTS_BIAS_INDICES,
Tim Hall79d07d22020-04-27 18:20:16 +0100546 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200547 BuiltinOperator.HASHTABLE_LOOKUP: (Op.HashtableLookup, None, TFLITE_NO_INDICES),
548 BuiltinOperator.L2_NORMALIZATION: (Op.L2Norm, OptionsSerializer("L2NormOptions", (fused_act,)), TFLITE_NO_INDICES),
Dwight Lidman8a12da12021-07-19 13:43:05 +0200549 BuiltinOperator.L2_POOL_2D: (
550 Op.L2Pool2D,
551 OptionsSerializer(
552 "Pool2DOptions", ("filter_height", "filter_width", fused_act, padding, "stride_h", "stride_w")
553 ),
554 TFLITE_NO_INDICES,
555 ),
Tim Hall79d07d22020-04-27 18:20:16 +0100556 BuiltinOperator.LOCAL_RESPONSE_NORMALIZATION: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200557 Op.LRN,
Dwight Lidman8a12da12021-07-19 13:43:05 +0200558 OptionsSerializer("LocalResponseNormalizationOptions", ("alpha", "beta", "bias", "radius")),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200559 TFLITE_NO_INDICES,
Tim Hall79d07d22020-04-27 18:20:16 +0100560 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200561 BuiltinOperator.LOGISTIC: (Op.Sigmoid, None, TFLITE_IFM_INDICES),
562 BuiltinOperator.LSH_PROJECTION: (
563 Op.LSHProjection,
564 OptionsSerializer("LSHProjectionOptions", ("type",)),
565 TFLITE_NO_INDICES,
566 ),
Dwight Lidman8a12da12021-07-19 13:43:05 +0200567 BuiltinOperator.LSTM: (
568 Op.Lstm,
569 OptionsSerializer(
570 "LSTMOptions", ("asymmetric_quantize_inputs", "cell_clip", fused_act, "kernel_type", "proj_clip")
571 ),
572 TFLITE_IFM_WEIGHTS_INDICES,
573 ),
574 BuiltinOperator.MAX_POOL_2D: (
575 Op.MaxPool,
576 OptionsSerializer(
577 "Pool2DOptions", ("filter_height", "filter_width", fused_act, padding, "stride_h", "stride_w")
578 ),
579 TFLITE_IFM_INDICES,
580 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200581 BuiltinOperator.MUL: (Op.Mul, OptionsSerializer("MulOptions", (fused_act,)), TFLITE_IFM_IFM2_INDICES),
582 BuiltinOperator.RELU: (Op.Relu, None, TFLITE_IFM_INDICES),
583 BuiltinOperator.RELU_N1_TO_1: (Op.ReluN1To1, None, TFLITE_IFM_INDICES),
584 BuiltinOperator.RELU6: (Op.Relu6, None, TFLITE_IFM_INDICES),
585 BuiltinOperator.RESHAPE: (
586 Op.Reshape,
Dwight Lidmancfb42622021-08-11 12:28:51 +0200587 OptionsSerializer("ReshapeOptions", (("new_shape", is_int_vec),)),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200588 TFLITE_IFM_INDICES,
589 ),
Tim Hall79d07d22020-04-27 18:20:16 +0100590 BuiltinOperator.RESIZE_BILINEAR: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200591 Op.ResizeBilinear,
Tim Hall79d07d22020-04-27 18:20:16 +0100592 OptionsSerializer("ResizeBilinearOptions", ("align_corners", "half_pixel_centers")),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200593 TFLITE_IFM_INDICES,
Tim Hall79d07d22020-04-27 18:20:16 +0100594 ),
Dwight Lidman8a12da12021-07-19 13:43:05 +0200595 BuiltinOperator.RNN: (
596 Op.Rnn,
597 OptionsSerializer("RNNOptions", ("asymmetric_quantize_inputs", fused_act)),
598 TFLITE_IFM_WEIGHTS_INDICES,
599 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200600 BuiltinOperator.SOFTMAX: (Op.Softmax, OptionsSerializer("SoftmaxOptions", ("beta",)), TFLITE_IFM_INDICES),
601 BuiltinOperator.SPACE_TO_DEPTH: (
602 Op.SpaceToDepth,
603 OptionsSerializer("SpaceToDepthOptions", ("block_size",)),
604 TFLITE_NO_INDICES,
605 ),
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200606 BuiltinOperator.SVDF: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200607 Op.Svdf,
Dwight Lidman8a12da12021-07-19 13:43:05 +0200608 OptionsSerializer("SVDFOptions", ("asymmetric_quantize_inputs", fused_act, "rank")),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200609 TFLITE_NO_INDICES,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200610 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200611 BuiltinOperator.TANH: (Op.Tanh, None, TFLITE_IFM_INDICES),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200612 BuiltinOperator.CONCAT_EMBEDDINGS: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200613 Op.ConcatEmbeddings,
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200614 OptionsSerializer(
615 "ConcatEmbeddingsOptions",
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200616 (
Dwight Lidman8a12da12021-07-19 13:43:05 +0200617 "embedding_dim_per_channel",
618 "embedding_dim_per_channel_as_numpy",
619 "embedding_dim_per_channel_is_none",
620 "embedding_dim_per_channel_length",
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200621 "num_channels",
622 "num_columns_per_channel",
623 "num_columns_per_channel_as_numpy",
Dwight Lidman8a12da12021-07-19 13:43:05 +0200624 "num_columns_per_channel_is_none",
625 "num_columns_per_channel_length",
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200626 ),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200627 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200628 TFLITE_NO_INDICES,
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200629 ),
630 BuiltinOperator.SKIP_GRAM: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200631 Op.SkipGram,
Dwight Lidman8a12da12021-07-19 13:43:05 +0200632 OptionsSerializer("SkipGramOptions", ("include_all_ngrams", "max_skip_size", "ngram_size")),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200633 TFLITE_NO_INDICES,
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200634 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200635 BuiltinOperator.CALL: (Op.Call, OptionsSerializer("CallOptions", ("subgraph",)), TFLITE_NO_INDICES),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200636 BuiltinOperator.EMBEDDING_LOOKUP_SPARSE: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200637 Op.EmbeddingLookupSparse,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200638 OptionsSerializer("EmbeddingLookupSparseOptions", ("combiner",)),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200639 TFLITE_NO_INDICES,
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200640 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200641 BuiltinOperator.PAD: (Op.Pad, OptionsSerializer("PadOptions"), TFLITE_IFM_INDICES),
642 BuiltinOperator.UNIDIRECTIONAL_SEQUENCE_RNN: (
643 Op.UnidirectionalSequenceRnn,
Dwight Lidman8a12da12021-07-19 13:43:05 +0200644 OptionsSerializer("SequenceRNNOptions", ("asymmetric_quantize_inputs", fused_act, "time_major")),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200645 TFLITE_IFM_WEIGHTS_INDICES,
646 ),
Dwight Lidman8a12da12021-07-19 13:43:05 +0200647 BuiltinOperator.GATHER: (
648 Op.GatherV2,
649 OptionsSerializer("GatherOptions", ("axis", "batch_dims")),
650 TFLITE_NO_INDICES,
651 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200652 BuiltinOperator.BATCH_TO_SPACE_ND: (
653 Op.BatchToSpaceND,
654 OptionsSerializer("BatchToSpaceNDOptions"),
655 TFLITE_NO_INDICES,
656 ),
657 BuiltinOperator.SPACE_TO_BATCH_ND: (
658 Op.SpaceToBatchND,
659 OptionsSerializer("SpaceToBatchNDOptions"),
660 TFLITE_NO_INDICES,
661 ),
James Ward6bf16132021-09-08 11:14:20 +0100662 BuiltinOperator.TRANSPOSE: (Op.Transpose, OptionsSerializer("TransposeOptions"), TFLITE_IFM_IFM2_INDICES),
Dwight Lidman8a12da12021-07-19 13:43:05 +0200663 BuiltinOperator.MEAN: (Op.Mean, OptionsSerializer("ReducerOptions", ("keep_dims",)), TFLITE_IFM_INDICES),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200664 BuiltinOperator.SUB: (
665 Op.Sub,
Dwight Lidman8a12da12021-07-19 13:43:05 +0200666 OptionsSerializer("SubOptions", (fused_act, "pot_scale_int16")),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200667 TFLITE_IFM_IFM2_INDICES,
668 ),
669 BuiltinOperator.DIV: (Op.Div, OptionsSerializer("DivOptions", (fused_act,)), TFLITE_NO_INDICES),
670 BuiltinOperator.SQUEEZE: (
671 Op.Squeeze,
Jonas Ohlssond8575072022-03-30 10:30:25 +0200672 OptionsSerializer(
673 "SqueezeOptions",
674 (("squeeze_dims", is_int_vec),),
675 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200676 TFLITE_IFM_INDICES,
677 ),
678 BuiltinOperator.UNIDIRECTIONAL_SEQUENCE_LSTM: (
679 Op.UnidirectionalSequenceLstm,
Dwight Lidman8a12da12021-07-19 13:43:05 +0200680 OptionsSerializer(
681 "UnidirectionalSequenceLSTMOptions",
682 ("asymmetric_quantize_inputs", "cell_clip", fused_act, "proj_clip", "time_major"),
683 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200684 TFLITE_IFM_WEIGHTS_INDICES,
685 ),
Tim Hall79d07d22020-04-27 18:20:16 +0100686 BuiltinOperator.STRIDED_SLICE: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200687 Op.StridedSlice,
Tim Hall79d07d22020-04-27 18:20:16 +0100688 OptionsSerializer(
Dwight Lidman8a12da12021-07-19 13:43:05 +0200689 "StridedSliceOptions", ("begin_mask", "ellipsis_mask", "end_mask", "new_axis_mask", "shrink_axis_mask")
Tim Hall79d07d22020-04-27 18:20:16 +0100690 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200691 TFLITE_IFM_INDICES,
Tim Hall79d07d22020-04-27 18:20:16 +0100692 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200693 BuiltinOperator.BIDIRECTIONAL_SEQUENCE_RNN: (
694 Op.BidirectionalSequenceRnn,
Dwight Lidman8a12da12021-07-19 13:43:05 +0200695 OptionsSerializer(
696 "BidirectionalSequenceRNNOptions", ("asymmetric_quantize_inputs", fused_act, "merge_outputs", "time_major")
697 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200698 TFLITE_IFM_WEIGHTS_INDICES,
699 ),
James Ward6bf16132021-09-08 11:14:20 +0100700 BuiltinOperator.EXP: (Op.Exp, OptionsSerializer("ExpOptions"), TFLITE_NO_INDICES),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200701 BuiltinOperator.TOPK_V2: (Op.TopKV2, OptionsSerializer("TopKV2Options"), TFLITE_NO_INDICES),
702 BuiltinOperator.SPLIT: (Op.Split, OptionsSerializer("SplitOptions", ("num_splits",)), TFLITE_SPLIT_IFM_INDICES),
703 BuiltinOperator.LOG_SOFTMAX: (Op.LogSoftmax, OptionsSerializer("LogSoftmaxOptions"), TFLITE_NO_INDICES),
704 BuiltinOperator.DELEGATE: (Op.Delegate, None, TFLITE_NO_INDICES),
705 BuiltinOperator.BIDIRECTIONAL_SEQUENCE_LSTM: (
706 Op.BidirectionalSequenceLstm,
Dwight Lidman8a12da12021-07-19 13:43:05 +0200707 OptionsSerializer(
708 "BidirectionalSequenceLSTMOptions",
709 ("asymmetric_quantize_inputs", "cell_clip", fused_act, "merge_outputs", "proj_clip", "time_major"),
710 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200711 TFLITE_IFM_WEIGHTS_INDICES,
712 ),
Tim Hall79d07d22020-04-27 18:20:16 +0100713 BuiltinOperator.CAST: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200714 Op.Cast,
Tim Hall79d07d22020-04-27 18:20:16 +0100715 OptionsSerializer(
716 "CastOptions",
717 (
718 ("in_data_type", datatype_deserialize, datatype_serialize),
719 ("out_data_type", datatype_deserialize, datatype_serialize),
720 ),
721 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200722 TFLITE_NO_INDICES,
Tim Hall79d07d22020-04-27 18:20:16 +0100723 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200724 BuiltinOperator.PRELU: (Op.Prelu, None, TFLITE_NO_INDICES),
725 BuiltinOperator.MAXIMUM: (Op.Maximum, OptionsSerializer("MaximumMinimumOptions"), TFLITE_IFM_IFM2_INDICES),
Tim Hall79d07d22020-04-27 18:20:16 +0100726 BuiltinOperator.ARG_MAX: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200727 Op.ArgMax,
Tim Hall79d07d22020-04-27 18:20:16 +0100728 OptionsSerializer("ArgMaxOptions", (("output_type", datatype_deserialize, datatype_serialize),)),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200729 TFLITE_NO_INDICES,
Tim Hall79d07d22020-04-27 18:20:16 +0100730 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200731 BuiltinOperator.MINIMUM: (Op.Minimum, OptionsSerializer("MaximumMinimumOptions"), TFLITE_IFM_IFM2_INDICES),
732 BuiltinOperator.LESS: (Op.Less, OptionsSerializer("LessOptions"), TFLITE_NO_INDICES),
733 BuiltinOperator.NEG: (Op.Neg, OptionsSerializer("NegOptions"), TFLITE_NO_INDICES),
734 BuiltinOperator.PADV2: (Op.PadV2, OptionsSerializer("PadV2Options"), TFLITE_NO_INDICES),
735 BuiltinOperator.GREATER: (Op.Greater, OptionsSerializer("GreaterOptions"), TFLITE_NO_INDICES),
736 BuiltinOperator.GREATER_EQUAL: (Op.GreaterEqual, OptionsSerializer("GreaterEqualOptions"), TFLITE_NO_INDICES),
737 BuiltinOperator.LESS_EQUAL: (Op.LessEqual, OptionsSerializer("LessEqualOptions"), TFLITE_NO_INDICES),
738 BuiltinOperator.SELECT: (Op.Select, OptionsSerializer("SelectOptions"), TFLITE_NO_INDICES),
739 BuiltinOperator.SLICE: (Op.Slice, OptionsSerializer("SliceOptions"), TFLITE_IFM_INDICES),
740 BuiltinOperator.SIN: (Op.Sin, None, TFLITE_NO_INDICES),
Tim Hall79d07d22020-04-27 18:20:16 +0100741 BuiltinOperator.TRANSPOSE_CONV: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200742 Op.Conv2DBackpropInput,
Dwight Lidman8a12da12021-07-19 13:43:05 +0200743 OptionsSerializer("TransposeConvOptions", (padding, "stride_h", "stride_w")),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200744 TFLITE_CONV2D_BACKPROP_INDICES,
Tim Hall79d07d22020-04-27 18:20:16 +0100745 ),
746 BuiltinOperator.SPARSE_TO_DENSE: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200747 Op.SparseToDense,
Tim Hall79d07d22020-04-27 18:20:16 +0100748 OptionsSerializer("SparseToDenseOptions", ("validate_indices",)),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200749 TFLITE_NO_INDICES,
Tim Hall79d07d22020-04-27 18:20:16 +0100750 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200751 BuiltinOperator.TILE: (Op.Tile, OptionsSerializer("TileOptions"), TFLITE_NO_INDICES),
752 BuiltinOperator.EXPAND_DIMS: (Op.ExpandDims, OptionsSerializer("ExpandDimsOptions"), TFLITE_IFM_INDICES),
753 BuiltinOperator.EQUAL: (Op.Equal, OptionsSerializer("EqualOptions"), TFLITE_NO_INDICES),
754 BuiltinOperator.NOT_EQUAL: (Op.NotEqual, OptionsSerializer("NotEqualOptions"), TFLITE_NO_INDICES),
755 BuiltinOperator.LOG: (Op.Log, None, TFLITE_NO_INDICES),
Dwight Lidman8a12da12021-07-19 13:43:05 +0200756 BuiltinOperator.SUM: (Op.Sum, OptionsSerializer("ReducerOptions", ("keep_dims",)), TFLITE_NO_INDICES),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200757 BuiltinOperator.SQRT: (Op.Sqrt, None, TFLITE_NO_INDICES),
758 BuiltinOperator.RSQRT: (Op.Rsqrt, None, TFLITE_NO_INDICES),
Tim Hall79d07d22020-04-27 18:20:16 +0100759 BuiltinOperator.SHAPE: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200760 Op.Shape,
Tim Hall79d07d22020-04-27 18:20:16 +0100761 OptionsSerializer("ShapeOptions", (("out_type", datatype_deserialize, datatype_serialize),)),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200762 TFLITE_NO_INDICES,
Tim Hall79d07d22020-04-27 18:20:16 +0100763 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200764 BuiltinOperator.POW: (Op.Pow, OptionsSerializer("PowOptions"), TFLITE_NO_INDICES),
Tim Hall79d07d22020-04-27 18:20:16 +0100765 BuiltinOperator.ARG_MIN: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200766 Op.ArgMin,
Tim Hall79d07d22020-04-27 18:20:16 +0100767 OptionsSerializer("ArgMinOptions", (("output_type", datatype_deserialize, datatype_serialize),)),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200768 TFLITE_NO_INDICES,
Tim Hall79d07d22020-04-27 18:20:16 +0100769 ),
770 BuiltinOperator.FAKE_QUANT: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200771 Op.FakeQuantWithMinMaxArgs,
Dwight Lidman8a12da12021-07-19 13:43:05 +0200772 OptionsSerializer("FakeQuantOptions", ("max", "min", "narrow_range", "num_bits")),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200773 TFLITE_NO_INDICES,
Tim Hall79d07d22020-04-27 18:20:16 +0100774 ),
Dwight Lidman8a12da12021-07-19 13:43:05 +0200775 BuiltinOperator.REDUCE_PROD: (Op.Prod, OptionsSerializer("ReducerOptions", ("keep_dims",)), TFLITE_NO_INDICES),
776 BuiltinOperator.REDUCE_MAX: (Op.Max, OptionsSerializer("ReducerOptions", ("keep_dims",)), TFLITE_NO_INDICES),
777 BuiltinOperator.PACK: (Op.Pack, OptionsSerializer("PackOptions", ("axis", "values_count")), TFLITE_IFM_INDICES),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200778 BuiltinOperator.LOGICAL_OR: (Op.LogicalOr, OptionsSerializer("LogicalOrOptions"), TFLITE_NO_INDICES),
779 BuiltinOperator.ONE_HOT: (Op.OneHot, OptionsSerializer("OneHotOptions", ("axis",)), TFLITE_NO_INDICES),
780 BuiltinOperator.LOGICAL_AND: (Op.LogicalAnd, OptionsSerializer("LogicalAndOptions"), TFLITE_NO_INDICES),
781 BuiltinOperator.LOGICAL_NOT: (Op.LogicalNot, OptionsSerializer("LogicalNotOptions"), TFLITE_NO_INDICES),
Dwight Lidman8a12da12021-07-19 13:43:05 +0200782 BuiltinOperator.UNPACK: (Op.Unpack, OptionsSerializer("UnpackOptions", ("axis", "num")), TFLITE_IFM_INDICES),
783 BuiltinOperator.REDUCE_MIN: (Op.Min, OptionsSerializer("ReducerOptions", ("keep_dims",)), TFLITE_NO_INDICES),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200784 BuiltinOperator.FLOOR_DIV: (Op.FloorDiv, OptionsSerializer("FloorDivOptions"), TFLITE_NO_INDICES),
Dwight Lidman8a12da12021-07-19 13:43:05 +0200785 BuiltinOperator.REDUCE_ANY: (Op.Any, OptionsSerializer("ReducerOptions", ("keep_dims",)), TFLITE_NO_INDICES),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200786 BuiltinOperator.SQUARE: (Op.Square, OptionsSerializer("SquareOptions"), TFLITE_NO_INDICES),
787 BuiltinOperator.ZEROS_LIKE: (Op.ZerosLike, OptionsSerializer("ZerosLikeOptions"), TFLITE_NO_INDICES),
788 BuiltinOperator.FILL: (Op.Fill, OptionsSerializer("FillOptions"), TFLITE_NO_INDICES),
789 BuiltinOperator.FLOOR_MOD: (Op.FloorMod, OptionsSerializer("FloorModOptions"), TFLITE_NO_INDICES),
790 BuiltinOperator.RANGE: (Op.Range, OptionsSerializer("RangeOptions"), TFLITE_NO_INDICES),
Tim Hall79d07d22020-04-27 18:20:16 +0100791 BuiltinOperator.RESIZE_NEAREST_NEIGHBOR: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200792 Op.ResizeNearestNeighbor,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200793 OptionsSerializer("ResizeNearestNeighborOptions", ("align_corners", "half_pixel_centers")),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200794 TFLITE_NO_INDICES,
Tim Hall79d07d22020-04-27 18:20:16 +0100795 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200796 BuiltinOperator.LEAKY_RELU: (Op.LeakyRelu, OptionsSerializer("LeakyReluOptions", ("alpha",)), TFLITE_IFM_INDICES),
797 BuiltinOperator.SQUARED_DIFFERENCE: (
798 Op.SquaredDifference,
799 OptionsSerializer("SquaredDifferenceOptions"),
800 TFLITE_NO_INDICES,
801 ),
802 BuiltinOperator.MIRROR_PAD: (Op.MirrorPad, OptionsSerializer("MirrorPadOptions", ("mode",)), TFLITE_NO_INDICES),
803 BuiltinOperator.ABS: (Op.Abs, OptionsSerializer("AbsOptions"), TFLITE_IFM_INDICES),
804 BuiltinOperator.SPLIT_V: (Op.SplitV, OptionsSerializer("SplitVOptions", ("num_splits",)), TFLITE_IFM_INDICES),
Tim Hall79d07d22020-04-27 18:20:16 +0100805 BuiltinOperator.UNIQUE: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200806 Op.Unique,
Tim Hall79d07d22020-04-27 18:20:16 +0100807 OptionsSerializer("UniqueOptions", (("idx_out_type", datatype_deserialize, datatype_serialize),)),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200808 TFLITE_NO_INDICES,
Tim Hall79d07d22020-04-27 18:20:16 +0100809 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200810 BuiltinOperator.CEIL: (Op.Ceil, None, TFLITE_NO_INDICES),
811 BuiltinOperator.REVERSE_V2: (Op.ReverseV2, OptionsSerializer("ReverseV2Options"), TFLITE_NO_INDICES),
812 BuiltinOperator.ADD_N: (Op.AddN, OptionsSerializer("AddNOptions"), TFLITE_NO_INDICES),
813 BuiltinOperator.GATHER_ND: (Op.GatherNd, OptionsSerializer("GatherNdOptions"), TFLITE_NO_INDICES),
814 BuiltinOperator.COS: (Op.Cos, OptionsSerializer("CosOptions"), TFLITE_NO_INDICES),
815 BuiltinOperator.WHERE: (Op.Where, OptionsSerializer("WhereOptions"), TFLITE_NO_INDICES),
816 BuiltinOperator.RANK: (Op.Rank, OptionsSerializer("RankOptions"), TFLITE_NO_INDICES),
817 BuiltinOperator.ELU: (Op.Elu, None, TFLITE_NO_INDICES),
Tim Hall79d07d22020-04-27 18:20:16 +0100818 BuiltinOperator.REVERSE_SEQUENCE: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200819 Op.ReverseSequence,
Dwight Lidman8a12da12021-07-19 13:43:05 +0200820 OptionsSerializer("ReverseSequenceOptions", ("batch_dim", "seq_dim")),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200821 TFLITE_NO_INDICES,
Tim Hall79d07d22020-04-27 18:20:16 +0100822 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200823 BuiltinOperator.MATRIX_DIAG: (Op.MatrixDiag, OptionsSerializer("MatrixDiagOptions"), TFLITE_NO_INDICES),
824 BuiltinOperator.QUANTIZE: (Op.Quantize, OptionsSerializer("QuantizeOptions"), TFLITE_IFM_INDICES),
825 BuiltinOperator.MATRIX_SET_DIAG: (Op.MatrixSetDiag, OptionsSerializer("MatrixSetDiagOptions"), TFLITE_NO_INDICES),
826 BuiltinOperator.ROUND: (Op.Round, None, TFLITE_NO_INDICES),
827 BuiltinOperator.HARD_SWISH: (Op.HardSwish, OptionsSerializer("HardSwishOptions"), TFLITE_IFM_INDICES),
828 BuiltinOperator.IF: (
829 Op.If,
Dwight Lidman8a12da12021-07-19 13:43:05 +0200830 OptionsSerializer("IfOptions", ("else_subgraph_index", "then_subgraph_index")),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200831 TFLITE_NO_INDICES,
832 ),
Louis Verhaardaee5d752020-09-30 09:01:52 +0200833 BuiltinOperator.WHILE: (
834 Op.While,
Dwight Lidman8a12da12021-07-19 13:43:05 +0200835 OptionsSerializer("WhileOptions", ("body_subgraph_index", "cond_subgraph_index")),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200836 TFLITE_NO_INDICES,
Louis Verhaardaee5d752020-09-30 09:01:52 +0200837 ),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200838 BuiltinOperator.NON_MAX_SUPPRESSION_V4: (
839 Op.NonMaxSuppressionV4,
840 OptionsSerializer("NonMaxSuppressionV4Options"),
841 TFLITE_NO_INDICES,
842 ),
843 BuiltinOperator.NON_MAX_SUPPRESSION_V5: (
844 Op.NonMaxSuppressionV5,
845 OptionsSerializer("NonMaxSuppressionV5Options"),
846 TFLITE_NO_INDICES,
847 ),
848 BuiltinOperator.SCATTER_ND: (Op.ScatterNd, OptionsSerializer("ScatterNdOptions"), TFLITE_NO_INDICES),
849 BuiltinOperator.SELECT_V2: (Op.SelectV2, OptionsSerializer("SelectV2Options"), TFLITE_NO_INDICES),
850 BuiltinOperator.DENSIFY: (Op.Densify, OptionsSerializer("DensifyOptions"), TFLITE_NO_INDICES),
851 BuiltinOperator.SEGMENT_SUM: (Op.SegmentSum, OptionsSerializer("SegmentSumOptions"), TFLITE_NO_INDICES),
852 BuiltinOperator.BATCH_MATMUL: (
853 Op.BatchMatMul,
Dwight Lidman8a12da12021-07-19 13:43:05 +0200854 OptionsSerializer("BatchMatMulOptions", ("adj_x", "adj_y", "asymmetric_quantize_inputs")),
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200855 TFLITE_NO_INDICES,
856 ),
857 BuiltinOperator.CUMSUM: (
858 Op.Cumsum,
859 OptionsSerializer("CumsumOptions", ("exclusive", "reverse")),
860 TFLITE_NO_INDICES,
861 ),
Dwight Lidman8a12da12021-07-19 13:43:05 +0200862 BuiltinOperator.CALL_ONCE: (
863 Op.CallOnce,
864 OptionsSerializer("CallOnceOptions", ("init_subgraph_index",)),
865 TFLITE_NO_INDICES,
866 ),
867 BuiltinOperator.BROADCAST_TO: (Op.BroadcastTo, OptionsSerializer("BroadcastToOptions"), TFLITE_NO_INDICES),
868 BuiltinOperator.RFFT2D: (Op.Rfft2D, OptionsSerializer("Rfft2dOptions"), TFLITE_NO_INDICES),
869 BuiltinOperator.CONV_3D: (
870 Op.Conv3D,
871 OptionsSerializer(
872 "Conv3DOptions",
873 (
874 "dilation_d_factor",
875 "dilation_h_factor",
876 "dilation_w_factor",
877 fused_act,
878 padding,
879 "stride_d",
880 "stride_h",
881 "stride_w",
882 ),
883 ),
884 TFLITE_NO_INDICES,
885 ),
886 BuiltinOperator.IMAG: (Op.Imag, None, TFLITE_NO_INDICES),
887 BuiltinOperator.REAL: (Op.Real, None, TFLITE_NO_INDICES),
888 BuiltinOperator.COMPLEX_ABS: (Op.ComplexAbs, None, TFLITE_NO_INDICES),
889 BuiltinOperator.HASHTABLE: (
890 Op.Hashtable,
891 OptionsSerializer("HashtableOptions", ("key_dtype", "table_id", "value_dtype")),
892 TFLITE_NO_INDICES,
893 ),
894 BuiltinOperator.HASHTABLE_FIND: (Op.HashtableFind, OptionsSerializer("HashtableFindOptions"), TFLITE_NO_INDICES),
895 BuiltinOperator.HASHTABLE_IMPORT: (
896 Op.HashtableImport,
897 OptionsSerializer("HashtableImportOptions"),
898 TFLITE_NO_INDICES,
899 ),
900 BuiltinOperator.HASHTABLE_SIZE: (Op.HashtableSize, OptionsSerializer("HashtableSizeOptions"), TFLITE_NO_INDICES),
901 BuiltinOperator.REDUCE_ALL: (Op.ReduceAll, OptionsSerializer("ReducerOptions", ("keep_dims",)), TFLITE_NO_INDICES),
902 BuiltinOperator.CONV_3D_TRANSPOSE: (
903 Op.Conv3DTranspose,
904 OptionsSerializer(
905 "Conv3DOptions",
906 (
907 "dilation_d_factor",
908 "dilation_h_factor",
909 "dilation_w_factor",
910 fused_act,
911 padding,
912 "stride_d",
913 "stride_h",
914 "stride_w",
915 ),
916 ),
917 TFLITE_NO_INDICES,
918 ),
Rickard Bolin2de898a2021-12-20 08:35:23 +0000919 BuiltinOperator.VAR_HANDLE: (
920 Op.VarHandle,
Jonas Ohlssond8575072022-03-30 10:30:25 +0200921 OptionsSerializer(
922 "VarHandleOptions",
923 (
924 "container",
925 "shared_name",
926 ),
927 ),
Rickard Bolin2de898a2021-12-20 08:35:23 +0000928 TFLITE_NO_INDICES,
929 ),
930 BuiltinOperator.READ_VARIABLE: (Op.ReadVariable, OptionsSerializer("ReadVariableOptions"), TFLITE_NO_INDICES),
931 BuiltinOperator.ASSIGN_VARIABLE: (Op.AssignVariable, OptionsSerializer("AssignVariableOptions"), TFLITE_NO_INDICES),
932 BuiltinOperator.BROADCAST_ARGS: (Op.BroadcastArgs, None, TFLITE_NO_INDICES),
933 BuiltinOperator.RANDOM_STANDARD_NORMAL: (
934 Op.RandomStandardNormal,
Jonas Ohlssond8575072022-03-30 10:30:25 +0200935 OptionsSerializer(
936 "RandomOptions",
937 (
938 "seed",
939 "seed2",
940 ),
941 ),
Rickard Bolin2de898a2021-12-20 08:35:23 +0000942 TFLITE_NO_INDICES,
943 ),
Tim Hallb94bcd42021-09-07 14:24:51 +0100944 BuiltinOperator.CUSTOM: (Op.Custom, CustomOptionsSerializer(), TFLITE_NO_INDICES),
Tim Hall79d07d22020-04-27 18:20:16 +0100945}
946
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200947builtin_operator_inv_map = {v[0]: (k, v[1], v[2]) for k, v in builtin_operator_map.items()}
Tim Hall79d07d22020-04-27 18:20:16 +0100948
Patrik Gustavsson5e26eda2021-06-30 09:07:16 +0200949builtin_operator_inv_map[Op.CustomNpuOp] = (BuiltinOperator.CUSTOM, CustomOptionsSerializer(), TFLITE_NO_INDICES)
Michael McGeagh219ec072020-11-09 11:11:26 +0000950
Michael McGeagh837dc1b2020-11-10 12:38:25 +0000951BUILTIN_OPERATOR_UNKNOWN = "UNKNOWN"
952
Tim Halla3fe6652022-03-03 17:43:16 +0000953builtin_operator_name_map = {v: k for k, v in vars(BuiltinOperator).items()}
Michael McGeagh837dc1b2020-11-10 12:38:25 +0000954
Michael McGeagh219ec072020-11-09 11:11:26 +0000955
Tim Halla3fe6652022-03-03 17:43:16 +0000956def optype_to_builtintype(op_type: Op):
Michael McGeagh219ec072020-11-09 11:11:26 +0000957 if op_type in builtin_operator_inv_map:
Tim Halla3fe6652022-03-03 17:43:16 +0000958 return builtin_operator_name_map[builtin_operator_inv_map[op_type][0]]
Michael McGeagh219ec072020-11-09 11:11:26 +0000959 else:
Michael McGeagh837dc1b2020-11-10 12:38:25 +0000960 return BUILTIN_OPERATOR_UNKNOWN