experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Refactor for case BuiltinOperator_GATHER in flatbuffer_conversions

Browse Source
This commit is contained in:
rsun 2021-01-06 12:13:45 -08:00
parent d57f956517
commit 0f1b6628fc
2 changed files with 25 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
tensorflow/lite/core/api/flatbuffer_conversions.cc
View File

@ -608,15 +608,7 @@ TfLiteStatus ParseOpDataTfLite(const Operator* op, BuiltinOperator op_type,
return kTfLiteOk;
}
case BuiltinOperator_GATHER: {
auto params = safe_allocator.Allocate<TfLiteGatherParams>();
TF_LITE_ENSURE(error_reporter, params != nullptr);
params->axis = 0;
if (const auto* gather_params = op->builtin_options_as_GatherOptions()) {
params->axis = gather_params->axis();
}
*builtin_data = params.release();
return kTfLiteOk;
return ParseGather(op, error_reporter, allocator, builtin_data);
}
case BuiltinOperator_SQUEEZE: {
@ -1196,6 +1188,27 @@ TfLiteStatus ParseFullyConnected(const Operator* op,
return kTfLiteOk;
}
// We have this parse function instead of directly returning kTfLiteOk from the
// switch-case in ParseOpData because this function is used as part of the
// selective registration for the OpResolver implementation in micro.
TfLiteStatus ParseGather(const Operator* op,
ErrorReporter* error_reporter,
BuiltinDataAllocator* allocator,
void** builtin_data) {
CheckParsePointerParams(op, error_reporter, allocator, builtin_data);
SafeBuiltinDataAllocator safe_allocator(allocator);
auto params = safe_allocator.Allocate<TfLiteGatherParams>();
TF_LITE_ENSURE(error_reporter, params != nullptr);
params->axis = 0;
if (const auto* gather_params = op->builtin_options_as_GatherOptions()) {
params->axis = gather_params->axis();
}
*builtin_data = params.release();
return kTfLiteOk;
}
// We have this parse function instead of directly returning kTfLiteOk from the
// switch-case in ParseOpData because this function is used as part of the
// selective registration for the OpResolver implementation in micro.

3
tensorflow/lite/core/api/flatbuffer_conversions.h
View File

@ -131,6 +131,9 @@ TfLiteStatus ParseFullyConnected(const Operator* op,
BuiltinDataAllocator* allocator,
void** builtin_data);
TfLiteStatus ParseGather(const Operator* op, ErrorReporter* error_reporter,
BuiltinDataAllocator* allocator, void** builtin_data);
TfLiteStatus ParseGreater(const Operator* op, ErrorReporter* error_reporter,
BuiltinDataAllocator* allocator, void** builtin_data);