blob: ea9e8a3605f3a5ef48c4e14493bb93a74a26f400 [file] [log] [blame]
Tim Hall79d07d22020-04-27 18:20:16 +01001# Copyright (C) 2020 Arm Limited or its affiliates. All rights reserved.
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.
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
Diego Russoe8a10452020-04-21 17:39:10 +010027from .tflite import AbsOptions
28from .tflite import AddNOptions
Tim Hall79d07d22020-04-27 18:20:16 +010029from .tflite import AddOptions
Tim Hall79d07d22020-04-27 18:20:16 +010030from .tflite import ArgMaxOptions
Tim Hall79d07d22020-04-27 18:20:16 +010031from .tflite import ArgMinOptions
Jacob Bohlin8daf6b72020-09-15 16:28:35 +020032from .tflite import BatchMatMulOptions
Diego Russoe8a10452020-04-21 17:39:10 +010033from .tflite import BatchToSpaceNDOptions
Tim Hall79d07d22020-04-27 18:20:16 +010034from .tflite import BidirectionalSequenceLSTMOptions
35from .tflite import BidirectionalSequenceRNNOptions
Diego Russoe8a10452020-04-21 17:39:10 +010036from .tflite import CallOptions
37from .tflite import CastOptions
38from .tflite import ConcatEmbeddingsOptions
39from .tflite import ConcatenationOptions
40from .tflite import Conv2DOptions
Tim Hall79d07d22020-04-27 18:20:16 +010041from .tflite import CosOptions
Tim Hall79d07d22020-04-27 18:20:16 +010042from .tflite import DensifyOptions
43from .tflite import DepthToSpaceOptions
Diego Russoe8a10452020-04-21 17:39:10 +010044from .tflite import DepthwiseConv2DOptions
45from .tflite import DequantizeOptions
46from .tflite import DivOptions
47from .tflite import EmbeddingLookupSparseOptions
48from .tflite import EqualOptions
49from .tflite import ExpandDimsOptions
50from .tflite import ExpOptions
51from .tflite import FakeQuantOptions
52from .tflite import FillOptions
53from .tflite import FloorDivOptions
54from .tflite import FloorModOptions
55from .tflite import FullyConnectedOptions
56from .tflite import GatherNdOptions
57from .tflite import GatherOptions
58from .tflite import GreaterEqualOptions
59from .tflite import GreaterOptions
Louis Verhaard678645b2020-06-15 15:22:47 +020060from .tflite import HardSwishOptions
Tim Hall79d07d22020-04-27 18:20:16 +010061from .tflite import IfOptions
Diego Russoe8a10452020-04-21 17:39:10 +010062from .tflite import L2NormOptions
63from .tflite import LeakyReluOptions
64from .tflite import LessEqualOptions
65from .tflite import LessOptions
66from .tflite import LocalResponseNormalizationOptions
67from .tflite import LogicalAndOptions
68from .tflite import LogicalNotOptions
69from .tflite import LogicalOrOptions
70from .tflite import LogSoftmaxOptions
71from .tflite import LSHProjectionOptions
72from .tflite import LSTMOptions
73from .tflite import MatrixDiagOptions
74from .tflite import MatrixSetDiagOptions
75from .tflite import MaximumMinimumOptions
76from .tflite import MirrorPadOptions
77from .tflite import MulOptions
78from .tflite import NegOptions
Tim Hall79d07d22020-04-27 18:20:16 +010079from .tflite import NonMaxSuppressionV4Options
80from .tflite import NonMaxSuppressionV5Options
Diego Russoe8a10452020-04-21 17:39:10 +010081from .tflite import NotEqualOptions
82from .tflite import OneHotOptions
83from .tflite import PackOptions
84from .tflite import PadOptions
85from .tflite import PadV2Options
86from .tflite import Pool2DOptions
87from .tflite import PowOptions
88from .tflite import QuantizeOptions
89from .tflite import RangeOptions
90from .tflite import RankOptions
91from .tflite import ReducerOptions
92from .tflite import ReshapeOptions
93from .tflite import ResizeBilinearOptions
94from .tflite import ResizeNearestNeighborOptions
95from .tflite import ReverseSequenceOptions
96from .tflite import ReverseV2Options
97from .tflite import RNNOptions
Tim Hall79d07d22020-04-27 18:20:16 +010098from .tflite import ScatterNdOptions
99from .tflite import SegmentSumOptions
Diego Russoe8a10452020-04-21 17:39:10 +0100100from .tflite import SelectOptions
Tim Hall79d07d22020-04-27 18:20:16 +0100101from .tflite import SelectV2Options
Diego Russoe8a10452020-04-21 17:39:10 +0100102from .tflite import SequenceRNNOptions
103from .tflite import ShapeOptions
104from .tflite import SkipGramOptions
105from .tflite import SliceOptions
106from .tflite import SoftmaxOptions
107from .tflite import SpaceToBatchNDOptions
108from .tflite import SpaceToDepthOptions
109from .tflite import SparseToDenseOptions
110from .tflite import SplitOptions
111from .tflite import SplitVOptions
112from .tflite import SquaredDifferenceOptions
113from .tflite import SquareOptions
114from .tflite import SqueezeOptions
115from .tflite import StridedSliceOptions
116from .tflite import SubOptions
117from .tflite import SVDFOptions
118from .tflite import TileOptions
119from .tflite import TopKV2Options
120from .tflite import TransposeConvOptions
121from .tflite import TransposeOptions
122from .tflite import UnidirectionalSequenceLSTMOptions
123from .tflite import UniqueOptions
124from .tflite import UnpackOptions
125from .tflite import WhereOptions
Tim Hall79d07d22020-04-27 18:20:16 +0100126from .tflite import WhileOptions
Diego Russoe8a10452020-04-21 17:39:10 +0100127from .tflite import ZerosLikeOptions
128from .tflite.ActivationFunctionType import ActivationFunctionType
Diego Russoea6111a2020-04-14 18:41:58 +0100129from .tflite.BuiltinOperator import BuiltinOperator
130from .tflite.BuiltinOptions import BuiltinOptions
131from .tflite.Padding import Padding
Diego Russoe8a10452020-04-21 17:39:10 +0100132from .tflite.TensorType import TensorType
Tim Hall79d07d22020-04-27 18:20:16 +0100133
134
135def inverse_map(map):
136 return {v: k for k, v in map.items()}
137
138
139datatype_map = {
140 TensorType.UINT8: DataType.uint8,
141 TensorType.INT8: DataType.int8,
142 TensorType.INT16: DataType.int16,
143 TensorType.INT32: DataType.int32,
144 TensorType.INT64: DataType.int64,
145 TensorType.FLOAT16: DataType.float16,
146 TensorType.FLOAT32: DataType.float32,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200147 TensorType.FLOAT64: DataType.float64,
Tim Hall79d07d22020-04-27 18:20:16 +0100148 TensorType.STRING: DataType.string,
149 TensorType.BOOL: DataType.bool,
Jacob Bohlinf767b932020-08-13 15:32:45 +0200150 TensorType.COMPLEX64: DataType.complex64,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200151 TensorType.COMPLEX128: DataType.complex128,
Tim Hall79d07d22020-04-27 18:20:16 +0100152}
153
154datatype_inv_map = inverse_map(datatype_map)
155datatype_inv_map[DataType.quint8] = TensorType.UINT8
156
157datatype_inv_map[DataType.qint8] = TensorType.INT8
158datatype_inv_map[DataType.qint16] = TensorType.INT16
159datatype_inv_map[DataType.qint32] = TensorType.INT32
160
161
162datatype_map_numpy = {
163 TensorType.UINT8: np.uint8,
164 TensorType.INT8: np.int8,
165 TensorType.INT16: np.int16,
166 TensorType.INT32: np.int32,
167 TensorType.INT64: np.int64,
168 TensorType.FLOAT16: np.float16,
169 TensorType.FLOAT32: np.float32,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200170 TensorType.FLOAT64: np.float64,
Tim Hall79d07d22020-04-27 18:20:16 +0100171 TensorType.BOOL: np.bool,
Jacob Bohlinf767b932020-08-13 15:32:45 +0200172 TensorType.COMPLEX64: np.complex64,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200173 TensorType.COMPLEX128: np.complex128,
Dwight Lidmane05de452020-11-05 15:56:08 +0100174 TensorType.STRING: np.dtype("S1"),
Tim Hall79d07d22020-04-27 18:20:16 +0100175}
176
177
178builtin_options_map = {
179 BuiltinOptions.Conv2DOptions: Conv2DOptions.Conv2DOptions,
180 BuiltinOptions.DepthwiseConv2DOptions: DepthwiseConv2DOptions.DepthwiseConv2DOptions,
181 BuiltinOptions.ConcatEmbeddingsOptions: ConcatEmbeddingsOptions.ConcatEmbeddingsOptions,
182 BuiltinOptions.LSHProjectionOptions: LSHProjectionOptions.LSHProjectionOptions,
183 BuiltinOptions.Pool2DOptions: Pool2DOptions.Pool2DOptions,
184 BuiltinOptions.SVDFOptions: SVDFOptions.SVDFOptions,
185 BuiltinOptions.RNNOptions: RNNOptions.RNNOptions,
186 BuiltinOptions.FullyConnectedOptions: FullyConnectedOptions.FullyConnectedOptions,
187 BuiltinOptions.SoftmaxOptions: SoftmaxOptions.SoftmaxOptions,
188 BuiltinOptions.ConcatenationOptions: ConcatenationOptions.ConcatenationOptions,
189 BuiltinOptions.AddOptions: AddOptions.AddOptions,
190 BuiltinOptions.L2NormOptions: L2NormOptions.L2NormOptions,
191 BuiltinOptions.LocalResponseNormalizationOptions: LocalResponseNormalizationOptions.LocalResponseNormalizationOptions, # noqa: E501
192 BuiltinOptions.LSTMOptions: LSTMOptions.LSTMOptions,
193 BuiltinOptions.ResizeBilinearOptions: ResizeBilinearOptions.ResizeBilinearOptions,
194 BuiltinOptions.CallOptions: CallOptions.CallOptions,
195 BuiltinOptions.ReshapeOptions: ReshapeOptions.ReshapeOptions,
196 BuiltinOptions.SkipGramOptions: SkipGramOptions.SkipGramOptions,
197 BuiltinOptions.SpaceToDepthOptions: SpaceToDepthOptions.SpaceToDepthOptions,
198 BuiltinOptions.EmbeddingLookupSparseOptions: EmbeddingLookupSparseOptions.EmbeddingLookupSparseOptions,
199 BuiltinOptions.MulOptions: MulOptions.MulOptions,
200 BuiltinOptions.PadOptions: PadOptions.PadOptions,
201 BuiltinOptions.GatherOptions: GatherOptions.GatherOptions,
202 BuiltinOptions.BatchToSpaceNDOptions: BatchToSpaceNDOptions.BatchToSpaceNDOptions,
203 BuiltinOptions.SpaceToBatchNDOptions: SpaceToBatchNDOptions.SpaceToBatchNDOptions,
204 BuiltinOptions.TransposeOptions: TransposeOptions.TransposeOptions,
205 BuiltinOptions.ReducerOptions: ReducerOptions.ReducerOptions,
206 BuiltinOptions.SubOptions: SubOptions.SubOptions,
207 BuiltinOptions.DivOptions: DivOptions.DivOptions,
208 BuiltinOptions.SqueezeOptions: SqueezeOptions.SqueezeOptions,
209 BuiltinOptions.SequenceRNNOptions: SequenceRNNOptions.SequenceRNNOptions,
210 BuiltinOptions.StridedSliceOptions: StridedSliceOptions.StridedSliceOptions,
211 BuiltinOptions.ExpOptions: ExpOptions.ExpOptions,
212 BuiltinOptions.TopKV2Options: TopKV2Options.TopKV2Options,
213 BuiltinOptions.SplitOptions: SplitOptions.SplitOptions,
214 BuiltinOptions.LogSoftmaxOptions: LogSoftmaxOptions.LogSoftmaxOptions,
215 BuiltinOptions.CastOptions: CastOptions.CastOptions,
216 BuiltinOptions.DequantizeOptions: DequantizeOptions.DequantizeOptions,
217 BuiltinOptions.MaximumMinimumOptions: MaximumMinimumOptions.MaximumMinimumOptions,
218 BuiltinOptions.ArgMaxOptions: ArgMaxOptions.ArgMaxOptions,
219 BuiltinOptions.LessOptions: LessOptions.LessOptions,
220 BuiltinOptions.NegOptions: NegOptions.NegOptions,
221 BuiltinOptions.PadV2Options: PadV2Options.PadV2Options,
222 BuiltinOptions.GreaterOptions: GreaterOptions.GreaterOptions,
223 BuiltinOptions.GreaterEqualOptions: GreaterEqualOptions.GreaterEqualOptions,
224 BuiltinOptions.LessEqualOptions: LessEqualOptions.LessEqualOptions,
225 BuiltinOptions.SelectOptions: SelectOptions.SelectOptions,
226 BuiltinOptions.SliceOptions: SliceOptions.SliceOptions,
227 BuiltinOptions.TransposeConvOptions: TransposeConvOptions.TransposeConvOptions,
228 BuiltinOptions.SparseToDenseOptions: SparseToDenseOptions.SparseToDenseOptions,
229 BuiltinOptions.TileOptions: TileOptions.TileOptions,
230 BuiltinOptions.ExpandDimsOptions: ExpandDimsOptions.ExpandDimsOptions,
231 BuiltinOptions.EqualOptions: EqualOptions.EqualOptions,
232 BuiltinOptions.NotEqualOptions: NotEqualOptions.NotEqualOptions,
233 BuiltinOptions.ShapeOptions: ShapeOptions.ShapeOptions,
234 BuiltinOptions.PowOptions: PowOptions.PowOptions,
235 BuiltinOptions.ArgMinOptions: ArgMinOptions.ArgMinOptions,
236 BuiltinOptions.FakeQuantOptions: FakeQuantOptions.FakeQuantOptions,
237 BuiltinOptions.PackOptions: PackOptions.PackOptions,
238 BuiltinOptions.LogicalOrOptions: LogicalOrOptions.LogicalOrOptions,
239 BuiltinOptions.OneHotOptions: OneHotOptions.OneHotOptions,
240 BuiltinOptions.LogicalAndOptions: LogicalAndOptions.LogicalAndOptions,
241 BuiltinOptions.LogicalNotOptions: LogicalNotOptions.LogicalNotOptions,
242 BuiltinOptions.UnpackOptions: UnpackOptions.UnpackOptions,
243 BuiltinOptions.FloorDivOptions: FloorDivOptions.FloorDivOptions,
244 BuiltinOptions.SquareOptions: SquareOptions.SquareOptions,
245 BuiltinOptions.ZerosLikeOptions: ZerosLikeOptions.ZerosLikeOptions,
246 BuiltinOptions.FillOptions: FillOptions.FillOptions,
247 BuiltinOptions.BidirectionalSequenceLSTMOptions: BidirectionalSequenceLSTMOptions.BidirectionalSequenceLSTMOptions,
248 BuiltinOptions.BidirectionalSequenceRNNOptions: BidirectionalSequenceRNNOptions.BidirectionalSequenceRNNOptions,
249 BuiltinOptions.UnidirectionalSequenceLSTMOptions: UnidirectionalSequenceLSTMOptions.UnidirectionalSequenceLSTMOptions, # noqa: E501
250 BuiltinOptions.FloorModOptions: FloorModOptions.FloorModOptions,
251 BuiltinOptions.RangeOptions: RangeOptions.RangeOptions,
252 BuiltinOptions.ResizeNearestNeighborOptions: ResizeNearestNeighborOptions.ResizeNearestNeighborOptions,
253 BuiltinOptions.LeakyReluOptions: LeakyReluOptions.LeakyReluOptions,
254 BuiltinOptions.SquaredDifferenceOptions: SquaredDifferenceOptions.SquaredDifferenceOptions,
255 BuiltinOptions.MirrorPadOptions: MirrorPadOptions.MirrorPadOptions,
256 BuiltinOptions.AbsOptions: AbsOptions.AbsOptions,
257 BuiltinOptions.SplitVOptions: SplitVOptions.SplitVOptions,
258 BuiltinOptions.UniqueOptions: UniqueOptions.UniqueOptions,
259 BuiltinOptions.ReverseV2Options: ReverseV2Options.ReverseV2Options,
260 BuiltinOptions.AddNOptions: AddNOptions.AddNOptions,
261 BuiltinOptions.GatherNdOptions: GatherNdOptions.GatherNdOptions,
262 BuiltinOptions.CosOptions: CosOptions.CosOptions,
263 BuiltinOptions.WhereOptions: WhereOptions.WhereOptions,
264 BuiltinOptions.RankOptions: RankOptions.RankOptions,
265 BuiltinOptions.ReverseSequenceOptions: ReverseSequenceOptions.ReverseSequenceOptions,
266 BuiltinOptions.MatrixDiagOptions: MatrixDiagOptions.MatrixDiagOptions,
267 BuiltinOptions.QuantizeOptions: QuantizeOptions.QuantizeOptions,
268 BuiltinOptions.MatrixSetDiagOptions: MatrixSetDiagOptions.MatrixSetDiagOptions,
269 BuiltinOptions.DensifyOptions: DensifyOptions.DensifyOptions,
270 BuiltinOptions.DepthToSpaceOptions: DepthToSpaceOptions.DepthToSpaceOptions,
Louis Verhaard678645b2020-06-15 15:22:47 +0200271 BuiltinOptions.HardSwishOptions: HardSwishOptions.HardSwishOptions,
Tim Hall79d07d22020-04-27 18:20:16 +0100272 BuiltinOptions.IfOptions: IfOptions.IfOptions,
273 BuiltinOptions.NonMaxSuppressionV4Options: NonMaxSuppressionV4Options.NonMaxSuppressionV4Options,
274 BuiltinOptions.NonMaxSuppressionV5Options: NonMaxSuppressionV5Options.NonMaxSuppressionV5Options,
275 BuiltinOptions.ScatterNdOptions: ScatterNdOptions.ScatterNdOptions,
276 BuiltinOptions.SegmentSumOptions: SegmentSumOptions.SegmentSumOptions,
277 BuiltinOptions.SelectV2Options: SelectV2Options.SelectV2Options,
278 BuiltinOptions.WhileOptions: WhileOptions.WhileOptions,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200279 BuiltinOptions.BatchMatMulOptions: BatchMatMulOptions.BatchMatMulOptions,
Tim Hall79d07d22020-04-27 18:20:16 +0100280}
281
282builtin_options_inv_map = inverse_map(builtin_options_map)
283
284
285def underscore_to_camel_case(s):
286 return "".join(x.title() for x in s.split("_"))
287
288
289def padding_deserialize(x):
290 return padding_map[x]
291
292
293def padding_serialize(builder, x):
294 return padding_inv_map[x]
295
296
297def activation_deserialize(x):
298 return activation_function_map[x]
299
300
301def activation_serialize(builder, x):
302 return activation_function_inv_map[x]
303
304
305def datatype_deserialize(x):
306 return datatype_map[x]
307
308
309def datatype_serialize(builder, x):
310 return datatype_inv_map[x]
311
312
313def identity(x):
314 return x
315
316
317def identity_serialize(builder, x):
318 return x
319
320
321def write_byte_vector(builder, v):
322 builder.StartVector(1, len(v), 1)
323 for e in v[::-1]:
324 builder.PrependByte(e)
325 return builder.EndVector(len(v))
326
327
328def write_int_vector(builder, v):
329 builder.StartVector(4, len(v), 4)
330 for e in v[::-1]:
331 builder.PrependInt32(e)
332 return builder.EndVector(len(v))
333
334
335class OptionsSerializer:
336 def __init__(self, name, members=[]):
337 self.name = name
338 self.module = globals()[self.name]
339 self.cls = getattr(self.module, self.name)
340 self.builtin_opt_type = builtin_options_inv_map[self.cls]
Tim Hall79d07d22020-04-27 18:20:16 +0100341 self.members = []
342 for mem in members:
343 deserialize = identity
344 serialize = identity_serialize
345 is_vector = False
346 if isinstance(mem, tuple):
347 if len(mem) == 3:
348 mem, deserialize, serialize = mem
349 elif len(mem) == 2:
350 mem, is_vector = mem
351 deserialize = tuple
352 serialize = write_int_vector
353 else:
354 assert 0
355 underscore_mem = mem
356 camelcase_mem = underscore_to_camel_case(mem)
357 self.members.append((underscore_mem, camelcase_mem, deserialize, serialize, is_vector))
358
Tim Hallc8310b12020-06-17 14:53:11 +0100359 def deserialize(self, op_data):
360 builtin_options = op_data.BuiltinOptions()
Tim Hall79d07d22020-04-27 18:20:16 +0100361 attrs = {}
Tim Hallc8310b12020-06-17 14:53:11 +0100362 if builtin_options:
Tim Hall79d07d22020-04-27 18:20:16 +0100363 tfattrs = self.cls()
Tim Hallc8310b12020-06-17 14:53:11 +0100364 tfattrs.Init(builtin_options.Bytes, builtin_options.Pos)
Tim Hall79d07d22020-04-27 18:20:16 +0100365 for underscore_mem, camelcase_mem, deserialize, serialize, is_vector in self.members:
366 fun = camelcase_mem
367 if is_vector:
368 fun += "AsNumpy"
369
370 a = deserialize(getattr(tfattrs, fun)())
371 attrs[underscore_mem] = a
372 return attrs
373
374 def serialize(self, builder, attrs):
375 ser_attrs = []
376 for underscore_mem, camelcase_mem, deserialize, serialize, is_vector in self.members:
377 a = serialize(builder, attrs[underscore_mem])
378 ser_attrs.append((camelcase_mem, a))
379
380 getattr(self.module, self.name + "Start")(builder)
381
382 for camelcase_mem, a in ser_attrs:
383 getattr(self.module, self.name + "Add" + camelcase_mem)(builder, a)
384
385 return getattr(self.module, self.name + "End")(builder), None
386
387
388class CustomOptionsSerializer:
Tim Hallc8310b12020-06-17 14:53:11 +0100389 CUSTOM_OPTIONS_NPU_OP = [0x01, 0x04, 0x01] # NpuOp=1, FlexbufferFormat.UINT8=4, byte length=1
390 CUSTOM_OPTIONS_FORMAT_DEFAULT = 0
391
Tim Hall79d07d22020-04-27 18:20:16 +0100392 def __init__(self):
Tim Hall79d07d22020-04-27 18:20:16 +0100393 self.custom_opt_format = 0
394
Tim Hallc8310b12020-06-17 14:53:11 +0100395 def deserialize(self, op_data):
Tim Hall79d07d22020-04-27 18:20:16 +0100396 attrs = {}
Tim Hallc8310b12020-06-17 14:53:11 +0100397 custom_options = op_data.CustomOptionsAsNumpy()
398 attrs["custom_options"] = custom_options
399 attrs["custom_options_format"] = op_data.CustomOptionsFormat()
400
401 if np.array_equal(custom_options, self.CUSTOM_OPTIONS_NPU_OP):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200402 attrs["custom_type"] = CustomType.ExistingNpuOp
Tim Hallc8310b12020-06-17 14:53:11 +0100403
Tim Hall79d07d22020-04-27 18:20:16 +0100404 return attrs
405
406 def serialize(self, builder, attrs):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200407 custom_type = attrs.get("custom_type", CustomType.ThirdPartyOp)
Tim Hallc8310b12020-06-17 14:53:11 +0100408 self.custom_opt_format = attrs.get("custom_options_format", self.CUSTOM_OPTIONS_FORMAT_DEFAULT)
Tim Hall79d07d22020-04-27 18:20:16 +0100409
410 # Set NPU op custom options for the TensorFlow Lite custom operator
Louis Verhaardaee5d752020-09-30 09:01:52 +0200411 if custom_type == CustomType.NpuOp:
Tim Hallc8310b12020-06-17 14:53:11 +0100412 custom_options = self.CUSTOM_OPTIONS_NPU_OP
413 else:
414 custom_options = attrs.get("custom_options", [])
Tim Hall79d07d22020-04-27 18:20:16 +0100415
Tim Hallc8310b12020-06-17 14:53:11 +0100416 custom_options_bytes = struct.pack("<{0}B".format(len(custom_options)), *custom_options)
417 custom_offset = write_byte_vector(builder, custom_options_bytes)
Tim Hall79d07d22020-04-27 18:20:16 +0100418
419 return None, custom_offset
420
421
422padding_map = {
423 Padding.SAME: b"SAME",
424 Padding.VALID: b"VALID",
425}
426
427padding_inv_map = inverse_map(padding_map)
428
429
430activation_function_map = {
431 ActivationFunctionType.NONE: None,
Louis Verhaardaee5d752020-09-30 09:01:52 +0200432 ActivationFunctionType.RELU: Op.Relu,
433 ActivationFunctionType.RELU_N1_TO_1: Op.ReluN1To1,
434 ActivationFunctionType.RELU6: Op.Relu6,
435 ActivationFunctionType.TANH: Op.Tanh,
436 ActivationFunctionType.SIGN_BIT: Op.SignBit,
Tim Hall79d07d22020-04-27 18:20:16 +0100437}
438
439activation_function_inv_map = inverse_map(activation_function_map)
440
441fused_act = ("fused_activation_function", activation_deserialize, activation_serialize)
442padding = ("padding", padding_deserialize, padding_serialize)
443
444pool2d_opts = OptionsSerializer(
445 "Pool2DOptions", (padding, "stride_w", "stride_h", "filter_width", "filter_height", fused_act,)
446)
447
448depthwise_opts = OptionsSerializer(
449 "DepthwiseConv2DOptions",
450 (padding, "stride_w", "stride_h", "depth_multiplier", fused_act, "dilation_w_factor", "dilation_h_factor",),
451)
452
453conv2d_opts = OptionsSerializer(
454 "Conv2DOptions", (padding, "stride_w", "stride_h", fused_act, "dilation_w_factor", "dilation_h_factor",)
455)
456
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200457lstm_opts = OptionsSerializer(
458 "LSTMOptions", (fused_act, "cell_clip", "proj_clip", "kernel_type", "asymmetric_quantize_inputs")
459)
Tim Hall79d07d22020-04-27 18:20:16 +0100460
461unidir_seq_lstm_opts = OptionsSerializer(
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200462 "UnidirectionalSequenceLSTMOptions",
463 (fused_act, "cell_clip", "proj_clip", "time_major", "asymmetric_quantize_inputs",),
Tim Hall79d07d22020-04-27 18:20:16 +0100464)
465
466bidir_seq_lstm_opts = OptionsSerializer(
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200467 "BidirectionalSequenceLSTMOptions",
468 (fused_act, "cell_clip", "proj_clip", "merge_outputs", "time_major", "asymmetric_quantize_inputs"),
Tim Hall79d07d22020-04-27 18:20:16 +0100469)
470
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200471rnn_opts = OptionsSerializer("RNNOptions", (fused_act, "asymmetric_quantize_inputs"))
Tim Hall79d07d22020-04-27 18:20:16 +0100472
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200473seq_rnn_opts = OptionsSerializer("SequenceRNNOptions", ("time_major", fused_act, "asymmetric_quantize_inputs",))
Tim Hall79d07d22020-04-27 18:20:16 +0100474
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200475bidir_seq_rnn_opts = OptionsSerializer(
476 "BidirectionalSequenceRNNOptions", ("time_major", fused_act, "merge_outputs", "asymmetric_quantize_inputs")
477)
Tim Hall79d07d22020-04-27 18:20:16 +0100478
479
480reducer_opts = OptionsSerializer("ReducerOptions", ("keep_dims",))
481
482is_int_vec = True
483
Tim Hall79d07d22020-04-27 18:20:16 +0100484builtin_operator_map = {
Louis Verhaardaee5d752020-09-30 09:01:52 +0200485 BuiltinOperator.ADD: (Op.Add, OptionsSerializer("AddOptions", (fused_act, "pot_scale_int16"))),
486 BuiltinOperator.AVERAGE_POOL_2D: (Op.AvgPool, pool2d_opts),
487 BuiltinOperator.CONCATENATION: (Op.ConcatTFLite, OptionsSerializer("ConcatenationOptions", ("axis", fused_act))),
488 BuiltinOperator.CONV_2D: (Op.Conv2DBias, conv2d_opts),
489 BuiltinOperator.DEPTHWISE_CONV_2D: (Op.DepthwiseConv2DBias, depthwise_opts),
490 BuiltinOperator.DEPTH_TO_SPACE: (Op.DepthToSpace, OptionsSerializer("DepthToSpaceOptions", ("block_size",))),
491 BuiltinOperator.DEQUANTIZE: (Op.Dequantize, OptionsSerializer("DequantizeOptions")),
492 BuiltinOperator.EMBEDDING_LOOKUP: (Op.EmbeddingLookup, None),
493 BuiltinOperator.FLOOR: (Op.Floor, None),
Tim Hall79d07d22020-04-27 18:20:16 +0100494 BuiltinOperator.FULLY_CONNECTED: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200495 Op.FullyConnected,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200496 OptionsSerializer("FullyConnectedOptions", (fused_act, "weights_format", "asymmetric_quantize_inputs")),
Tim Hall79d07d22020-04-27 18:20:16 +0100497 ),
Louis Verhaardaee5d752020-09-30 09:01:52 +0200498 BuiltinOperator.HASHTABLE_LOOKUP: (Op.HashtableLookup, None),
499 BuiltinOperator.L2_NORMALIZATION: (Op.L2Norm, OptionsSerializer("L2NormOptions", (fused_act,))),
500 BuiltinOperator.L2_POOL_2D: (Op.L2Pool2D, pool2d_opts),
Tim Hall79d07d22020-04-27 18:20:16 +0100501 BuiltinOperator.LOCAL_RESPONSE_NORMALIZATION: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200502 Op.LRN,
Tim Hall79d07d22020-04-27 18:20:16 +0100503 OptionsSerializer("LocalResponseNormalizationOptions", ("radius", "bias", "alpha", "beta")),
504 ),
Louis Verhaardaee5d752020-09-30 09:01:52 +0200505 BuiltinOperator.LOGISTIC: (Op.Sigmoid, None),
506 BuiltinOperator.LSH_PROJECTION: (Op.LSHProjection, OptionsSerializer("LSHProjectionOptions", ("type",))),
507 BuiltinOperator.LSTM: (Op.Lstm, lstm_opts),
508 BuiltinOperator.MAX_POOL_2D: (Op.MaxPool, pool2d_opts),
509 BuiltinOperator.MUL: (Op.Mul, OptionsSerializer("MulOptions", (fused_act,))),
510 BuiltinOperator.RELU: (Op.Relu, None),
511 BuiltinOperator.RELU_N1_TO_1: (Op.ReluN1To1, None),
512 BuiltinOperator.RELU6: (Op.Relu6, None),
513 BuiltinOperator.RESHAPE: (Op.Reshape, OptionsSerializer("ReshapeOptions", (("new_shape", is_int_vec),))),
Tim Hall79d07d22020-04-27 18:20:16 +0100514 BuiltinOperator.RESIZE_BILINEAR: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200515 Op.ResizeBilinear,
Tim Hall79d07d22020-04-27 18:20:16 +0100516 OptionsSerializer("ResizeBilinearOptions", ("align_corners", "half_pixel_centers")),
517 ),
Louis Verhaardaee5d752020-09-30 09:01:52 +0200518 BuiltinOperator.RNN: (Op.Rnn, rnn_opts),
519 BuiltinOperator.SOFTMAX: (Op.Softmax, OptionsSerializer("SoftmaxOptions", ("beta",))),
520 BuiltinOperator.SPACE_TO_DEPTH: (Op.SpaceToDepth, OptionsSerializer("SpaceToDepthOptions", ("block_size",))),
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200521 BuiltinOperator.SVDF: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200522 Op.Svdf,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200523 OptionsSerializer("SVDFOptions", ("rank", fused_act, "asymmetric_quantize_inputs")),
524 ),
Louis Verhaardaee5d752020-09-30 09:01:52 +0200525 BuiltinOperator.TANH: (Op.Tanh, None),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200526 BuiltinOperator.CONCAT_EMBEDDINGS: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200527 Op.ConcatEmbeddings,
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200528 OptionsSerializer(
529 "ConcatEmbeddingsOptions",
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200530 (
531 "num_channels",
532 "num_columns_per_channel",
533 "num_columns_per_channel_as_numpy",
534 "num_columns_per_channel_as_length",
535 "embedding_dim_per_channel",
536 "embedding_dim_per_channel_as_numpy",
537 "embedding_dim_per_channel_as_length",
538 ),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200539 ),
540 ),
541 BuiltinOperator.SKIP_GRAM: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200542 Op.SkipGram,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200543 OptionsSerializer("SkipGramOptions", ("ngram_size", "max_skip_size", "include_all_ngrams")),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200544 ),
Louis Verhaardaee5d752020-09-30 09:01:52 +0200545 BuiltinOperator.CALL: (Op.Call, OptionsSerializer("CallOptions", ("subgraph",))),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200546 BuiltinOperator.EMBEDDING_LOOKUP_SPARSE: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200547 Op.EmbeddingLookupSparse,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200548 OptionsSerializer("EmbeddingLookupSparseOptions", ("combiner",)),
Dwight Lidmand2a972d2020-08-18 14:08:10 +0200549 ),
Louis Verhaardaee5d752020-09-30 09:01:52 +0200550 BuiltinOperator.PAD: (Op.Pad, OptionsSerializer("PadOptions")),
551 BuiltinOperator.UNIDIRECTIONAL_SEQUENCE_RNN: (Op.UnidirectionalSequenceRnn, seq_rnn_opts),
552 BuiltinOperator.GATHER: (Op.GatherV2, OptionsSerializer("GatherOptions", ("axis",))),
553 BuiltinOperator.BATCH_TO_SPACE_ND: (Op.BatchToSpaceND, OptionsSerializer("BatchToSpaceNDOptions")),
554 BuiltinOperator.SPACE_TO_BATCH_ND: (Op.SpaceToBatchND, OptionsSerializer("SpaceToBatchNDOptions")),
555 BuiltinOperator.TRANSPOSE: (Op.Transpose, OptionsSerializer("TransposeOptions")),
556 BuiltinOperator.MEAN: (Op.Mean, None),
557 BuiltinOperator.SUB: (Op.Sub, OptionsSerializer("SubOptions", (fused_act, "pot_scale_int16",))),
558 BuiltinOperator.DIV: (Op.Div, OptionsSerializer("DivOptions", (fused_act,))),
559 BuiltinOperator.SQUEEZE: (Op.Squeeze, OptionsSerializer("SqueezeOptions", (("squeeze_dims", is_int_vec),))),
560 BuiltinOperator.UNIDIRECTIONAL_SEQUENCE_LSTM: (Op.UnidirectionalSequenceLstm, unidir_seq_lstm_opts),
Tim Hall79d07d22020-04-27 18:20:16 +0100561 BuiltinOperator.STRIDED_SLICE: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200562 Op.StridedSlice,
Tim Hall79d07d22020-04-27 18:20:16 +0100563 OptionsSerializer(
564 "StridedSliceOptions", ("begin_mask", "end_mask", "ellipsis_mask", "new_axis_mask", "shrink_axis_mask")
565 ),
566 ),
Louis Verhaardaee5d752020-09-30 09:01:52 +0200567 BuiltinOperator.BIDIRECTIONAL_SEQUENCE_RNN: (Op.BidirectionalSequenceRnn, bidir_seq_rnn_opts),
568 BuiltinOperator.EXP: (Op.Exp, OptionsSerializer("ExpOptions")),
569 BuiltinOperator.TOPK_V2: (Op.TopKV2, OptionsSerializer("TopKV2Options")),
570 BuiltinOperator.SPLIT: (Op.Split, OptionsSerializer("SplitOptions", ("num_splits",))),
571 BuiltinOperator.LOG_SOFTMAX: (Op.LogSoftmax, OptionsSerializer("LogSoftmaxOptions")),
572 BuiltinOperator.DELEGATE: (Op.Delegate, None),
573 BuiltinOperator.BIDIRECTIONAL_SEQUENCE_LSTM: (Op.BidirectionalSequenceLstm, bidir_seq_lstm_opts),
Tim Hall79d07d22020-04-27 18:20:16 +0100574 BuiltinOperator.CAST: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200575 Op.Cast,
Tim Hall79d07d22020-04-27 18:20:16 +0100576 OptionsSerializer(
577 "CastOptions",
578 (
579 ("in_data_type", datatype_deserialize, datatype_serialize),
580 ("out_data_type", datatype_deserialize, datatype_serialize),
581 ),
582 ),
583 ),
Louis Verhaardaee5d752020-09-30 09:01:52 +0200584 BuiltinOperator.PRELU: (Op.Prelu, None),
585 BuiltinOperator.MAXIMUM: (Op.Maximum, OptionsSerializer("MaximumMinimumOptions")),
Tim Hall79d07d22020-04-27 18:20:16 +0100586 BuiltinOperator.ARG_MAX: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200587 Op.ArgMax,
Tim Hall79d07d22020-04-27 18:20:16 +0100588 OptionsSerializer("ArgMaxOptions", (("output_type", datatype_deserialize, datatype_serialize),)),
589 ),
Louis Verhaardaee5d752020-09-30 09:01:52 +0200590 BuiltinOperator.MINIMUM: (Op.Minimum, OptionsSerializer("MaximumMinimumOptions")),
591 BuiltinOperator.LESS: (Op.Less, OptionsSerializer("LessOptions")),
592 BuiltinOperator.NEG: (Op.Neg, OptionsSerializer("NegOptions")),
593 BuiltinOperator.PADV2: (Op.PadV2, OptionsSerializer("PadV2Options")),
594 BuiltinOperator.GREATER: (Op.Greater, OptionsSerializer("GreaterOptions")),
595 BuiltinOperator.GREATER_EQUAL: (Op.GreaterEqual, OptionsSerializer("GreaterEqualOptions")),
596 BuiltinOperator.LESS_EQUAL: (Op.LessEqual, OptionsSerializer("LessEqualOptions")),
597 BuiltinOperator.SELECT: (Op.Select, OptionsSerializer("SelectOptions")),
598 BuiltinOperator.SLICE: (Op.Slice, OptionsSerializer("SliceOptions")),
599 BuiltinOperator.SIN: (Op.Sin, None),
Tim Hall79d07d22020-04-27 18:20:16 +0100600 BuiltinOperator.TRANSPOSE_CONV: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200601 Op.Conv2DBackpropInput,
Tim Hall79d07d22020-04-27 18:20:16 +0100602 OptionsSerializer("TransposeConvOptions", (padding, "stride_w", "stride_h")),
603 ),
604 BuiltinOperator.SPARSE_TO_DENSE: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200605 Op.SparseToDense,
Tim Hall79d07d22020-04-27 18:20:16 +0100606 OptionsSerializer("SparseToDenseOptions", ("validate_indices",)),
607 ),
Louis Verhaardaee5d752020-09-30 09:01:52 +0200608 BuiltinOperator.TILE: (Op.Tile, OptionsSerializer("TileOptions")),
609 BuiltinOperator.EXPAND_DIMS: (Op.ExpandDims, OptionsSerializer("ExpandDimsOptions")),
610 BuiltinOperator.EQUAL: (Op.Equal, OptionsSerializer("EqualOptions")),
611 BuiltinOperator.NOT_EQUAL: (Op.NotEqual, OptionsSerializer("NotEqualOptions")),
612 BuiltinOperator.LOG: (Op.Log, None),
613 BuiltinOperator.SUM: (Op.Sum, None),
614 BuiltinOperator.SQRT: (Op.Sqrt, None),
615 BuiltinOperator.RSQRT: (Op.Rsqrt, None),
Tim Hall79d07d22020-04-27 18:20:16 +0100616 BuiltinOperator.SHAPE: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200617 Op.Shape,
Tim Hall79d07d22020-04-27 18:20:16 +0100618 OptionsSerializer("ShapeOptions", (("out_type", datatype_deserialize, datatype_serialize),)),
619 ),
Louis Verhaardaee5d752020-09-30 09:01:52 +0200620 BuiltinOperator.POW: (Op.Pow, OptionsSerializer("PowOptions")),
Tim Hall79d07d22020-04-27 18:20:16 +0100621 BuiltinOperator.ARG_MIN: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200622 Op.ArgMin,
Tim Hall79d07d22020-04-27 18:20:16 +0100623 OptionsSerializer("ArgMinOptions", (("output_type", datatype_deserialize, datatype_serialize),)),
624 ),
625 BuiltinOperator.FAKE_QUANT: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200626 Op.FakeQuantWithMinMaxArgs,
Tim Hall79d07d22020-04-27 18:20:16 +0100627 OptionsSerializer("FakeQuantOptions", ("min", "max", "num_bits", "narrow_range")),
628 ),
Louis Verhaardaee5d752020-09-30 09:01:52 +0200629 BuiltinOperator.REDUCE_PROD: (Op.Prod, reducer_opts),
630 BuiltinOperator.REDUCE_MAX: (Op.Max, reducer_opts),
631 BuiltinOperator.PACK: (Op.Pack, OptionsSerializer("PackOptions", ("values_count", "axis"))),
632 BuiltinOperator.LOGICAL_OR: (Op.LogicalOr, OptionsSerializer("LogicalOrOptions")),
633 BuiltinOperator.ONE_HOT: (Op.OneHot, OptionsSerializer("OneHotOptions", ("axis",))),
634 BuiltinOperator.LOGICAL_AND: (Op.LogicalAnd, OptionsSerializer("LogicalAndOptions")),
635 BuiltinOperator.LOGICAL_NOT: (Op.LogicalNot, OptionsSerializer("LogicalNotOptions")),
636 BuiltinOperator.UNPACK: (Op.Unpack, OptionsSerializer("UnpackOptions", ("num", "axis"))),
637 BuiltinOperator.REDUCE_MIN: (Op.Min, reducer_opts),
638 BuiltinOperator.FLOOR_DIV: (Op.FloorDiv, OptionsSerializer("FloorDivOptions")),
639 BuiltinOperator.REDUCE_ANY: (Op.Any, reducer_opts),
640 BuiltinOperator.SQUARE: (Op.Square, OptionsSerializer("SquareOptions")),
641 BuiltinOperator.ZEROS_LIKE: (Op.ZerosLike, OptionsSerializer("ZerosLikeOptions")),
642 BuiltinOperator.FILL: (Op.Fill, OptionsSerializer("FillOptions")),
643 BuiltinOperator.FLOOR_MOD: (Op.FloorMod, OptionsSerializer("FloorModOptions")),
644 BuiltinOperator.RANGE: (Op.Range, OptionsSerializer("RangeOptions")),
Tim Hall79d07d22020-04-27 18:20:16 +0100645 BuiltinOperator.RESIZE_NEAREST_NEIGHBOR: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200646 Op.ResizeNearestNeighbor,
Jacob Bohlin8daf6b72020-09-15 16:28:35 +0200647 OptionsSerializer("ResizeNearestNeighborOptions", ("align_corners", "half_pixel_centers")),
Tim Hall79d07d22020-04-27 18:20:16 +0100648 ),
Louis Verhaardaee5d752020-09-30 09:01:52 +0200649 BuiltinOperator.LEAKY_RELU: (Op.LeakyRelu, OptionsSerializer("LeakyReluOptions", ("alpha",))),
650 BuiltinOperator.SQUARED_DIFFERENCE: (Op.SquaredDifference, OptionsSerializer("SquaredDifferenceOptions")),
651 BuiltinOperator.MIRROR_PAD: (Op.MirrorPad, OptionsSerializer("MirrorPadOptions", ("mode",))),
652 BuiltinOperator.ABS: (Op.Abs, OptionsSerializer("AbsOptions")),
653 BuiltinOperator.SPLIT_V: (Op.SplitV, OptionsSerializer("SplitVOptions", ("num_splits",))),
Tim Hall79d07d22020-04-27 18:20:16 +0100654 BuiltinOperator.UNIQUE: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200655 Op.Unique,
Tim Hall79d07d22020-04-27 18:20:16 +0100656 OptionsSerializer("UniqueOptions", (("idx_out_type", datatype_deserialize, datatype_serialize),)),
657 ),
Louis Verhaardaee5d752020-09-30 09:01:52 +0200658 BuiltinOperator.CEIL: (Op.Ceil, None),
659 BuiltinOperator.REVERSE_V2: (Op.ReverseV2, OptionsSerializer("ReverseV2Options")),
660 BuiltinOperator.ADD_N: (Op.AddN, OptionsSerializer("AddNOptions")),
661 BuiltinOperator.GATHER_ND: (Op.GatherNd, OptionsSerializer("GatherNdOptions")),
662 BuiltinOperator.COS: (Op.Cos, OptionsSerializer("CosOptions")),
663 BuiltinOperator.WHERE: (Op.Where, OptionsSerializer("WhereOptions")),
664 BuiltinOperator.RANK: (Op.Rank, OptionsSerializer("RankOptions")),
665 BuiltinOperator.ELU: (Op.Elu, None),
Tim Hall79d07d22020-04-27 18:20:16 +0100666 BuiltinOperator.REVERSE_SEQUENCE: (
Louis Verhaardaee5d752020-09-30 09:01:52 +0200667 Op.ReverseSequence,
Tim Hall79d07d22020-04-27 18:20:16 +0100668 OptionsSerializer("ReverseSequenceOptions", ("seq_dim", "batch_dim")),
669 ),
Louis Verhaardaee5d752020-09-30 09:01:52 +0200670 BuiltinOperator.MATRIX_DIAG: (Op.MatrixDiag, OptionsSerializer("MatrixDiagOptions")),
671 BuiltinOperator.QUANTIZE: (Op.Quantize, OptionsSerializer("QuantizeOptions")),
672 BuiltinOperator.MATRIX_SET_DIAG: (Op.MatrixSetDiag, OptionsSerializer("MatrixSetDiagOptions")),
673 BuiltinOperator.ROUND: (Op.Round, None),
674 BuiltinOperator.HARD_SWISH: (Op.HardSwish, OptionsSerializer("HardSwishOptions")),
675 BuiltinOperator.IF: (Op.If, OptionsSerializer("IfOptions", ("then_subgraph_index", "else_subgraph_index"))),
676 BuiltinOperator.WHILE: (
677 Op.While,
678 OptionsSerializer("WhileOptions", ("cond_subgraph_index", "body_subgraph_index")),
679 ),
680 BuiltinOperator.NON_MAX_SUPPRESSION_V4: (Op.NonMaxSuppressionV4, OptionsSerializer("NonMaxSuppressionV4Options")),
681 BuiltinOperator.NON_MAX_SUPPRESSION_V5: (Op.NonMaxSuppressionV5, OptionsSerializer("NonMaxSuppressionV5Options")),
682 BuiltinOperator.SCATTER_ND: (Op.ScatterNd, OptionsSerializer("ScatterNdOptions")),
683 BuiltinOperator.SELECT_V2: (Op.SelectV2, OptionsSerializer("SelectV2Options")),
684 BuiltinOperator.DENSIFY: (Op.Densify, OptionsSerializer("DensifyOptions")),
685 BuiltinOperator.SEGMENT_SUM: (Op.SegmentSum, OptionsSerializer("SegmentSumOptions")),
686 BuiltinOperator.BATCH_MATMUL: (Op.BatchMatMul, OptionsSerializer("BatchMatMulOptions", ("adj_x", "adj_y"))),
687 BuiltinOperator.CUSTOM: (Op.Custom, CustomOptionsSerializer()),
Tim Hall79d07d22020-04-27 18:20:16 +0100688}
689
690builtin_operator_inv_map = {v[0]: (k, v[1]) for k, v in builtin_operator_map.items()}
691
Louis Verhaardaee5d752020-09-30 09:01:52 +0200692builtin_operator_inv_map[Op.CustomNpuOp] = (BuiltinOperator.CUSTOM, CustomOptionsSerializer())
Michael McGeagh219ec072020-11-09 11:11:26 +0000693
Michael McGeagh837dc1b2020-11-10 12:38:25 +0000694BUILTIN_OPERATOR_UNKNOWN = "UNKNOWN"
695
696
697def builtin_type_name(builtin):
698 return next(k for k, v in vars(BuiltinOperator).items() if v == builtin)
699
Michael McGeagh219ec072020-11-09 11:11:26 +0000700
701def optype_to_builtintype(op_type):
702 if op_type in builtin_operator_inv_map:
Michael McGeagh837dc1b2020-11-10 12:38:25 +0000703 return builtin_type_name(builtin_operator_inv_map[op_type][0])
Michael McGeagh219ec072020-11-09 11:11:26 +0000704 else:
Michael McGeagh837dc1b2020-11-10 12:38:25 +0000705 return BUILTIN_OPERATOR_UNKNOWN