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

String input support on TFLite Gather_ND op

Browse Source 
PiperOrigin-RevId: 304934439
Change-Id: Ifd8fbfd330275c882e2a69a95d24a68f3de59945
This commit is contained in:
Hyeonjong Ryu 2020-04-05 17:40:04 -07:00 committed by TensorFlower Gardener
parent a930563d96
commit 64abc48387
9 changed files with 255 additions and 157 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/compiler/mlir/lite/ir/tfl_ops.td
View File

@ -830,12 +830,12 @@ def TFL_GatherNdOp : TFL_Op<"gather_nd", [NoSideEffect]> {
}];
let arguments = (ins
TFL_TensorOf<[F32, I8, I64, I32, TFL_Uint8]>:$params,
TFL_TensorOf<[F32, I8, I64, I32, TFL_Uint8, TFL_Str]>:$params,
TFL_I32OrI64Tensor:$indices
);
let results = (outs
TFL_TensorOf<[F32, I8, I64, I32, TFL_Uint8]>:$output
TFL_TensorOf<[F32, I8, I64, I32, TFL_Uint8, TFL_Str]>:$output
);
}

12
tensorflow/lite/kernels/gather_nd.cc
View File

@ -41,6 +41,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
case kTfLiteInt8:
case kTfLiteInt64:
case kTfLiteInt32:
case kTfLiteString:
break;
default:
context->ReportError(
@ -103,6 +104,15 @@ TfLiteStatus GatherNd(const TfLiteTensor* params, const TfLiteTensor* indices,
return kTfLiteOk;
}
template <typename IndicesT>
TfLiteStatus GatherNdString(const TfLiteTensor* params,
const TfLiteTensor* indices, TfLiteTensor* output) {
reference_ops::GatherNdString(
GetTensorShape(params), params, GetTensorShape(indices),
GetTensorData<IndicesT>(indices), GetTensorShape(output), output);
return kTfLiteOk;
}
template <typename IndicesT>
TfLiteStatus EvalGatherNd(TfLiteContext* context, const TfLiteTensor* params,
const TfLiteTensor* indices, TfLiteTensor* output) {
@ -117,6 +127,8 @@ TfLiteStatus EvalGatherNd(TfLiteContext* context, const TfLiteTensor* params,
return GatherNd<int32_t, IndicesT>(params, indices, output);
case kTfLiteInt64:
return GatherNd<int64_t, IndicesT>(params, indices, output);
case kTfLiteString:
return GatherNdString<IndicesT>(params, indices, output);
default:
context->ReportError(context,
"Params type '%s' are not supported by gather_nd.",

23
tensorflow/lite/kernels/gather_nd_test.cc
View File

@ -313,5 +313,28 @@ TEST(GatherNdOpTest, Int64Int64) {
ElementsAreArray({-2LL, 2LL, 2LL, 3LL, 3LL, -3LL}));
}
TEST(GatherNdOpTest, StringInt32) {
GatherNdOpModel m({TensorType_STRING, {3, 2, 3}}, {TensorType_INT32, {2, 2}});
m.SetInput<std::string>({"A", "B", "C", "D", "E", "F", //
"G", "H", "I", "J", "K", "L", //
"M", "N", "O", "P", "Q", "R"});
m.SetPositions<int32_t>({0, 1, 1, 0});
m.Invoke();
EXPECT_THAT(m.GetOutput<std::string>(),
ElementsAreArray({"D", "E", "F", "G", "H", "I"}));
}
TEST(GatherNdOpTest, StringInt64) {
GatherNdOpModel m({TensorType_STRING, {3, 2, 3}}, {TensorType_INT64, {2, 2}});
m.SetInput<std::string>({"A", "B", "C", "D", "E", "F", //
"G", "H", "I", "J", "K", "L", //
"M", "N", "O", "P", "Q", "R"});
m.SetPositions<int64_t>({0LL, 1LL, 1LL, 0LL});
m.Invoke();
EXPECT_THAT(m.GetOutput<std::string>(),
ElementsAreArray({"D", "E", "F", "G", "H", "I"}));
}
} // namespace
} // namespace tflite

87
tensorflow/lite/kernels/internal/reference/reference_ops.h
View File

@ -1561,6 +1561,40 @@ inline void Gather(const tflite::GatherParams& op_params,
}
}
// Common subroutine for both `GatherNd` and `GatherNdString`.
struct GatherNdHelperResult {
int n_slices;
int slice_size;
int indices_nd;
std::vector<int> dims_to_count;
};
// Returns common values being used on both `GatherNd` and `GatherNdString`.
inline GatherNdHelperResult GatherNdHelper(const RuntimeShape& params_shape,
const RuntimeShape& indices_shape) {
GatherNdHelperResult ret;
ret.n_slices = 1;
ret.slice_size = 1;
const int indices_dims = indices_shape.DimensionsCount();
ret.indices_nd = indices_shape.Dims(indices_dims - 1);
const int params_dims = params_shape.DimensionsCount();
for (int i = 0; i < indices_dims - 1; ++i) {
ret.n_slices *= indices_shape.Dims(i);
}
for (int i = ret.indices_nd; i < params_dims; ++i) {
ret.slice_size *= params_shape.Dims(i);
}
int remain_flat_size = params_shape.FlatSize();
ret.dims_to_count = std::vector<int>(ret.indices_nd, 0);
for (int i = 0; i < ret.indices_nd; ++i) {
ret.dims_to_count[i] = remain_flat_size / params_shape.Dims(i);
remain_flat_size = ret.dims_to_count[i];
}
return ret;
}
template <typename ParamsT, typename IndicesT = int32>
inline void GatherNd(const RuntimeShape& params_shape,
const ParamsT* params_data,
@ -1569,35 +1603,40 @@ inline void GatherNd(const RuntimeShape& params_shape,
const RuntimeShape& output_shape, ParamsT* output_data) {
ruy::profiler::ScopeLabel label("GatherNd");
int n_slices = 1;
int slice_size = 1;
const int indices_dims = indices_shape.DimensionsCount();
const int indices_nd = indices_shape.Dims(indices_dims - 1);
const int params_dims = params_shape.DimensionsCount();
for (int i = 0; i < indices_dims - 1; ++i) {
n_slices *= indices_shape.Dims(i);
}
for (int i = indices_nd; i < params_dims; ++i) {
slice_size *= params_shape.Dims(i);
}
int remain_flat_size = params_shape.FlatSize();
std::vector<int> dims_to_count(indices_nd, 0);
for (int i = 0; i < indices_nd; ++i) {
dims_to_count[i] = remain_flat_size / params_shape.Dims(i);
remain_flat_size = dims_to_count[i];
}
for (int i = 0; i < n_slices; ++i) {
const GatherNdHelperResult res = GatherNdHelper(params_shape, indices_shape);
for (int i = 0; i < res.n_slices; ++i) {
int from_pos = 0;
for (int j = 0; j < indices_nd; ++j) {
from_pos += indices_data[i * indices_nd + j] * dims_to_count[j];
for (int j = 0; j < res.indices_nd; ++j) {
from_pos += indices_data[i * res.indices_nd + j] * res.dims_to_count[j];
}
std::memcpy(output_data + i * slice_size, params_data + from_pos,
sizeof(ParamsT) * slice_size);
std::memcpy(output_data + i * res.slice_size, params_data + from_pos,
sizeof(ParamsT) * res.slice_size);
}
}
template <typename IndicesT = int32>
inline void GatherNdString(const RuntimeShape& params_shape,
const TfLiteTensor* params_data,
const RuntimeShape& indices_shape,
const IndicesT* indices_data,
const RuntimeShape& output_shape,
TfLiteTensor* output_data) {
ruy::profiler::ScopeLabel label("GatherNdString");
const GatherNdHelperResult res = GatherNdHelper(params_shape, indices_shape);
DynamicBuffer buffer;
for (int i = 0; i < res.n_slices; ++i) {
int from_pos = 0;
for (int j = 0; j < res.indices_nd; ++j) {
from_pos += indices_data[i * res.indices_nd + j] * res.dims_to_count[j];
}
for (int j = 0; j < res.slice_size; ++j) {
buffer.AddString(GetString(params_data, from_pos + j));
}
}
buffer.WriteToTensor(output_data, /*new_shape=*/nullptr);
}
template <typename IndicesT, typename UpdatesT>
inline void ScatterNd(const RuntimeShape& indices_shape,
const IndicesT* indices_data,

256
tensorflow/lite/kernels/register.cc
View File

@ -33,199 +33,199 @@ namespace builtin {
BuiltinOpResolver::BuiltinOpResolver() {
AddBuiltin(BuiltinOperator_ABS, Register_ABS());
AddBuiltin(BuiltinOperator_HARD_SWISH, Register_HARD_SWISH());
AddBuiltin(BuiltinOperator_RELU, Register_RELU(), /* min_version */ 1,
/* max_version */ 2);
AddBuiltin(BuiltinOperator_RELU, Register_RELU(), /* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_RELU_N1_TO_1, Register_RELU_N1_TO_1());
AddBuiltin(BuiltinOperator_RELU6, Register_RELU6(), /* min_version */ 1,
/* max_version */ 2);
AddBuiltin(BuiltinOperator_TANH, Register_TANH(), /* min_version */ 1,
/* max_version */ 2);
AddBuiltin(BuiltinOperator_RELU6, Register_RELU6(), /* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_TANH, Register_TANH(), /* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_LOGISTIC, Register_LOGISTIC(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_AVERAGE_POOL_2D, Register_AVERAGE_POOL_2D(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_MAX_POOL_2D, Register_MAX_POOL_2D(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_L2_POOL_2D, Register_L2_POOL_2D());
AddBuiltin(BuiltinOperator_CONV_2D, Register_CONV_2D(),
/* min_version */ 1,
/* max_version */ 3);
/* min_version = */ 1,
/* max_version = */ 3);
AddBuiltin(BuiltinOperator_DEPTHWISE_CONV_2D, Register_DEPTHWISE_CONV_2D(),
/* min_version */ 1,
/* max_version */ 3);
/* min_version = */ 1,
/* max_version = */ 3);
AddBuiltin(BuiltinOperator_SVDF, Register_SVDF(),
/* min_version */ 1,
/* max_version */ 3);
/* min_version = */ 1,
/* max_version = */ 3);
AddBuiltin(BuiltinOperator_RNN, Register_RNN(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN,
Register_BIDIRECTIONAL_SEQUENCE_RNN(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN,
Register_UNIDIRECTIONAL_SEQUENCE_RNN(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_EMBEDDING_LOOKUP, Register_EMBEDDING_LOOKUP(),
/* min_version */ 1,
/* max_version */ 3);
/* min_version = */ 1,
/* max_version = */ 3);
AddBuiltin(BuiltinOperator_EMBEDDING_LOOKUP_SPARSE,
Register_EMBEDDING_LOOKUP_SPARSE());
AddBuiltin(BuiltinOperator_FULLY_CONNECTED, Register_FULLY_CONNECTED(),
/* min_version */ 1,
/* max_version */ 6);
/* min_version = */ 1,
/* max_version = */ 6);
AddBuiltin(BuiltinOperator_LSH_PROJECTION, Register_LSH_PROJECTION());
AddBuiltin(BuiltinOperator_HASHTABLE_LOOKUP, Register_HASHTABLE_LOOKUP());
AddBuiltin(BuiltinOperator_SOFTMAX, Register_SOFTMAX(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_CONCATENATION, Register_CONCATENATION(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_ADD, Register_ADD(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_SPACE_TO_BATCH_ND, Register_SPACE_TO_BATCH_ND(),
/* min_version */ 1,
/* max_version */ 3);
/* min_version = */ 1,
/* max_version = */ 3);
AddBuiltin(BuiltinOperator_BATCH_TO_SPACE_ND, Register_BATCH_TO_SPACE_ND(),
/* min_version */ 1,
/* max_version */ 3);
AddBuiltin(BuiltinOperator_MUL, Register_MUL(), /* min_version */ 1,
/* max_version */ 3);
/* min_version = */ 1,
/* max_version = */ 3);
AddBuiltin(BuiltinOperator_MUL, Register_MUL(), /* min_version = */ 1,
/* max_version = */ 3);
AddBuiltin(BuiltinOperator_L2_NORMALIZATION, Register_L2_NORMALIZATION(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_LOCAL_RESPONSE_NORMALIZATION,
Register_LOCAL_RESPONSE_NORMALIZATION());
AddBuiltin(BuiltinOperator_LSTM, Register_LSTM(), /* min_version */ 1,
/* max_version */ 3);
AddBuiltin(BuiltinOperator_LSTM, Register_LSTM(), /* min_version = */ 1,
/* max_version = */ 3);
AddBuiltin(BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM,
Register_BIDIRECTIONAL_SEQUENCE_LSTM(), /* min_version */ 1,
/* max_version */ 3);
Register_BIDIRECTIONAL_SEQUENCE_LSTM(), /* min_version = */ 1,
/* max_version = */ 3);
AddBuiltin(BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM,
Register_UNIDIRECTIONAL_SEQUENCE_LSTM(), /* min_version */ 1,
/* max_version */ 2);
AddBuiltin(BuiltinOperator_PAD, Register_PAD(), /* min_version */ 1,
/* max_version */ 2);
AddBuiltin(BuiltinOperator_PADV2, Register_PADV2(), /* min_version */ 1,
/* max_version */ 2);
Register_UNIDIRECTIONAL_SEQUENCE_LSTM(), /* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_PAD, Register_PAD(), /* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_PADV2, Register_PADV2(), /* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_RESHAPE, Register_RESHAPE());
AddBuiltin(BuiltinOperator_RESIZE_BILINEAR, Register_RESIZE_BILINEAR(),
/* min_version */ 1,
/* max_version */ 3);
/* min_version = */ 1,
/* max_version = */ 3);
AddBuiltin(BuiltinOperator_RESIZE_NEAREST_NEIGHBOR,
Register_RESIZE_NEAREST_NEIGHBOR(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_SKIP_GRAM, Register_SKIP_GRAM());
AddBuiltin(BuiltinOperator_SPACE_TO_DEPTH, Register_SPACE_TO_DEPTH(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_DEPTH_TO_SPACE, Register_DEPTH_TO_SPACE());
AddBuiltin(BuiltinOperator_GATHER, Register_GATHER(),
/* min_version */ 1,
/* max_version */ 3);
/* min_version = */ 1,
/* max_version = */ 3);
AddBuiltin(BuiltinOperator_TRANSPOSE, Register_TRANSPOSE(),
/* min_version */ 1,
/* max_version */ 4);
/* min_version = */ 1,
/* max_version = */ 4);
AddBuiltin(BuiltinOperator_MEAN, Register_MEAN(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_DIV, Register_DIV());
AddBuiltin(BuiltinOperator_SUB, Register_SUB(),
/* min_version */ 1,
/* max_version */ 3);
AddBuiltin(BuiltinOperator_SPLIT, Register_SPLIT(), /* min_version */ 1,
/* max_version */ 3);
/* min_version = */ 1,
/* max_version = */ 3);
AddBuiltin(BuiltinOperator_SPLIT, Register_SPLIT(), /* min_version = */ 1,
/* max_version = */ 3);
AddBuiltin(BuiltinOperator_SPLIT_V, Register_SPLIT_V(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_SQUEEZE, Register_SQUEEZE());
AddBuiltin(BuiltinOperator_STRIDED_SLICE, Register_STRIDED_SLICE(),
/* min_version */ 1,
/* max_version */ 4);
/* min_version = */ 1,
/* max_version = */ 4);
AddBuiltin(BuiltinOperator_EXP, Register_EXP());
AddBuiltin(BuiltinOperator_TOPK_V2, Register_TOPK_V2(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_LOG, Register_LOG());
AddBuiltin(BuiltinOperator_LOG_SOFTMAX, Register_LOG_SOFTMAX(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_CAST, Register_CAST());
AddBuiltin(BuiltinOperator_DEQUANTIZE, Register_DEQUANTIZE(),
/* min_version */ 1,
/* max_version */ 4);
/* min_version = */ 1,
/* max_version = */ 4);
AddBuiltin(BuiltinOperator_PRELU, Register_PRELU());
AddBuiltin(BuiltinOperator_MAXIMUM, Register_MAXIMUM(),
/* min_version */ 1,
/* max_version */ 4);
/* min_version = */ 1,
/* max_version = */ 4);
AddBuiltin(BuiltinOperator_MINIMUM, Register_MINIMUM(),
/* min_version */ 1,
/* max_version */ 4);
/* min_version = */ 1,
/* max_version = */ 4);
AddBuiltin(BuiltinOperator_ARG_MAX, Register_ARG_MAX(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_ARG_MIN, Register_ARG_MIN(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_GREATER, Register_GREATER(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_GREATER_EQUAL, Register_GREATER_EQUAL(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_LESS, Register_LESS(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_LESS_EQUAL, Register_LESS_EQUAL(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_FLOOR, Register_FLOOR());
AddBuiltin(BuiltinOperator_CEIL, Register_CEIL());
AddBuiltin(BuiltinOperator_ROUND, Register_ROUND());
AddBuiltin(BuiltinOperator_NEG, Register_NEG());
AddBuiltin(BuiltinOperator_SELECT, Register_SELECT(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_SELECT_V2, Register_SELECT_V2());
AddBuiltin(BuiltinOperator_SLICE, Register_SLICE(),
/* min_version */ 1,
/* max_version */ 3);
/* min_version = */ 1,
/* max_version = */ 3);
AddBuiltin(BuiltinOperator_SIN, Register_SIN());
AddBuiltin(BuiltinOperator_COS, Register_COS());
AddBuiltin(BuiltinOperator_TRANSPOSE_CONV, Register_TRANSPOSE_CONV(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_TILE, Register_TILE(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_SUM, Register_SUM(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_REDUCE_PROD, Register_REDUCE_PROD());
AddBuiltin(BuiltinOperator_REDUCE_MAX, Register_REDUCE_MAX(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_REDUCE_MIN, Register_REDUCE_MIN(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_REDUCE_ANY, Register_REDUCE_ANY());
AddBuiltin(BuiltinOperator_EXPAND_DIMS, Register_EXPAND_DIMS());
AddBuiltin(BuiltinOperator_SPARSE_TO_DENSE, Register_SPARSE_TO_DENSE(),
/* min_version */ 1,
/* max_version */ 3);
/* min_version = */ 1,
/* max_version = */ 3);
AddBuiltin(BuiltinOperator_EQUAL, Register_EQUAL(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_NOT_EQUAL, Register_NOT_EQUAL(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_SQRT, Register_SQRT());
AddBuiltin(BuiltinOperator_RSQRT, Register_RSQRT());
AddBuiltin(BuiltinOperator_SHAPE, Register_SHAPE());
@ -233,41 +233,43 @@ BuiltinOpResolver::BuiltinOpResolver() {
AddBuiltin(BuiltinOperator_POW, Register_POW());
AddBuiltin(BuiltinOperator_FAKE_QUANT, Register_FAKE_QUANT(), 1, 2);
AddBuiltin(BuiltinOperator_PACK, Register_PACK(),
/* min_version */ 1,
/* max_version */ 3);
/* min_version = */ 1,
/* max_version = */ 3);
AddBuiltin(BuiltinOperator_ONE_HOT, Register_ONE_HOT());
AddBuiltin(BuiltinOperator_LOGICAL_OR, Register_LOGICAL_OR());
AddBuiltin(BuiltinOperator_LOGICAL_AND, Register_LOGICAL_AND());
AddBuiltin(BuiltinOperator_LOGICAL_NOT, Register_LOGICAL_NOT());
AddBuiltin(BuiltinOperator_UNPACK, Register_UNPACK(),
/* min_version */ 1,
/* max_version */ 4);
/* min_version = */ 1,
/* max_version = */ 4);
AddBuiltin(BuiltinOperator_FLOOR_DIV, Register_FLOOR_DIV(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_SQUARE, Register_SQUARE());
AddBuiltin(BuiltinOperator_ZEROS_LIKE, Register_ZEROS_LIKE());
AddBuiltin(BuiltinOperator_FLOOR_MOD, Register_FLOOR_MOD());
AddBuiltin(BuiltinOperator_RANGE, Register_RANGE());
AddBuiltin(BuiltinOperator_LEAKY_RELU, Register_LEAKY_RELU(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_SQUARED_DIFFERENCE, Register_SQUARED_DIFFERENCE());
AddBuiltin(BuiltinOperator_FILL, Register_FILL());
AddBuiltin(BuiltinOperator_MIRROR_PAD, Register_MIRROR_PAD());
AddBuiltin(BuiltinOperator_UNIQUE, Register_UNIQUE());
AddBuiltin(BuiltinOperator_REVERSE_V2, Register_REVERSE_V2(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_ADD_N, Register_ADD_N());
AddBuiltin(BuiltinOperator_GATHER_ND, Register_GATHER_ND());
AddBuiltin(BuiltinOperator_GATHER_ND, Register_GATHER_ND(),
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_WHERE, Register_WHERE());
AddBuiltin(BuiltinOperator_ELU, Register_ELU());
AddBuiltin(BuiltinOperator_REVERSE_SEQUENCE, Register_REVERSE_SEQUENCE());
AddBuiltin(BuiltinOperator_MATRIX_DIAG, Register_MATRIX_DIAG());
AddBuiltin(BuiltinOperator_QUANTIZE, Register_QUANTIZE(),
/* min_version */ 1,
/* max_version */ 2);
/* min_version = */ 1,
/* max_version = */ 2);
AddBuiltin(BuiltinOperator_MATRIX_SET_DIAG, Register_MATRIX_SET_DIAG());
AddBuiltin(BuiltinOperator_IF, tflite::ops::builtin::Register_IF());
AddBuiltin(BuiltinOperator_WHILE, tflite::ops::builtin::Register_WHILE());

6
tensorflow/lite/testing/op_tests/gather_nd.py
View File

@ -29,19 +29,19 @@ def make_gather_nd_tests(options):
test_parameters = [
{
"params_dtype": [tf.float32, tf.int32, tf.int64],
"params_dtype": [tf.float32, tf.int32, tf.int64, tf.string],
"params_shape": [[5, 1]],
"indices_dtype": [tf.int32, tf.int64],
"indices_shape": [[1, 1]],
},
{
"params_dtype": [tf.float32, tf.int32, tf.int64],
"params_dtype": [tf.float32, tf.int32, tf.int64, tf.string],
"params_shape": [[5, 5]],
"indices_dtype": [tf.int32, tf.int64],
"indices_shape": [[2, 1], [2, 2]],
},
{
"params_dtype": [tf.float32, tf.int32, tf.int64],
"params_dtype": [tf.float32, tf.int32, tf.int64, tf.string],
"params_shape": [[5, 5, 10]],
"indices_dtype": [tf.int32, tf.int64],
"indices_shape": [[3, 1], [2, 2], [2, 3], [2, 1, 3]],

2
tensorflow/lite/toco/tflite/op_version.cc
View File

@ -73,7 +73,7 @@ string GetMinimumRuntimeVersionForModel(const Model& model) {
{{OperatorType::kGather, 1}, "1.6.0"},
{{OperatorType::kGather, 2}, "1.14.0"},
{{OperatorType::kGather, 3}, "1.15.0"},
{{OperatorType::kGatherNd, 1}, "1.14.0"},
{{OperatorType::kGatherNd, 2}, kPendingReleaseOpVersion},
{{OperatorType::kSvdf, 1}, "1.5.0"},
{{OperatorType::kSvdf, 2}, "1.14.0"},
{{OperatorType::kSvdf, 3}, "2.2.0"},

7
tensorflow/lite/tools/versioning/op_version.cc
View File

@ -373,6 +373,13 @@ int GetBuiltinOperatorVersion(const OpSignature& op_sig) {
}
return 1;
case BuiltinOperator_GATHER_ND:
if (!op_sig.input_types.empty() &&
op_sig.input_types.at(0) == TensorType_STRING) {
return 2;
}
return 1;
case BuiltinOperator_AVERAGE_POOL_2D:
case BuiltinOperator_ADD:
case BuiltinOperator_CONCATENATION:

15
tensorflow/lite/tools/versioning/op_version_test.cc
View File

@ -525,4 +525,19 @@ TEST(OpVersionTest, VersioningTransposeTest) {
};
EXPECT_EQ(GetBuiltinOperatorVersion(fake_op_sig), 1);
}
TEST(OpVersionTest, VersioningGatherNdOperatorTest) {
OpSignature fake_op_sig = {
.op = BuiltinOperator_GATHER_ND,
.input_types =
std::vector<TensorType>{TensorType_INT32, TensorType_INT32},
};
EXPECT_EQ(GetBuiltinOperatorVersion(fake_op_sig), 1);
fake_op_sig = {
.op = BuiltinOperator_GATHER_ND,
.input_types =
std::vector<TensorType>{TensorType_STRING, TensorType_INT32},
};
EXPECT_EQ(GetBuiltinOperatorVersion(fake_op_sig), 2);
}
} // namespace tflite