wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 5 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/lite/kernels/variants/BUILD
+376
View File
@@ -0,0 +1,376 @@
load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("@rules_cc//cc:cc_test.bzl", "cc_test")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
load("//tensorflow/lite:build_def.bzl", "tflite_copts", "tflite_copts_warnings", "tflite_linkopts")
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:LICENSE"])
common_copts = tflite_copts() + tflite_copts_warnings()
# Library of list ops.
cc_library(
name = "list_ops_lib",
srcs = glob(
["list_kernels/*.cc"],
exclude = ["list_kernels/*test*"],
) + ["register_list_ops.cc"],
hdrs = [
"list_ops_lib.h",
"register_list_ops.h",
],
visibility = ["//visibility:public"],
deps = [
":list_ops_util",
":tensor_array",
"//tensorflow/lite:array",
"//tensorflow/lite:mutable_op_resolver",
"//tensorflow/lite:util",
"//tensorflow/lite/c:c_api_types",
"//tensorflow/lite/core/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/kernels:cpu_backend_context",
"//tensorflow/lite/kernels:kernel_util",
"//tensorflow/lite/kernels/internal:optimized_base",
"//tensorflow/lite/kernels/internal:runtime_shape",
"//tensorflow/lite/kernels/internal:tensor",
"//tensorflow/lite/kernels/internal:tensor_ctypes",
"//tensorflow/lite/schema:schema_fbs",
],
)
# Utils for simple kernel unit tests.
cc_library(
name = "test_util",
testonly = True,
srcs = ["list_kernels/test_util.cc"],
hdrs = ["list_kernels/test_util.h"],
deps = [
":tensor_array",
"//tensorflow/lite:array",
"//tensorflow/lite:type_to_tflitetype",
"//tensorflow/lite:util",
"//tensorflow/lite/c:common",
"//tensorflow/lite/core/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/kernels:kernel_util",
"//tensorflow/lite/kernels:op_macros",
"//tensorflow/lite/kernels:test_util",
"//tensorflow/lite/schema:schema_fbs",
"@com_google_absl//absl/types:span",
"@com_google_googletest//:gtest_main",
],
)
cc_test(
name = "list_reserve_test",
srcs = ["list_kernels/list_reserve_test.cc"],
deps = [
":list_ops_lib",
":tensor_array",
"//tensorflow/lite/core/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/kernels:test_util",
"//tensorflow/lite/schema:schema_fbs",
"@com_google_googletest//:gtest_main",
],
)
cc_test(
name = "list_stack_test",
srcs = ["list_kernels/list_stack_test.cc"],
deps = [
":list_ops_lib",
":test_util",
"//tensorflow/lite/core/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/kernels:test_util",
"//tensorflow/lite/schema:schema_fbs",
"@com_google_googletest//:gtest_main",
],
)
cc_test(
name = "list_set_item_test",
srcs = ["list_kernels/list_set_item_test.cc"],
deps = [
":list_ops_lib",
":tensor_array",
":test_util",
"//tensorflow/lite:type_to_tflitetype",
"//tensorflow/lite/core/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/kernels:test_util",
"//tensorflow/lite/kernels/internal:compatibility",
"//tensorflow/lite/schema:schema_fbs",
"@com_google_googletest//:gtest_main",
],
)
cc_test(
name = "list_from_tensor_test",
srcs = ["list_kernels/list_from_tensor_test.cc"],
deps = [
":list_ops_lib",
":tensor_array",
"//tensorflow/lite/c:common",
"//tensorflow/lite/core/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/kernels:test_util",
"//tensorflow/lite/kernels/internal:compatibility",
"//tensorflow/lite/kernels/internal:tensor_ctypes",
"//tensorflow/lite/schema:schema_fbs",
"@com_google_googletest//:gtest_main",
],
)
cc_test(
name = "list_get_item_test",
srcs = ["list_kernels/list_get_item_test.cc"],
deps = [
":list_ops_lib",
":test_util",
"//tensorflow/lite/core/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/kernels:test_util",
"//tensorflow/lite/kernels/internal:tensor_ctypes",
"//tensorflow/lite/schema:schema_fbs",
"@com_google_googletest//:gtest_main",
],
)
cc_test(
name = "list_length_test",
srcs = ["list_kernels/list_length_test.cc"],
deps = [
":list_ops_lib",
":tensor_array",
":test_util",
"//tensorflow/lite/core/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/kernels:test_util",
"//tensorflow/lite/schema:schema_fbs",
"@com_google_googletest//:gtest_main",
],
)
cc_test(
name = "list_element_shape_test",
srcs = ["list_kernels/list_element_shape_test.cc"],
deps = [
":list_ops_lib",
":tensor_array",
":test_util",
"//tensorflow/lite/core/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/kernels:test_util",
"//tensorflow/lite/schema:schema_fbs",
"@com_google_googletest//:gtest_main",
],
)
cc_test(
name = "list_pop_back_test",
srcs = ["list_kernels/list_pop_back_test.cc"],
deps = [
":list_ops_lib",
":tensor_array",
":test_util",
"//tensorflow/lite/core/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/kernels:test_util",
"//tensorflow/lite/kernels/internal:tensor_ctypes",
"//tensorflow/lite/schema:schema_fbs",
"@com_google_googletest//:gtest_main",
],
)
cc_test(
name = "list_push_back_test",
srcs = ["list_kernels/list_push_back_test.cc"],
deps = [
":list_ops_lib",
":tensor_array",
":test_util",
"//tensorflow/lite:type_to_tflitetype",
"//tensorflow/lite/core/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/kernels:kernel_util",
"//tensorflow/lite/kernels:test_util",
"//tensorflow/lite/kernels/internal:compatibility",
"//tensorflow/lite/schema:schema_fbs",
"@com_google_googletest//:gtest_main",
],
)
cc_test(
name = "variant_add_n_test",
srcs = ["list_kernels/variant_add_n_test.cc"],
deps = [
":list_ops_lib",
":tensor_array",
":test_util",
"//tensorflow/lite:type_to_tflitetype",
"//tensorflow/lite/core/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/kernels:kernel_util",
"//tensorflow/lite/kernels:test_util",
"//tensorflow/lite/kernels/internal:compatibility",
"//tensorflow/lite/schema:schema_fbs",
"@com_google_googletest//:gtest_main",
],
)
cc_test(
name = "variant_zeros_like_test",
srcs = ["list_kernels/variant_zeros_like_test.cc"],
deps = [
":list_ops_lib",
":tensor_array",
":test_util",
"//tensorflow/lite/core/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/kernels:kernel_util",
"//tensorflow/lite/kernels:test_util",
"//tensorflow/lite/kernels/internal:compatibility",
"//tensorflow/lite/schema:schema_fbs",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "list_ops_util",
srcs = ["list_ops_util.cc"],
hdrs = ["list_ops_util.h"],
visibility = ["//tensorflow/lite/kernels/variants:__subpackages__"],
deps = [
":tensor_array",
"//tensorflow/lite:array",
"//tensorflow/lite:util",
"//tensorflow/lite/c:c_api_types",
"//tensorflow/lite/core/c:common",
],
)
cc_test(
name = "list_ops_util_test",
srcs = ["list_ops_util_test.cc"],
deps = [
":list_ops_util",
":tensor_array",
"//tensorflow/lite:array",
"//tensorflow/lite:util",
"//tensorflow/lite/core/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/kernels:test_util",
"@com_google_googletest//:gtest_main",
],
)
# Utils for testing complex multi-op subgraphs.
cc_library(
name = "list_ops_subgraph_test_util",
testonly = True,
srcs = ["list_ops_subgraph_test_util.cc"],
hdrs = ["list_ops_subgraph_test_util.h"],
deps = [
":list_ops_lib",
"//tensorflow/lite:builtin_ops",
"//tensorflow/lite/c:c_api_types",
"//tensorflow/lite/c:common",
"//tensorflow/lite/core:cc_api_stable",
"//tensorflow/lite/core:subgraph",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/kernels:builtin_ops",
"//tensorflow/lite/kernels:op_macros",
"//tensorflow/lite/kernels:subgraph_test_util",
"//tensorflow/lite/schema:schema_fbs",
"@com_google_absl//absl/types:span",
"@com_google_googletest//:gtest",
],
)
cc_test(
name = "list_ops_subgraph_test",
srcs = ["list_ops_subgraph_test.cc"],
deps = [
":list_ops_subgraph_test_util",
":tensor_array",
"//tensorflow/lite:array",
"//tensorflow/lite:interpreter_test_util",
"//tensorflow/lite:util",
"//tensorflow/lite/c:c_api_types",
"//tensorflow/lite/c:common",
"//tensorflow/lite/kernels:kernel_util",
"//tensorflow/lite/kernels:op_macros",
"//tensorflow/lite/kernels:test_util",
"//tensorflow/lite/kernels/internal:compatibility",
"//tensorflow/lite/schema:schema_fbs",
"@com_google_absl//absl/types:span",
"@com_google_googletest//:gtest_main",
],
)
# Type erased class that sits behind the `data.data` of variant tensors.
cc_library(
name = "tensor_array",
srcs = ["tensor_array.cc"],
hdrs = ["tensor_array.h"],
compatible_with = get_compatible_with_portable(),
visibility = ["//visibility:private"],
deps = [
"//tensorflow/lite:array",
"//tensorflow/lite:util",
"//tensorflow/lite/c:common",
"//tensorflow/lite/core/c:common",
],
)
cc_test(
name = "tensor_array_test",
srcs = ["tensor_array_test.cc"],
deps = [
":tensor_array",
"//tensorflow/lite:array",
"//tensorflow/lite:type_to_tflitetype",
"//tensorflow/lite:util",
"//tensorflow/lite/c:c_api_types",
"//tensorflow/lite/c:common",
"//tensorflow/lite/core/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/kernels:test_util",
"@com_google_googletest//:gtest_main",
],
)
# Benchmark tool override with registered custom tensorlist ops.
cc_library(
name = "register_list_ops_for_benchmark",
srcs = ["register_list_ops.cc"],
visibility = ["//visibility:public"],
deps = [
":list_ops_lib",
"//tensorflow/lite:mutable_op_resolver",
"//tensorflow/lite:op_resolver",
"//tensorflow/lite/tools/benchmark:register_custom_op",
],
alwayslink = 1,
)
cc_binary(
name = "benchmark_model_with_list_ops",
copts = common_copts,
linkopts = tflite_linkopts() + select({
"//tensorflow:android": [
"-pie", # Android 5.0 and later supports only PIE
"-lm", # some builtin ops, e.g., tanh, need -lm
"-Wl,--rpath=/data/local/tmp/", # Hexagon delegate libraries should be in /data/local/tmp
],
"//conditions:default": [],
}),
tags = ["builder_default_android_arm64"],
deps = [
":register_list_ops_for_benchmark",
"//tensorflow/lite/tools/benchmark:benchmark_model_main",
],
)
tensorflow/lite/kernels/variants/README.md
+31
View File
@@ -0,0 +1,31 @@
# TfLite Variants
Variant tensors in TensorFlow wrap and store arbitrary C++ objects within their
data members. Common usage regards non-trivial and potentially referential
buffer semantics (TensorLists and DataSets being cannonical examples).
This directory contains implementations for these containers
and kernels for the tflite runtime.
Currently tflite supplies performant custom kernels for a subset of
`tf.list_ops`.
Also refer to [variant tensor in the tflite common api](https://github.com/tensorflow/tensorflow/blob/61c76427561a46d03605370fc685d810c1c3e717/tensorflow/lite/core/c/common.h#L1320C9-L1320C9)
, [tensor list legalization](https://github.com/tensorflow/tensorflow/blob/d36fb81bc1ef258d5024b791d61cdd5136ca09af/tensorflow/compiler/mlir/lite/transforms/legalize_tensorlist.cc)
, and [end to end python tests](https://github.com/tensorflow/tensorflow/blob/d36fb81bc1ef258d5024b791d61cdd5136ca09af/tensorflow/lite/kernels/variants/py/end_to_end_test.py)
for example usage.
**api**
* `./py/register_list_ops_py.py` : Bindings for registering ops in python.
* `./list_ops_lib` : Include for tensorlist kernel registrations.
* `./register_list_ops` : Register all kernels with op resolver in C++.
**implementations**
* `./list_kernels/` : Custom `TensorList*` kernels for tflite.
* `./tensor_array` : A variable length array of reference counted `TfLiteTensor`
.
**tests**
* `/list_ops_subgraph_test` : Multi-Op tests through C++ api.
* `/py/end_to_end_test.py` : Tests through python api and compare to `tf.list_ops`.
tensorflow/lite/kernels/variants/list_kernels/list_element_shape.cc
+111
View File
@@ -0,0 +1,111 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstdint>
#include <cstring>
#include "tensorflow/lite/array.h"
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
namespace tflite {
namespace variants {
namespace ops {
namespace list_element_shape {
namespace {
// This kernel returns a `TfLiteTensor` which represents the shape signature
// of the tensorlist elements stored in the `TensorArray::ElementShape` field as
// a `TfLiteIntArray`. Encoding shape signatures in tensors works as follows:
//
// An unranked shape signature:
// `TfLiteIntArray` : []
// `TfLiteTensor` : shape = [], data = [-1]
//
// A scalar shape signature:
// `TfLiteIntArray` : [0]
// `TfLiteTensor` : shape = [0], data = []
//
// A ranked tensor shape signature (with possibly dynamic dimensions):
// `TfLiteIntArray` : [Dim1, Dim2 ... DimRank] (DimI >= -1)
// `TfLiteTensor` : shape = [Rank], data = [Dim1, Dim2 ... DimRank]
using ::tflite::variants::TensorArray;
constexpr int kListInput = 0;
constexpr int kShapeOut = 0;
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
const TfLiteTensor* list_input;
TF_LITE_ENSURE_OK(context,
GetInputSafe(context, node, kListInput, &list_input));
TF_LITE_ENSURE(context, list_input->type == kTfLiteVariant);
TfLiteTensor* shape_out;
TF_LITE_ENSURE_OK(context,
GetOutputSafe(context, node, kShapeOut, &shape_out));
TF_LITE_ENSURE_TYPES_EQ(context, shape_out->type, kTfLiteInt32);
SetTensorToDynamic(shape_out);
return kTfLiteOk;
}
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
const TfLiteTensor* list_input;
TF_LITE_ENSURE_OK(context,
GetInputSafe(context, node, kListInput, &list_input));
const TensorArray* const list =
reinterpret_cast<const TensorArray*>(list_input->data.data);
const TfLiteIntArray& element_shape = *list->ElementShape();
TfLiteTensor* shape_out;
TF_LITE_ENSURE_OK(context,
GetOutputSafe(context, node, kShapeOut, &shape_out));
if (element_shape.size == 0) {
// Unranked
context->ResizeTensor(context, shape_out, BuildTfLiteArray(0).release());
GetTensorData<int32_t>(shape_out)[0] = -1;
} else if (element_shape.data[0] == 0) {
// Scalar
context->ResizeTensor(context, shape_out, BuildTfLiteArray({0}).release());
} else {
// Ranked
context->ResizeTensor(context, shape_out,
BuildTfLiteArray({element_shape.size}).release());
memcpy(GetTensorData<int32_t>(shape_out), element_shape.data,
element_shape.size * sizeof(int32_t));
}
return kTfLiteOk;
}
} // namespace
} // namespace list_element_shape
TfLiteRegistration* Register_LIST_ELEMENT_SHAPE() {
static TfLiteRegistration r = {nullptr, nullptr, list_element_shape::Prepare,
list_element_shape::Eval};
return &r;
}
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/list_element_shape_test.cc
+98
View File
@@ -0,0 +1,98 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/schema/schema_generated.h"
namespace tflite {
namespace variants {
namespace ops {
namespace {
using ::testing::ElementsAreArray;
class ListElementShapeModel : public ListOpModel {
public:
ListElementShapeModel() {
list_input_ = AddInput({TensorType_VARIANT, {}});
shape_output_ = AddOutput({TensorType_INT32, {}});
SetCustomOp("ListElementShape", {}, Register_LIST_ELEMENT_SHAPE);
BuildInterpreter({{}});
}
const TfLiteTensor* GetOutputTensor(int index) {
return interpreter_->tensor(index);
}
int list_input_;
int shape_output_;
};
TEST(ListElementShapeTest, MultiDimStaticShape) {
ListElementShapeModel m;
m.PopulateListTensor(0, {2, 2}, 10, kTfLiteInt32);
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* const out = m.GetOutputTensor(m.shape_output_);
ASSERT_THAT(out, DimsAre({2}));
ASSERT_THAT(std::vector<int>(out->data.i32, out->data.i32 + 2),
ElementsAreArray({2, 2}));
}
TEST(ListElementShapeTest, MultiDimWithDynamicDims) {
ListElementShapeModel m;
m.PopulateListTensor(0, {2, -1, 3}, 10, kTfLiteInt32);
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* const out = m.GetOutputTensor(m.shape_output_);
ASSERT_THAT(out, DimsAre({3}));
ASSERT_THAT(std::vector<int>(out->data.i32, out->data.i32 + 3),
ElementsAreArray({2, -1, 3}));
}
TEST(ListElementShapeTest, ScalarShape) {
ListElementShapeModel m;
m.PopulateListTensor(0, {0}, 10, kTfLiteInt32);
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* const out = m.GetOutputTensor(m.shape_output_);
ASSERT_THAT(out, DimsAre({0}));
ASSERT_EQ(out->bytes, 0);
}
TEST(ListElementShapeTest, UnrankedShape) {
ListElementShapeModel m;
m.PopulateListTensor(0, {}, 10, kTfLiteInt32);
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* const out = m.GetOutputTensor(m.shape_output_);
ASSERT_THAT(out, DimsAre({}));
ASSERT_EQ(out->bytes, sizeof(int));
ASSERT_EQ(out->data.i32[0], -1);
}
} // namespace
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/list_from_tensor.cc
+144
View File
@@ -0,0 +1,144 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstring>
#include <utility>
#include "tensorflow/lite/array.h"
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/kernels/variants/list_ops_util.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/util.h"
namespace tflite {
namespace variants {
namespace ops {
namespace {
constexpr int kTensorInput = 0;
constexpr int kElementShapeInput = 1;
constexpr int kListOut = 0;
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
const TfLiteTensor* element_shape;
TF_LITE_ENSURE_OK(
context, GetInputSafe(context, node, kElementShapeInput, &element_shape));
TF_LITE_ENSURE_TYPES_EQ(context, element_shape->type, kTfLiteInt32);
TfLiteTensor* output;
TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, kListOut, &output));
TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteVariant);
output->allocation_type = kTfLiteVariantObject;
return kTfLiteOk;
}
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
const TfLiteTensor* tensor_input;
TF_LITE_ENSURE_OK(context,
GetInputSafe(context, node, kTensorInput, &tensor_input));
const int rank = tensor_input->dims->size;
// As in Tensorflow, input is not permitted be a scalar.
TF_LITE_ENSURE(context, rank > 0);
// Output list has `num_elements` equal to the first dim of `tensor_input`,
// and element tensors with shape equal to `Shape(tensor_input)[1:]`.
const int list_len = tensor_input->dims->data[0];
IntArrayUniquePtr element_shape_for_tensors =
BuildTfLiteArray(rank - 1, tensor_input->dims->data + 1);
TF_LITE_ENSURE(context, element_shape_for_tensors != nullptr);
// `element_shape_tensor` is an auxiliary input shape signature which
// is to be used as the `ElementShape()` attribute of the resulting
// `TensorArray`.
const TfLiteTensor* element_shape_tensor;
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kElementShapeInput,
&element_shape_tensor));
TF_LITE_ENSURE(context, (element_shape_tensor->dims->size == 1 &&
element_shape_tensor->dims->data[0] == rank - 1) ||
element_shape_tensor->dims->size == 0);
IntArrayUniquePtr element_shape_for_list;
TF_LITE_ENSURE_OK(context, TensorAsShape(context, *element_shape_tensor,
element_shape_for_list));
TF_LITE_ENSURE(context, element_shape_for_list != nullptr);
// Check given element shape is compatible with the suffix of input tensor's
// shape. TODO(b/257472333) consider wrapping this in `#ifndef NDEBUG`.
if (element_shape_for_list->size > 0) {
TF_LITE_ENSURE_EQ(context, element_shape_for_list->size,
element_shape_for_tensors->size);
for (int i = 0; i < element_shape_for_tensors->size; ++i) {
const int lhs = element_shape_for_list->data[i];
const int rhs = element_shape_for_tensors->data[i];
TF_LITE_ENSURE(context, lhs == -1 || rhs == -1 || lhs == rhs);
}
}
TfLiteTensor* output;
TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, kListOut, &output));
// Build and retrieve output list.
IntArrayUniquePtr copied_element_shape =
BuildTfLiteArray(*element_shape_for_list);
TF_LITE_ENSURE(context, copied_element_shape != nullptr);
TF_LITE_ENSURE_OK(context, TfLiteTensorVariantRealloc<TensorArray>(
output, tensor_input->type,
std::move(copied_element_shape)));
TensorArray* arr =
static_cast<TensorArray*>(static_cast<VariantData*>(output->data.data));
TF_LITE_ENSURE(context, arr->Resize(list_len));
// Copy each row of input into the elements of the new list.
size_t data_offset = 0;
for (int i = 0; i < list_len; ++i) {
IntArrayUniquePtr cur_shape = BuildTfLiteArray(*element_shape_for_tensors);
TF_LITE_ENSURE(context, cur_shape != nullptr);
TensorUniquePtr tensor_to_set = BuildTfLiteTensor(
tensor_input->type, std::move(cur_shape), kTfLiteDynamic);
TF_LITE_ENSURE(context, tensor_to_set != nullptr);
if (tensor_to_set->bytes > 0) {
TF_LITE_ENSURE(context, tensor_to_set->data.raw != nullptr);
TF_LITE_ENSURE(context, tensor_input->data.raw != nullptr);
TF_LITE_ENSURE(context,
tensor_to_set->bytes <= tensor_input->bytes - data_offset);
memcpy(tensor_to_set->data.raw, tensor_input->data.raw + data_offset,
tensor_to_set->bytes);
}
data_offset += tensor_to_set->bytes;
TF_LITE_ENSURE(context, arr->Set(i, std::move(tensor_to_set)));
}
return kTfLiteOk;
}
} // namespace
TfLiteRegistration* Register_LIST_FROM_TENSOR() {
static TfLiteRegistration r = {nullptr, nullptr, Prepare, Eval};
return &r;
}
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/list_from_tensor_test.cc
+169
View File
@@ -0,0 +1,169 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <tuple>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/kernels/internal/compatibility.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/schema/schema_generated.h"
using ::testing::ElementsAre;
namespace tflite {
namespace variants {
namespace ops {
namespace {
class ListFromTensorModel : public SingleOpModel {
public:
ListFromTensorModel(TensorData tensor_data, TensorData shape_data) {
tensor_id_ = AddInput(tensor_data);
shape_id_ = AddInput(shape_data);
list_id_ = AddOutput({TensorType_VARIANT, {1}});
SetCustomOp("TensorListFromTensor", /*custom_option=*/{},
Register_LIST_FROM_TENSOR);
BuildInterpreter({tensor_data.shape, shape_data.shape});
}
const TensorArray* GetOutputTensorArray(int tensor_id) {
TfLiteTensor* tensor = interpreter_->tensor(tensor_id);
TFLITE_CHECK(tensor != nullptr && tensor->type == kTfLiteVariant &&
tensor->allocation_type == kTfLiteVariantObject);
return static_cast<const TensorArray*>(
static_cast<const VariantData*>(tensor->data.data));
}
int tensor_id_;
int shape_id_;
int list_id_;
};
TEST(ListFromTensorTest, MatrixInput_ReturnsListWithVectorElements) {
ListFromTensorModel m({TensorType_INT32, {2, 2}}, {TensorType_INT32, {1}});
m.PopulateTensor<int>(m.tensor_id_, {1, 2, 3, 4});
m.PopulateTensor<int>(m.shape_id_, {2});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* arr = m.GetOutputTensorArray(m.list_id_);
ASSERT_EQ(arr->NumElements(), 2);
ASSERT_THAT(arr->ElementShape(), DimsAre({2}));
ASSERT_EQ(arr->ElementType(), kTfLiteInt32);
{
const TfLiteTensor* element = arr->At(0);
ASSERT_THAT(element, DimsAre({2}));
EXPECT_THAT(std::make_tuple(GetTensorData<int>(element), 2),
ElementsAre(1, 2));
}
{
const TfLiteTensor* element = arr->At(1);
ASSERT_THAT(element, DimsAre({2}));
EXPECT_THAT(std::make_tuple(GetTensorData<int>(element), 2),
ElementsAre(3, 4));
}
}
TEST(ListFromTensorTest, VectorInput_ReturnsListWithScalarElements) {
ListFromTensorModel m({TensorType_INT32, {2}}, {TensorType_INT32, {0}});
m.PopulateTensor<int>(m.tensor_id_, {1, 2});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* arr = m.GetOutputTensorArray(m.list_id_);
ASSERT_EQ(arr->NumElements(), 2);
ASSERT_THAT(arr->ElementShape(), DimsAre({}));
ASSERT_EQ(arr->ElementType(), kTfLiteInt32);
{
const TfLiteTensor* element = arr->At(0);
ASSERT_THAT(element, DimsAre({}));
EXPECT_THAT(std::make_tuple(GetTensorData<int>(element), 1),
ElementsAre(1));
}
{
const TfLiteTensor* element = arr->At(1);
ASSERT_THAT(element, DimsAre({}));
EXPECT_THAT(std::make_tuple(GetTensorData<int>(element), 1),
ElementsAre(2));
}
}
TEST(ListFromTensorTest, 3DInput_ReturnsListWithMatrixElements) {
ListFromTensorModel m({TensorType_INT32, {2, 2, 2}}, {TensorType_INT32, {2}});
m.PopulateTensor<int>(m.tensor_id_, {1, 2, 3, 4, 5, 6, 7, 8});
m.PopulateTensor<int>(m.shape_id_, {2, 2});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* arr = m.GetOutputTensorArray(m.list_id_);
ASSERT_EQ(arr->NumElements(), 2);
ASSERT_THAT(arr->ElementShape(), DimsAre({2, 2}));
ASSERT_EQ(arr->ElementType(), kTfLiteInt32);
{
const TfLiteTensor* element = arr->At(0);
ASSERT_THAT(element, DimsAre({2, 2}));
EXPECT_THAT(std::make_tuple(GetTensorData<int>(element), 4),
ElementsAre(1, 2, 3, 4));
}
{
const TfLiteTensor* element = arr->At(1);
ASSERT_THAT(element, DimsAre({2, 2}));
EXPECT_THAT(std::make_tuple(GetTensorData<int>(element), 4),
ElementsAre(5, 6, 7, 8));
}
}
TEST(ListFromTensorTest, MismatchedShapeInputTensorShape_Fails) {
ListFromTensorModel m({TensorType_INT32, {2, 2, 2}}, {TensorType_INT32, {2}});
m.PopulateTensor<int>(m.shape_id_, {2, 3});
ASSERT_EQ(m.Invoke(), kTfLiteError);
}
TEST(ListFromTensorTest, ScalarInput_Fails) {
ListFromTensorModel m({TensorType_INT32, {}}, {TensorType_INT32, {}});
ASSERT_EQ(m.Invoke(), kTfLiteError);
}
} // namespace
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/list_get_item.cc
+252
View File
@@ -0,0 +1,252 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstddef>
#include <cstring>
#include <utility>
#include "tensorflow/lite/array.h"
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/kernels/variants/list_ops_util.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/util.h"
namespace tflite {
namespace variants {
namespace ops {
namespace {
constexpr int kListInput = 0;
class GetItemSemantic {
public:
GetItemSemantic(TfLiteContext* ctx, TfLiteNode* node)
: ctx_(ctx), node_(node) {}
static constexpr int kElementShapeInputIdx = 2;
static constexpr int kTensorOutputIdx = 0;
static constexpr int kIndexInputIdx = 1;
[[nodiscard]] TfLiteStatus CheckAndHandleTensors() const {
TF_LITE_ENSURE(ctx_, NumInputs(node_) == 3 && NumOutputs(node_) == 1);
const TfLiteTensor* index_input;
TF_LITE_ENSURE_OK(ctx_,
GetInputSafe(ctx_, node_, kIndexInputIdx, &index_input));
TF_LITE_ENSURE_TYPES_EQ(ctx_, index_input->type, kTfLiteInt32);
return kTfLiteOk;
}
[[nodiscard]] TfLiteStatus GetIndexVal(const TensorArray* const arr,
int& result) const {
const TfLiteTensor* index_input;
TF_LITE_ENSURE_OK(ctx_,
GetInputSafe(ctx_, node_, kIndexInputIdx, &index_input));
TF_LITE_ENSURE_EQ(ctx_, index_input->bytes, sizeof(int));
TF_LITE_ENSURE(ctx_, index_input->data.raw != nullptr);
result = *GetTensorData<int>(index_input);
return kTfLiteOk;
}
[[nodiscard]] TfLiteStatus HandleOutput(const TensorArray* const arr) const {
return kTfLiteOk;
}
private:
TfLiteContext* const ctx_;
TfLiteNode* const node_;
};
class PopBackSemantic {
public:
PopBackSemantic(TfLiteContext* ctx, TfLiteNode* node)
: ctx_(ctx), node_(node) {}
static constexpr int kElementShapeInputIdx = 1;
static constexpr int kTensorOutputIdx = 1;
static constexpr int kListOutputIdx = 0;
[[nodiscard]] TfLiteStatus CheckAndHandleTensors() const {
TF_LITE_ENSURE(ctx_, NumInputs(node_) == 2 && NumOutputs(node_) == 2);
TfLiteTensor* list_output;
TF_LITE_ENSURE_OK(ctx_,
GetOutputSafe(ctx_, node_, kListOutputIdx, &list_output));
TF_LITE_ENSURE_TYPES_EQ(ctx_, list_output->type, kTfLiteVariant);
list_output->allocation_type = kTfLiteVariantObject;
return kTfLiteOk;
}
[[nodiscard]] TfLiteStatus GetIndexVal(const TensorArray* const arr,
int& result) const {
result = arr->NumElements() - 1;
return kTfLiteOk;
}
[[nodiscard]] TfLiteStatus HandleOutput(const TensorArray* const arr) const {
TfLiteTensor* list_output;
TF_LITE_ENSURE_OK(ctx_,
GetOutputSafe(ctx_, node_, kListOutputIdx, &list_output));
TensorArray* output_arr = static_cast<TensorArray*>(
arr->CloneTo(static_cast<VariantData*>(list_output->data.data)));
output_arr->Resize(output_arr->NumElements() - 1);
list_output->data.data = static_cast<VariantData*>(output_arr);
return kTfLiteOk;
}
private:
TfLiteContext* const ctx_;
TfLiteNode* const node_;
};
template <class Semantic>
TfLiteStatus Prepare(TfLiteContext* ctx, TfLiteNode* node) {
const auto semantic = Semantic(ctx, node);
TF_LITE_ENSURE_OK(ctx, semantic.CheckAndHandleTensors());
const TfLiteTensor* list_input;
TF_LITE_ENSURE_OK(ctx, GetInputSafe(ctx, node, kListInput, &list_input));
TF_LITE_ENSURE_TYPES_EQ(ctx, list_input->type, kTfLiteVariant);
const TfLiteTensor* element_shape_input;
TF_LITE_ENSURE_OK(ctx, GetInputSafe(ctx, node, semantic.kElementShapeInputIdx,
&element_shape_input));
TF_LITE_ENSURE_TYPES_EQ(ctx, element_shape_input->type, kTfLiteInt32);
TfLiteTensor* output;
TF_LITE_ENSURE_OK(
ctx, GetOutputSafe(ctx, node, semantic.kTensorOutputIdx, &output));
const TfLiteIntArray* const out_dims_sig = output->dims_signature;
// Tensorflow python allows the passing of a `tf.TensorShape` at list
// initialization and as an argument to this function as a hint. The compiler
// will also give this information to `output`. Currently tflite has no
// way to distinguish between `output->dims_signature` encoding a true scalar
// shape, or an "unranked" shape signature which is meant to be compatible
// with any shape (`tf.TensorShape(None)`). Because of this we need to fall
// back to dynamic if `output->dims_signature` looks like a scalar.
// Update this function after TODO(b/290271484) has been addressed.
if (out_dims_sig == nullptr || out_dims_sig->size <= 0 ||
!IsShapeFullyDefined(*out_dims_sig)) {
SetTensorToDynamic(output);
}
return kTfLiteOk;
}
template <class Semantic>
TfLiteStatus Eval(TfLiteContext* ctx, TfLiteNode* node) {
const auto semantic = Semantic(ctx, node);
const TfLiteTensor* list_input;
TF_LITE_ENSURE_OK(ctx, GetInputSafe(ctx, node, kListInput, &list_input));
TF_LITE_ENSURE_EQ(ctx, list_input->allocation_type, kTfLiteVariantObject);
TF_LITE_ENSURE(ctx, list_input->data.data != nullptr);
const auto* arr = static_cast<const TensorArray*>(
static_cast<VariantData*>(list_input->data.data));
int idx;
TF_LITE_ENSURE_OK(ctx, semantic.GetIndexVal(arr, idx));
TF_LITE_ENSURE(ctx, idx >= 0 && idx < arr->NumElements());
TfLiteTensor* output;
TF_LITE_ENSURE_OK(
ctx, GetOutputSafe(ctx, node, semantic.kTensorOutputIdx, &output));
TF_LITE_ENSURE_TYPES_EQ(ctx, arr->ElementType(), output->type);
const TfLiteTensor* const element = arr->At(idx);
if (element != nullptr) {
if (IsDynamicTensor(output)) {
size_t bytes;
TF_LITE_ENSURE_OK(ctx, BytesRequired(output->type, element->dims->data,
element->dims->size, &bytes, ctx));
TF_LITE_ENSURE_OK(ctx, TfLiteTensorResizeMaybeCopy(bytes, output, false));
}
TF_LITE_ENSURE_OK(ctx, TfLiteTensorCopy(element, output));
return semantic.HandleOutput(arr);
}
// As in tensorflow, it is possible to "get" an empty element in the list.
// In the dynamic case, try to infer the output shape first through input and
// list shape, and then through the elements in the list. Otherwise,
// the shape will be known at compile time so we can just use that.
if (!IsDynamicTensor(output)) {
if (output->bytes > 0) {
TF_LITE_ENSURE(ctx, output->data.data != nullptr);
memset(output->data.data, 0, output->bytes);
}
return semantic.HandleOutput(arr);
}
const TfLiteTensor* element_shape_input;
TF_LITE_ENSURE_OK(ctx, GetInputSafe(ctx, node, semantic.kElementShapeInputIdx,
&element_shape_input));
IntArrayUniquePtr element_shape;
TF_LITE_ENSURE_OK(ctx,
TensorAsShape(ctx, *element_shape_input, element_shape));
TF_LITE_ENSURE(ctx, element_shape != nullptr);
IntArrayUniquePtr output_shape = MergeShapesOrNull(
BuildTfLiteArray(*arr->ElementShape()), std::move(element_shape));
TF_LITE_ENSURE(ctx, output_shape != nullptr);
const bool can_infer_shape = (element_shape_input->dims->size != 0 ||
arr->ElementShape()->size != 0) &&
IsShapeFullyDefined(*output_shape);
if (!can_infer_shape) {
TF_LITE_ENSURE_MSG(
ctx,
GetShapeIfAllEqual(*arr, output_shape) == kTfLiteOk &&
output_shape != nullptr,
"Failed to infer the output shape for an item which has not been set.");
}
TF_LITE_ENSURE_OK(ctx,
ctx->ResizeTensor(ctx, output, output_shape.release()));
if (output->bytes > 0) {
TF_LITE_ENSURE(ctx, output->data.data != nullptr);
memset(output->data.data, 0, output->bytes);
}
return semantic.HandleOutput(arr);
}
} // namespace
TfLiteRegistration* Register_LIST_GET_ITEM() {
static TfLiteRegistration r = {nullptr, nullptr, Prepare<GetItemSemantic>,
Eval<GetItemSemantic>};
return &r;
}
TfLiteRegistration* Register_LIST_POP_BACK() {
static TfLiteRegistration r = {nullptr, nullptr, Prepare<PopBackSemantic>,
Eval<PopBackSemantic>};
return &r;
}
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/list_get_item_test.cc
+261
View File
@@ -0,0 +1,261 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <tuple>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/schema/schema_generated.h"
namespace tflite {
namespace variants {
namespace ops {
namespace {
using ::testing::ElementsAre;
// TODO add output param
class ListGetItemModel : public ListOpModel {
public:
ListGetItemModel(TensorData index, TensorData element_shape,
TensorData output) {
list_input_ = AddInput({TensorType_VARIANT, {}});
index_input_ = AddInput(index);
element_shape_input_ = AddInput(element_shape);
output_ = AddOutput(output);
SetCustomOp("TensorListGetItem", /*custom_option=*/{},
Register_LIST_GET_ITEM);
BuildInterpreter({{}, index.shape, element_shape.shape});
}
// Simplified constructor for creating valid models.
ListGetItemModel(int index, const std::vector<int>& element_shape) {
list_input_ = AddInput({TensorType_VARIANT, {}});
index_input_ = AddInput({TensorType_INT32, {1}});
element_shape_input_ =
AddInput({TensorType_INT32, {static_cast<int>(element_shape.size())}});
output_ = AddOutput({TensorType_INT32, element_shape});
SetCustomOp("TensorListGetItem", /*custom_option=*/{},
Register_LIST_GET_ITEM);
BuildInterpreter({{}, {1}, {static_cast<int>(element_shape.size())}});
PopulateListTensor(list_input_, {}, 2, kTfLiteInt32);
PopulateTensor(index_input_, {index});
PopulateTensor(element_shape_input_, element_shape);
}
const TfLiteTensor* GetOutput(int idx) { return interpreter_->tensor(idx); }
int list_input_;
int index_input_;
int element_shape_input_;
int output_;
};
TEST(ListGetItemTest, IndexOOB_Fails) {
ListGetItemModel m(-1, {2, 2});
ASSERT_EQ(m.Invoke(), kTfLiteError);
}
TEST(ListGetItemTest, GetPresentItem_ReturnsElement_Dynamic) {
ListGetItemModel m({TensorType_INT32, {}}, {TensorType_INT32, {2}},
{TensorType_INT32, {}});
m.PopulateListTensor(m.list_input_, {2, 2}, 3, kTfLiteInt32);
m.PopulateTensor(m.element_shape_input_, {2, 2});
m.PopulateTensor(m.index_input_, {0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutput(m.output_);
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
EXPECT_THAT(std::tuple(GetTensorData<int>(output), 4),
ElementsAre(0, 0, 0, 0));
}
TEST(ListGetItemTest, GetUnsetItem_InferShapeFromListShape_Dynamic) {
ListGetItemModel m({TensorType_INT32, {}}, {TensorType_INT32, {}},
{TensorType_INT32, {}});
m.PopulateListTensor(m.list_input_, {2, 2}, 2, kTfLiteInt32);
m.PopulateTensor(m.index_input_, {0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutput(m.output_);
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
EXPECT_THAT(std::tuple(GetTensorData<int>(output), 4),
ElementsAre(0, 0, 0, 0));
}
TEST(ListGetItemTest, GetUnsetItem_InferShapeFromGivenShape_Dynamic) {
ListGetItemModel m({TensorType_INT32, {}}, {TensorType_INT32, {2}},
{TensorType_INT32, {}});
m.PopulateListTensor(m.list_input_, {}, 2, kTfLiteInt32);
m.PopulateTensor(m.index_input_, {0});
m.PopulateTensor(m.element_shape_input_, {2, 2});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutput(m.output_);
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
EXPECT_THAT(std::tuple(GetTensorData<int>(output), 4),
ElementsAre(0, 0, 0, 0));
}
TEST(ListGetItemTest, GetUnsetItem_InferShapeFromOtherElements_Dynamic) {
ListGetItemModel m({TensorType_INT32, {}}, {TensorType_INT32, {}},
{TensorType_INT32, {}});
m.PopulateListTensor(m.list_input_, {}, 3, kTfLiteInt32);
m.ListSetItem(m.list_input_, 1, {2, 2}, kTfLiteInt32,
std::vector<int>{1, 2, 3, 4}.data());
m.ListSetItem(m.list_input_, 2, {2, 2}, kTfLiteInt32,
std::vector<int>{5, 6, 7, 8}.data());
m.PopulateTensor(m.index_input_, {0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutput(m.output_);
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
EXPECT_THAT(std::tuple(GetTensorData<int>(output), 4),
ElementsAre(0, 0, 0, 0));
}
TEST(ListGetItemTest,
GetUnsetItem_InferShapeFromMergedListShapeGivenShape_Dynamic) {
ListGetItemModel m({TensorType_INT32, {}}, {TensorType_INT32, {2}},
{TensorType_INT32, {}});
m.PopulateListTensor(m.list_input_, {2, -1}, 3, kTfLiteInt32);
m.PopulateTensor(m.element_shape_input_, {-1, 2});
m.PopulateTensor(m.index_input_, {0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutput(m.output_);
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
EXPECT_THAT(std::tuple(GetTensorData<int>(output), 4),
ElementsAre(0, 0, 0, 0));
}
TEST(ListGetItemTest, GetPresentItem_ReturnsElement_ScalarFallsBackDynamic) {
ListGetItemModel m({TensorType_INT32, {}}, {TensorType_INT32, {}},
{TensorType_INT32, {}});
m.PopulateListTensor(m.list_input_, {}, 3, kTfLiteInt32);
m.ListSetItem(m.list_input_, 1, {}, kTfLiteInt32, std::vector<int>{1}.data());
m.PopulateTensor(m.index_input_, {1});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutput(m.output_);
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_THAT(output, DimsAre({}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
EXPECT_THAT(std::tuple(GetTensorData<int>(output), 1), ElementsAre(1));
}
TEST(ListGetItemTest, GetPresentItem_ReturnsElement_Static) {
TensorData output_spec({TensorType_INT32, {2, 2}});
output_spec.shape_signature = {2, 2};
ListGetItemModel m({TensorType_INT32, {}}, {TensorType_INT32, {2}},
output_spec);
m.PopulateListTensor(m.list_input_, {2, 2}, 3, kTfLiteInt32);
m.ListSetItem(m.list_input_, 1, {2, 2}, kTfLiteInt32,
std::vector<int>{1, 2, 3, 4}.data());
m.PopulateTensor(m.element_shape_input_, {2, 2});
m.PopulateTensor(m.index_input_, {1});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutput(m.output_);
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_EQ(output->allocation_type, kTfLiteArenaRw);
EXPECT_THAT(std::tuple(GetTensorData<int>(output), 4),
ElementsAre(1, 2, 3, 4));
}
TEST(ListGetItemTest, GetPresentItem_OutputShapeMismatched_Fails_Static) {
TensorData output_spec({TensorType_INT32, {2, 2}});
output_spec.shape_signature = {2, 2};
ListGetItemModel m({TensorType_INT32, {}}, {TensorType_INT32, {}},
output_spec);
m.PopulateListTensor(m.list_input_, {}, 3, kTfLiteInt32);
m.ListSetItem(m.list_input_, 1, {3, 3}, kTfLiteInt32,
std::vector<int>{1, 2, 3, 4, 5, 6, 7, 8, 9}.data());
m.PopulateTensor(m.index_input_, {1});
ASSERT_EQ(m.Invoke(), kTfLiteError);
}
TEST(ListGetItemTest, GetUnsetItem_Static) {
TensorData output_spec({TensorType_INT32, {2, 2}});
output_spec.shape_signature = {2, 2};
ListGetItemModel m({TensorType_INT32, {}}, {TensorType_INT32, {2}},
output_spec);
m.PopulateListTensor(m.list_input_, {2, 2}, 3, kTfLiteInt32);
m.PopulateTensor(m.element_shape_input_, {2, 2});
m.PopulateTensor(m.index_input_, {0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutput(m.output_);
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_EQ(output->allocation_type, kTfLiteArenaRw);
EXPECT_THAT(std::tuple(GetTensorData<int>(output), 4),
ElementsAre(0, 0, 0, 0));
}
} // namespace
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/list_length.cc
+74
View File
@@ -0,0 +1,74 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
namespace tflite {
namespace variants {
namespace ops {
namespace list_length {
namespace {
using ::tflite::variants::TensorArray;
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
const TfLiteTensor* list_input;
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &list_input));
TF_LITE_ENSURE_TYPES_EQ(context, list_input->type, kTfLiteVariant);
TfLiteTensor* output;
TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteInt32);
TF_LITE_ENSURE_EQ(context, output->dims->size, 0);
output->allocation_type = kTfLiteArenaRw;
return kTfLiteOk;
}
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
const TfLiteTensor* list_input;
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &list_input));
TF_LITE_ENSURE_EQ(context, list_input->allocation_type, kTfLiteVariantObject);
const TensorArray* const input_arr =
reinterpret_cast<TensorArray*>(list_input->data.data);
TfLiteTensor* output;
TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
const int length = input_arr->NumElements();
output->data.i32[0] = length;
return kTfLiteOk;
}
} // namespace
} // namespace list_length
TfLiteRegistration* Register_LIST_LENGTH() {
static TfLiteRegistration r = {nullptr, nullptr, list_length::Prepare,
list_length::Eval};
return &r;
}
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/list_length_test.cc
+68
View File
@@ -0,0 +1,68 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/schema/schema_generated.h"
namespace tflite {
namespace variants {
namespace ops {
namespace {
class ListLengthModel : public ListOpModel {
public:
ListLengthModel() {
list_input_ = AddInput({TensorType_VARIANT, {}});
length_output_ = AddOutput({TensorType_INT32, {}});
SetCustomOp("ListLength", {}, Register_LIST_LENGTH);
BuildInterpreter({{}});
}
const TfLiteTensor* GetOutputTensor() {
return interpreter_->tensor(length_output_);
}
int list_input_;
int length_output_;
};
class ListLengthTest : public ::testing::TestWithParam<int> {};
TEST_P(ListLengthTest, OutputIsListLength) {
const int length = GetParam();
ListLengthModel m;
m.PopulateListTensor(m.list_input_, {2, 2}, length, kTfLiteInt32);
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* const output = m.GetOutputTensor();
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_EQ(output->allocation_type, kTfLiteArenaRw);
ASSERT_THAT(output, DimsAre({}));
ASSERT_EQ(output->data.i32[0], length);
}
INSTANTIATE_TEST_SUITE_P(ListLengthTests, ListLengthTest,
testing::Values(0, 1, 5, 10, 100));
} // namespace
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/list_pop_back_test.cc
+239
View File
@@ -0,0 +1,239 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <tuple>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/schema/schema_generated.h"
namespace tflite {
namespace variants {
namespace ops {
namespace {
using ::testing::ElementsAre;
class ListPopBackModel : public ListOpModel {
public:
ListPopBackModel(const std::vector<int>& list_element_shape, int num_elements,
const std::vector<int>& tensor_element_shape,
TensorData output = {TensorType_INT32, {}}) {
list_input_ = AddInput({TensorType_VARIANT, {}});
const std::vector<int> tensor_element_shape_shape =
tensor_element_shape.empty()
? std::vector<int>()
: std::vector<int>{static_cast<int>(tensor_element_shape.size())};
element_shape_input_ =
AddInput({TensorType_INT32, tensor_element_shape_shape});
list_output_ = AddOutput({TensorType_VARIANT, {}});
output_ = AddOutput(output);
SetCustomOp("ListPopBack", /*custom_option=*/{}, Register_LIST_POP_BACK);
BuildInterpreter({{}, tensor_element_shape_shape});
PopulateListTensor(list_input_, list_element_shape, num_elements,
kTfLiteInt32);
PopulateTensor(element_shape_input_, tensor_element_shape);
}
const TfLiteTensor* GetTensorOutput() {
return interpreter_->tensor(output_);
}
const TensorArray* GetTensorListOutput() {
return reinterpret_cast<const TensorArray*>(
interpreter_->tensor(list_output_)->data.data);
}
int list_input_;
int index_input_;
int element_shape_input_;
int output_;
int list_output_;
};
TEST(ListPopBackTest, ZeroLenListFails) {
ListPopBackModel m({2, 2}, 0, {});
ASSERT_EQ(m.Invoke(), kTfLiteError);
}
TEST(ListPopBackTest, GetUnsetItem_InferShapeFromListShape_Dynamic) {
ListPopBackModel m({2, 2}, 2, {});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* const output = m.GetTensorOutput();
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
EXPECT_THAT(std::tuple(GetTensorData<int>(output), 4),
ElementsAre(0, 0, 0, 0));
const TensorArray* const arr = m.GetTensorListOutput();
EXPECT_EQ(arr->NumElements(), 1);
}
TEST(ListPopBackTest, GetUnsetItem_InferShapeFromGivenShape_Dynamic) {
ListPopBackModel m({}, 2, {2, 2});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetTensorOutput();
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
EXPECT_THAT(std::tuple(GetTensorData<int>(output), 4),
ElementsAre(0, 0, 0, 0));
const TensorArray* const arr = m.GetTensorListOutput();
EXPECT_EQ(arr->NumElements(), 1);
}
TEST(ListPopBackTest, GetUnsetItem_InferShapeFromOtherElements_Dynamic) {
ListPopBackModel m({}, 3, {});
m.ListSetItem(m.list_input_, 0, {2, 2}, kTfLiteInt32,
std::vector<int>{1, 2, 3, 4}.data());
m.ListSetItem(m.list_input_, 1, {2, 2}, kTfLiteInt32,
std::vector<int>{5, 6, 7, 8}.data());
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetTensorOutput();
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
EXPECT_THAT(std::tuple(GetTensorData<int>(output), 4),
ElementsAre(0, 0, 0, 0));
const TensorArray* const arr = m.GetTensorListOutput();
EXPECT_EQ(arr->NumElements(), 2);
const TfLiteTensor* const item0 = arr->At(0);
ASSERT_NE(item0, nullptr);
ASSERT_THAT(item0, DimsAre({2, 2}));
ASSERT_THAT(std::tuple(GetTensorData<int>(item0), 4),
ElementsAre(1, 2, 3, 4));
const TfLiteTensor* const item1 = arr->At(1);
ASSERT_NE(item1, nullptr);
ASSERT_THAT(item1, DimsAre({2, 2}));
ASSERT_THAT(std::tuple(GetTensorData<int>(item1), 4),
ElementsAre(5, 6, 7, 8));
}
TEST(ListPopBackTest,
GetUnsetItem_InferShapeFromMergedListShapeGivenShape_Dynamic) {
ListPopBackModel m({2, -1}, 2, {-1, 2});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* const output = m.GetTensorOutput();
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
EXPECT_THAT(std::tuple(GetTensorData<int>(output), 4),
ElementsAre(0, 0, 0, 0));
const TensorArray* const arr = m.GetTensorListOutput();
EXPECT_EQ(arr->NumElements(), 1);
}
TEST(ListPopBackTest, GetPresentItem_ReturnsElement_ScalarFallsBackDynamic) {
ListPopBackModel m({}, 2, {});
m.ListSetItem(m.list_input_, 1, {}, kTfLiteInt32, std::vector<int>{1}.data());
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetTensorOutput();
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_THAT(output, DimsAre({}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
EXPECT_THAT(std::tuple(GetTensorData<int>(output), 1), ElementsAre(1));
const TensorArray* const arr = m.GetTensorListOutput();
EXPECT_EQ(arr->NumElements(), 1);
}
TEST(ListPopBackTest, GetPresentItem_ReturnsElement_Static) {
TensorData output_spec({TensorType_INT32, {2, 2}});
output_spec.shape_signature = {2, 2};
ListPopBackModel m({2, 2}, 2, {}, output_spec);
m.ListSetItem(m.list_input_, 1, {2, 2}, kTfLiteInt32,
std::vector<int>{1, 2, 3, 4}.data());
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetTensorOutput();
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_EQ(output->allocation_type, kTfLiteArenaRw);
EXPECT_THAT(std::tuple(GetTensorData<int>(output), 4),
ElementsAre(1, 2, 3, 4));
const TensorArray* const arr = m.GetTensorListOutput();
EXPECT_EQ(arr->NumElements(), 1);
}
TEST(ListPopBackTest, GetPresentItem_OutputShapeMismatched_Fails_Static) {
TensorData output_spec({TensorType_INT32, {2, 2}});
output_spec.shape_signature = {2, 2};
ListPopBackModel m({}, 2, {}, output_spec);
m.ListSetItem(m.list_input_, 1, {3, 3}, kTfLiteInt32,
std::vector<int>{1, 2, 3, 4, 5, 6, 7, 8, 9}.data());
ASSERT_EQ(m.Invoke(), kTfLiteError);
}
TEST(ListPopBackTest, GetUnsetItem_Static) {
TensorData output_spec({TensorType_INT32, {2, 2}});
output_spec.shape_signature = {2, 2};
ListPopBackModel m({}, 2, {}, output_spec);
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetTensorOutput();
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_EQ(output->allocation_type, kTfLiteArenaRw);
EXPECT_THAT(std::tuple(GetTensorData<int>(output), 4),
ElementsAre(0, 0, 0, 0));
const TensorArray* const arr = m.GetTensorListOutput();
EXPECT_EQ(arr->NumElements(), 1);
}
} // namespace
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/list_push_back_test.cc
+180
View File
@@ -0,0 +1,180 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstdint>
#include <optional>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/internal/compatibility.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/portable_type_to_tflitetype.h"
#include "tensorflow/lite/schema/schema_generated.h"
namespace tflite {
namespace variants {
namespace ops {
namespace {
using ::testing::AllOf;
template <typename T>
class PushBackWithTypeTest : public ::testing::Test {};
class ListPushBackModel : public ListOpModel {
public:
explicit ListPushBackModel(TensorData item_data) {
list_input_ = AddInput({TensorType_VARIANT, {}});
tensor_input_ = AddInput(item_data);
list_output_ = AddOutput({TensorType_VARIANT, {}});
SetCustomOp("ListPushBack", {}, Register_LIST_PUSH_BACK);
BuildInterpreter({{}, item_data.shape});
interpreter_->input_tensor(0)->allocation_type = kTfLiteVariantObject;
}
const TensorArray* GetOutputTensorArray() {
TfLiteTensor* tensor = interpreter_->tensor(list_output_);
TFLITE_CHECK(tensor != nullptr && tensor->type == kTfLiteVariant &&
tensor->allocation_type == kTfLiteVariantObject);
return static_cast<const TensorArray*>(
static_cast<const VariantData*>(tensor->data.data));
}
int list_input_;
int tensor_input_;
int list_output_;
};
constexpr int kNumElements = 4;
TYPED_TEST_SUITE_P(PushBackWithTypeTest);
TYPED_TEST_P(PushBackWithTypeTest, PushBackOnEmptyTensorList_ListShapeDefined) {
TfLiteType tfl_type = typeToTfLiteType<TypeParam>();
std::optional<TensorType> tensor_type = TflToTensorType(tfl_type);
ASSERT_TRUE(tensor_type.has_value());
ListPushBackModel m({tensor_type.value(), {2, 2}});
m.PopulateListTensor(m.list_input_, {2, 2}, kNumElements, tfl_type);
m.PopulateTensor<TypeParam>(m.tensor_input_, {0, 0, 0, 0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* const arr = m.GetOutputTensorArray();
ASSERT_EQ(arr->NumElements(), kNumElements + 1);
ASSERT_EQ(arr->ElementType(), tfl_type);
for (int i = 0; i < arr->NumElements() - 1; ++i) {
EXPECT_EQ(arr->At(i), nullptr);
}
// Sometimes sanitizers will fail if this is inlined within macro.
const int last_ind = arr->NumElements() - 1;
EXPECT_THAT(arr->At(last_ind), AllOf(IsAllocatedAs(tfl_type), DimsAre({2, 2}),
FilledWith<TypeParam>(0)));
}
TYPED_TEST_P(PushBackWithTypeTest,
PushBackOnEmptyTensorList_ListShapeUnranked) {
TfLiteType tfl_type = typeToTfLiteType<TypeParam>();
std::optional<TensorType> tensor_type = TflToTensorType(tfl_type);
ASSERT_TRUE(tensor_type.has_value());
ListPushBackModel m({tensor_type.value(), {2, 2}});
m.PopulateListTensor(m.list_input_, {}, kNumElements, tfl_type);
m.PopulateTensor<TypeParam>(m.tensor_input_, {0, 0, 0, 0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* const arr = m.GetOutputTensorArray();
ASSERT_EQ(arr->NumElements(), kNumElements + 1);
ASSERT_EQ(arr->ElementType(), tfl_type);
for (int i = 0; i < arr->NumElements() - 1; ++i) {
EXPECT_EQ(arr->At(i), nullptr);
}
// Sometimes sanitizers will fail if this is inlined within macro.
const int last_ind = arr->NumElements() - 1;
EXPECT_THAT(arr->At(last_ind), AllOf(IsAllocatedAs(tfl_type), DimsAre({2, 2}),
FilledWith<TypeParam>(0)));
}
TYPED_TEST_P(PushBackWithTypeTest,
PushBackOnNonEmptyListAtEmptyIndex_ItemsSameShape) {
TfLiteType tfl_type = typeToTfLiteType<TypeParam>();
std::optional<TensorType> tensor_type = TflToTensorType(tfl_type);
ASSERT_TRUE(tensor_type.has_value());
ListPushBackModel m({tensor_type.value(), {2, 2}});
m.PopulateListTensor(m.list_input_, {}, kNumElements, tfl_type);
TypeParam init_item_data[4] = {1, 1, 1, 1};
m.ListSetItem(m.list_input_, 0, {2, 2}, tfl_type, init_item_data);
m.PopulateTensor<TypeParam>(m.tensor_input_, {0, 0, 0, 0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* const arr = m.GetOutputTensorArray();
ASSERT_EQ(arr->NumElements(), kNumElements + 1);
ASSERT_EQ(arr->ElementType(), tfl_type);
for (int i = 1; i < arr->NumElements() - 1; ++i) {
EXPECT_EQ(arr->At(i), nullptr);
}
EXPECT_THAT(arr->At(0), AllOf(IsAllocatedAs(tfl_type), DimsAre({2, 2}),
FilledWith<TypeParam>(1)));
// Sometimes sanitizers will fail if this is inlined within macro.
const int last_ind = arr->NumElements() - 1;
EXPECT_THAT(arr->At(last_ind), AllOf(IsAllocatedAs(tfl_type), DimsAre({2, 2}),
FilledWith<TypeParam>(0)));
}
REGISTER_TYPED_TEST_SUITE_P(PushBackWithTypeTest,
PushBackOnEmptyTensorList_ListShapeDefined,
PushBackOnEmptyTensorList_ListShapeUnranked,
PushBackOnNonEmptyListAtEmptyIndex_ItemsSameShape);
using ValidTypes = ::testing::Types<int, int64_t, bool, float>;
INSTANTIATE_TYPED_TEST_SUITE_P(PushBackTests, PushBackWithTypeTest, ValidTypes);
TEST(ListPushBackTest, ItemNotSameTypeAsList_Fails) {
ListPushBackModel m{{TensorType_INT32, {2, 2}}};
m.PopulateListTensor(m.list_input_, {2, 2}, kNumElements, kTfLiteInt64);
ASSERT_EQ(m.Invoke(), kTfLiteError);
}
} // namespace
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/list_reserve.cc
+241
View File
@@ -0,0 +1,241 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstring>
#include <utility>
#include "tensorflow/lite/array.h"
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/kernels/variants/list_ops_util.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/util.h"
namespace tflite {
namespace variants {
namespace ops {
namespace list_reserve {
namespace {
using ::tflite::variants::TensorArray;
using ::tflite::variants::detail::ListReserveOptions;
TfLiteType ConvertTensorType(TensorType src) {
switch (src) {
case TensorType_INT32:
return kTfLiteInt32;
case TensorType_FLOAT32:
return kTfLiteFloat32;
case TensorType_BOOL:
return kTfLiteBool;
case TensorType_INT64:
return kTfLiteInt64;
default:
return kTfLiteNoType;
}
}
constexpr int kListOut = 0;
struct SemanticOutType {
TfLiteType element_type;
IntArrayUniquePtr element_shape;
int num_elements;
};
class ReserveSemantic {
public:
ReserveSemantic(TfLiteContext* context, TfLiteNode* node)
: context_(context), node_(node) {}
constexpr static int kElementShapeInput = 0;
constexpr static int kNumElementsInput = 1;
TfLiteStatus CheckInputs() const {
TF_LITE_ENSURE_EQ(context_, NumInputs(node_), 2);
const TfLiteTensor* element_shape;
TF_LITE_ENSURE_OK(
context_,
GetInputSafe(context_, node_, kElementShapeInput, &element_shape));
TF_LITE_ENSURE_TYPES_EQ(context_, element_shape->type, kTfLiteInt32);
const TfLiteTensor* num_elements;
TF_LITE_ENSURE_OK(context_, GetInputSafe(context_, node_, kNumElementsInput,
&num_elements));
TF_LITE_ENSURE_TYPES_EQ(context_, num_elements->type, kTfLiteInt32);
return kTfLiteOk;
}
TfLiteStatus Compute(SemanticOutType& result) const {
// Parse element type from custom options.
TF_LITE_ENSURE(context_, node_->custom_initial_data != nullptr);
TF_LITE_ENSURE(context_, node_->custom_initial_data_size >=
sizeof(ListReserveOptions));
auto* options =
reinterpret_cast<const ListReserveOptions*>(node_->custom_initial_data);
TfLiteType element_type = ConvertTensorType(options->element_type);
TF_LITE_ENSURE(context_, element_type != kTfLiteNoType);
const TfLiteTensor* num_elements;
TF_LITE_ENSURE_OK(context_, GetInputSafe(context_, node_, kNumElementsInput,
&num_elements));
TF_LITE_ENSURE_TYPES_EQ(context_, num_elements->type, kTfLiteInt32);
TF_LITE_ENSURE_EQ(context_, num_elements->dims->size, 0);
TF_LITE_ENSURE(context_, num_elements->data.data != nullptr);
const int num_elements_value = num_elements->data.i32[0];
TF_LITE_ENSURE(context_, num_elements_value >= 0);
// Create int array representing constraint on list's constituent elements.
const TfLiteTensor* element_shape_tensor;
TF_LITE_ENSURE_OK(context_,
GetInputSafe(context_, node_, kElementShapeInput,
&element_shape_tensor));
IntArrayUniquePtr element_shape;
TF_LITE_ENSURE_OK(context_, TensorAsShape(context_, *element_shape_tensor,
element_shape));
TF_LITE_ENSURE(context_, element_shape != nullptr);
result = SemanticOutType{element_type, std::move(element_shape),
num_elements_value};
return kTfLiteOk;
}
TfLiteStatus PopulateOutput(TensorArray* const output) const {
return kTfLiteOk;
}
private:
TfLiteContext* const context_;
TfLiteNode* const node_;
};
class ZerosLikeSemantic {
public:
ZerosLikeSemantic(TfLiteContext* context, TfLiteNode* node)
: context_(context), node_(node) {}
constexpr static int kListInput = 0;
TfLiteStatus CheckInputs() const {
TF_LITE_ENSURE_EQ(context_, NumInputs(node_), 1);
const TfLiteTensor* list_input;
TF_LITE_ENSURE_OK(context_,
GetInputSafe(context_, node_, kListInput, &list_input));
TF_LITE_ENSURE_TYPES_EQ(context_, list_input->type, kTfLiteVariant);
return kTfLiteOk;
}
TfLiteStatus Compute(SemanticOutType& result) const {
const TfLiteTensor* list_input;
TF_LITE_ENSURE_OK(context_,
GetInputSafe(context_, node_, kListInput, &list_input));
TF_LITE_ENSURE(context_, list_input->data.data != nullptr);
const TensorArray* const input = static_cast<const TensorArray*>(
static_cast<const VariantData*>(list_input->data.data));
TF_LITE_ENSURE(context_, input->ElementShape() != nullptr);
IntArrayUniquePtr element_shape = BuildTfLiteArray(*input->ElementShape());
TF_LITE_ENSURE(context_, element_shape != nullptr);
result = SemanticOutType{input->ElementType(), std::move(element_shape),
input->NumElements()};
return kTfLiteOk;
}
TfLiteStatus PopulateOutput(TensorArray* const output) const {
const TfLiteTensor* list_input;
TF_LITE_ENSURE_OK(context_,
GetInputSafe(context_, node_, kListInput, &list_input));
TF_LITE_ENSURE(context_, list_input->data.data != nullptr);
const TensorArray* const input = static_cast<const TensorArray*>(
static_cast<const VariantData*>(list_input->data.data));
for (int i = 0; i < input->NumElements(); ++i) {
const TfLiteTensor* const at = input->At(i);
if (at == nullptr) continue;
// Tensorflow supports lazy allocation in this case which is not possible
// with tflite tensors. If this proves to be a performance bottleneck we
// can investigate storing more info in TensorArray putting off allocation
// for later.
TensorUniquePtr output_at = BuildTfLiteTensor(
at->type, BuildTfLiteArray(*at->dims), kTfLiteDynamic);
TF_LITE_ENSURE(context_, output_at != nullptr);
if (output_at->bytes > 0) {
TF_LITE_ENSURE(context_, output_at->data.data != nullptr);
memset(output_at->data.data, 0, output_at->bytes);
}
TF_LITE_ENSURE(context_, output->Set(i, std::move(output_at)));
}
return kTfLiteOk;
}
private:
TfLiteContext* const context_;
TfLiteNode* const node_;
};
template <class Semantic>
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
const Semantic sem(context, node);
TF_LITE_ENSURE_OK(context, sem.CheckInputs());
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
TfLiteTensor* output;
TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, kListOut, &output));
TF_LITE_ENSURE_TYPES_EQ(context, output->type, kTfLiteVariant);
output->allocation_type = kTfLiteVariantObject;
return kTfLiteOk;
}
template <class Semantic>
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
const Semantic sem(context, node);
SemanticOutType data;
TF_LITE_ENSURE_OK(context, sem.Compute(data));
TfLiteTensor* output;
TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, kListOut, &output));
// Construct new `TensorArray` underneath the output tensor.
TfLiteStatus stat = TfLiteTensorVariantRealloc<TensorArray>(
output, data.element_type, std::move(data.element_shape));
TF_LITE_ENSURE_OK(context, stat);
// Set size of array.
TensorArray* const arr =
static_cast<TensorArray*>(static_cast<VariantData*>(output->data.data));
TF_LITE_ENSURE(context, arr->Resize(data.num_elements));
TF_LITE_ENSURE_OK(context, sem.PopulateOutput(arr));
return kTfLiteOk;
}
} // namespace
} // namespace list_reserve
TfLiteRegistration* Register_LIST_RESERVE() {
static TfLiteRegistration r = {
nullptr, nullptr, list_reserve::Prepare<list_reserve::ReserveSemantic>,
list_reserve::Eval<list_reserve::ReserveSemantic>};
return &r;
}
TfLiteRegistration* Register_VARIANT_ZEROS_LIKE() {
static TfLiteRegistration r = {
nullptr, nullptr, list_reserve::Prepare<list_reserve::ZerosLikeSemantic>,
list_reserve::Eval<list_reserve::ZerosLikeSemantic>};
return &r;
}
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/list_reserve_test.cc
+115
View File
@@ -0,0 +1,115 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstdint>
#include <cstring>
#include <vector>
#include <gtest/gtest.h>
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/schema/schema_generated.h"
namespace tflite {
namespace variants {
namespace ops {
namespace {
using ::tflite::variants::TensorArray;
std::vector<uint8_t> CustomOptionsToRaw(const std::vector<int32_t>& options) {
std::vector<uint8_t> raw(options.size() * sizeof(int32_t));
std::memcpy(raw.data(), options.data(), raw.size());
return raw;
}
class ListReserveModel : public SingleOpModel {
public:
explicit ListReserveModel(TensorType element_type) {
element_shape_input_ = AddInput({TensorType_INT32, {1}});
list_len_input_ = AddInput({TensorType_INT32, {}});
reserve_output_ = AddOutput({TensorType_VARIANT, {}});
SetCustomOp("ListReserve", CustomOptionsToRaw({element_type}),
Register_LIST_RESERVE);
BuildInterpreter({{1}, {}});
}
const TfLiteTensor* GetOutputTensor(int index) {
return interpreter_->tensor(index);
}
int list_len_input_;
int reserve_output_;
int element_shape_input_;
};
TEST(ListReserveTest, NonZeroNumElements_StaticShape) {
ListReserveModel m(TensorType_INT32);
m.PopulateTensor(m.list_len_input_, {5});
m.PopulateTensor(m.element_shape_input_, {2});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* tensor = m.GetOutputTensor(m.reserve_output_);
EXPECT_EQ(tensor->type, kTfLiteVariant);
EXPECT_EQ(tensor->allocation_type, kTfLiteVariantObject);
TensorArray* arr = reinterpret_cast<TensorArray*>(tensor->data.data);
EXPECT_EQ(arr->ElementType(), kTfLiteInt32);
EXPECT_EQ(arr->ElementShape()->size, 1);
ASSERT_EQ(arr->ElementShape()->data[0], 2);
ASSERT_EQ(arr->NumElements(), 5);
for (int i = 0; i < 5; ++i) {
ASSERT_EQ(arr->At(i), nullptr);
}
}
TEST(ListReserveTest, NegativeNumElements_Fails) {
ListReserveModel m(TensorType_INT32);
m.PopulateTensor(m.list_len_input_, {-1});
m.PopulateTensor(m.element_shape_input_, {2});
ASSERT_EQ(m.Invoke(), kTfLiteError);
}
TEST(ListReserveTest, NumElements0_StaticShape_Succeeds) {
ListReserveModel m(TensorType_INT32);
m.PopulateTensor(m.list_len_input_, {0});
m.PopulateTensor(m.element_shape_input_, {2});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* tensor = m.GetOutputTensor(m.reserve_output_);
TensorArray* arr = reinterpret_cast<TensorArray*>(tensor->data.data);
EXPECT_EQ(arr->NumElements(), 0);
EXPECT_EQ(arr->ElementType(), kTfLiteInt32);
}
TEST(ListReserveTest, NumElements0_StaticShape_FloatType) {
ListReserveModel m(TensorType_FLOAT32);
m.PopulateTensor(m.list_len_input_, {0});
m.PopulateTensor(m.element_shape_input_, {2});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* tensor = m.GetOutputTensor(m.reserve_output_);
TensorArray* arr = reinterpret_cast<TensorArray*>(tensor->data.data);
EXPECT_EQ(arr->NumElements(), 0);
EXPECT_EQ(arr->ElementType(), kTfLiteFloat32);
}
TEST(ListReserveTest, UnsupportedDataType_Fails) {
ListReserveModel m(TensorType_COMPLEX64);
m.PopulateTensor(m.list_len_input_, {0});
m.PopulateTensor(m.element_shape_input_, {2});
ASSERT_EQ(m.Invoke(), kTfLiteError);
}
} // namespace
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/list_set_item.cc
+175
View File
@@ -0,0 +1,175 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <utility>
#include "tensorflow/lite/array.h"
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/util.h"
namespace tflite {
namespace variants {
namespace ops {
namespace {
constexpr int kListInputIdx = 0;
constexpr int kIndexInputIdx = 1;
constexpr int kListOutputIdx = 0;
class SetItemSemantic {
public:
SetItemSemantic(TfLiteContext* ctx, TfLiteNode* node)
: ctx_(ctx), node_(node) {}
static constexpr int kItemInputIdx = 2;
TfLiteStatus CheckIndexInput() const {
const TfLiteTensor* index_input;
TF_LITE_ENSURE_OK(ctx_,
GetInputSafe(ctx_, node_, kIndexInputIdx, &index_input));
TF_LITE_ENSURE_TYPES_EQ(ctx_, index_input->type, kTfLiteInt32);
return kTfLiteOk;
}
TfLiteStatus GetIndexVal(const TensorArray& /*arr*/, int& result) const {
const TfLiteTensor* index_input;
TF_LITE_ENSURE_OK(ctx_,
GetInputSafe(ctx_, node_, kIndexInputIdx, &index_input));
TF_LITE_ENSURE_EQ(ctx_, index_input->bytes, sizeof(int));
const int* index_data = GetTensorData<int>(index_input);
TF_LITE_ENSURE(ctx_, index_data != nullptr);
const int index = *index_data;
TF_LITE_ENSURE(ctx_, index >= 0);
result = index;
return kTfLiteOk;
}
private:
TfLiteContext* const ctx_;
TfLiteNode* const node_;
};
class PushBackSemantic {
public:
PushBackSemantic(TfLiteContext* /*ctx*/, TfLiteNode* /*node*/) {}
static constexpr int kItemInputIdx = 1;
TfLiteStatus CheckIndexInput() const { return kTfLiteOk; }
TfLiteStatus GetIndexVal(const TensorArray& arr, int& result) const {
result = arr.NumElements();
return kTfLiteOk;
}
};
template <class Semantic>
TfLiteStatus Prepare(TfLiteContext* ctx, TfLiteNode* node) {
const auto semantic = Semantic(ctx, node);
TF_LITE_ENSURE(ctx, NumInputs(node) >= semantic.kItemInputIdx + 1);
const TfLiteTensor* list_input;
TF_LITE_ENSURE_OK(ctx, GetInputSafe(ctx, node, kListInputIdx, &list_input));
TF_LITE_ENSURE_TYPES_EQ(ctx, list_input->type, kTfLiteVariant);
TF_LITE_ENSURE_OK(ctx, semantic.CheckIndexInput());
const TfLiteTensor* item_input;
TF_LITE_ENSURE_OK(
ctx, GetInputSafe(ctx, node, semantic.kItemInputIdx, &item_input));
TfLiteTensor* output;
TF_LITE_ENSURE_OK(ctx, GetOutputSafe(ctx, node, kListOutputIdx, &output));
TF_LITE_ENSURE_TYPES_EQ(ctx, output->type, kTfLiteVariant);
output->allocation_type = kTfLiteVariantObject;
return kTfLiteOk;
}
template <class Semantic>
TfLiteStatus Eval(TfLiteContext* ctx, TfLiteNode* node) {
const auto semantic = Semantic(ctx, node);
const TfLiteTensor* list_input;
TF_LITE_ENSURE_OK(ctx, GetInputSafe(ctx, node, kListInputIdx, &list_input));
TF_LITE_ENSURE_EQ(ctx, list_input->allocation_type, kTfLiteVariantObject);
TF_LITE_ENSURE(ctx, list_input->data.data != nullptr);
TensorArray* input_arr = static_cast<TensorArray*>(
static_cast<VariantData*>(list_input->data.data));
TF_LITE_ENSURE(ctx, input_arr != nullptr);
int index;
TF_LITE_ENSURE_OK(ctx, semantic.GetIndexVal(*input_arr, index));
const TfLiteTensor* item_input;
TF_LITE_ENSURE_OK(
ctx, GetInputSafe(ctx, node, semantic.kItemInputIdx, &item_input));
TF_LITE_ENSURE_TYPES_EQ(ctx, input_arr->ElementType(), item_input->type);
TfLiteTensor* output;
TF_LITE_ENSURE_OK(ctx, GetOutputSafe(ctx, node, kListOutputIdx, &output));
TensorArray* output_arr = static_cast<TensorArray*>(
input_arr->CloneTo(static_cast<VariantData*>(output->data.data)));
TF_LITE_ENSURE(ctx, output_arr != nullptr);
output->data.data = static_cast<VariantData*>(output_arr);
// TODO(b/288302706) Skip copy when tensor is used only once.
TensorUniquePtr item_copy;
if (item_input->type == kTfLiteVariant) {
item_copy = BuildTfLiteTensor();
TF_LITE_ENSURE(ctx, item_copy != nullptr);
item_copy->type = kTfLiteVariant;
item_copy->bytes = item_input->bytes;
item_copy->allocation_type = kTfLiteVariantObject;
} else {
TF_LITE_ENSURE(ctx, item_input->dims != nullptr);
IntArrayUniquePtr item_dims = BuildTfLiteArray(*item_input->dims);
TF_LITE_ENSURE(ctx, item_dims != nullptr);
item_copy = BuildTfLiteTensor(item_input->type, std::move(item_dims),
kTfLiteDynamic);
TF_LITE_ENSURE(ctx, item_copy != nullptr);
if (item_copy->bytes > 0) {
TF_LITE_ENSURE(ctx, item_copy->data.data != nullptr);
}
}
TF_LITE_ENSURE_OK(ctx, TfLiteTensorCopy(item_input, item_copy.get()));
if (index >= output_arr->NumElements()) {
TF_LITE_ENSURE(ctx, output_arr->Resize(index + 1));
}
TF_LITE_ENSURE(ctx, output_arr->Set(index, std::move(item_copy)));
return kTfLiteOk;
}
} // namespace
TfLiteRegistration* Register_LIST_SET_ITEM() {
static TfLiteRegistration r = {nullptr, nullptr, Prepare<SetItemSemantic>,
Eval<SetItemSemantic>};
return &r;
}
TfLiteRegistration* Register_LIST_PUSH_BACK() {
static TfLiteRegistration r = {nullptr, nullptr, Prepare<PushBackSemantic>,
Eval<PushBackSemantic>};
return &r;
}
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/list_set_item_test.cc
+265
View File
@@ -0,0 +1,265 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstdint>
#include <optional>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/internal/compatibility.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/portable_type_to_tflitetype.h"
#include "tensorflow/lite/schema/schema_generated.h"
namespace tflite {
namespace variants {
namespace ops {
namespace {
using ::testing::AllOf;
template <typename T>
class SetItemWithTypeTest : public ::testing::Test {};
class ListSetItemModel : public ListOpModel {
public:
explicit ListSetItemModel(TensorData item_data) {
list_input_ = AddInput({TensorType_VARIANT, {}});
index_input_ = AddInput({TensorType_INT32, {1}});
tensor_input_ = AddInput(item_data);
list_output_ = AddOutput({TensorType_VARIANT, {}});
SetCustomOp("ListSetItem", {}, Register_LIST_SET_ITEM);
BuildInterpreter({{}, {1}, item_data.shape});
interpreter_->input_tensor(0)->allocation_type = kTfLiteVariantObject;
}
const TensorArray* GetOutputTensorArray(int tensor_id) {
TfLiteTensor* tensor = interpreter_->tensor(tensor_id);
TFLITE_CHECK(tensor != nullptr && tensor->type == kTfLiteVariant &&
tensor->allocation_type == kTfLiteVariantObject);
return static_cast<const TensorArray*>(
static_cast<const VariantData*>(tensor->data.data));
}
int index_input_;
int list_input_;
int tensor_input_;
int list_output_;
};
constexpr int kNumElements = 4;
TYPED_TEST_SUITE_P(SetItemWithTypeTest);
TYPED_TEST_P(SetItemWithTypeTest, SetItemOnEmptyTensorList_ListShapeDefined) {
TfLiteType tfl_type = typeToTfLiteType<TypeParam>();
std::optional<TensorType> tensor_type = TflToTensorType(tfl_type);
ASSERT_TRUE(tensor_type.has_value());
ListSetItemModel m({tensor_type.value(), {2, 2}});
m.PopulateTensor(m.index_input_, {0});
m.PopulateListTensor(m.list_input_, {2, 2}, kNumElements, tfl_type);
m.PopulateTensor<TypeParam>(m.tensor_input_, {0, 0, 0, 0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* arr = m.GetOutputTensorArray(m.list_output_);
ASSERT_EQ(arr->NumElements(), kNumElements);
ASSERT_EQ(arr->ElementType(), tfl_type);
for (int i = 1; i < arr->NumElements(); ++i) {
EXPECT_EQ(arr->At(i), nullptr);
}
EXPECT_THAT(arr->At(0), AllOf(IsAllocatedAs(tfl_type), DimsAre({2, 2}),
FilledWith(static_cast<TypeParam>(0))));
}
TYPED_TEST_P(SetItemWithTypeTest, SetItemOnEmptyTensorList_ListShapeUnranked) {
TfLiteType tfl_type = typeToTfLiteType<TypeParam>();
std::optional<TensorType> tensor_type = TflToTensorType(tfl_type);
ASSERT_TRUE(tensor_type.has_value());
ListSetItemModel m({tensor_type.value(), {2, 2}});
m.PopulateTensor(m.index_input_, {0});
m.PopulateListTensor(m.list_input_, {}, kNumElements, tfl_type);
m.PopulateTensor<TypeParam>(m.tensor_input_, {0, 0, 0, 0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* arr = m.GetOutputTensorArray(m.list_output_);
ASSERT_EQ(arr->NumElements(), kNumElements);
ASSERT_EQ(arr->ElementType(), tfl_type);
for (int i = 1; i < arr->NumElements(); ++i) {
EXPECT_EQ(arr->At(i), nullptr);
}
EXPECT_THAT(arr->At(0), AllOf(IsAllocatedAs(tfl_type), DimsAre({2, 2}),
FilledWith(static_cast<TypeParam>(0))));
}
TYPED_TEST_P(SetItemWithTypeTest, OverwriteSetItem_ItemsSameShape) {
TfLiteType tfl_type = typeToTfLiteType<TypeParam>();
std::optional<TensorType> tensor_type = TflToTensorType(tfl_type);
ASSERT_TRUE(tensor_type.has_value());
ListSetItemModel m({tensor_type.value(), {2, 2}});
m.PopulateTensor(m.index_input_, {0});
m.PopulateListTensor(m.list_input_, {}, kNumElements, tfl_type);
TypeParam init_item_data[4] = {1, 1, 1, 1};
m.ListSetItem(m.list_input_, 0, {2, 2}, tfl_type, init_item_data);
m.PopulateTensor<TypeParam>(m.tensor_input_, {0, 0, 0, 0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* arr = m.GetOutputTensorArray(m.list_output_);
ASSERT_EQ(arr->NumElements(), kNumElements);
ASSERT_EQ(arr->ElementType(), tfl_type);
for (int i = 1; i < arr->NumElements(); ++i) {
EXPECT_EQ(arr->At(i), nullptr);
}
EXPECT_THAT(arr->At(0), AllOf(IsAllocatedAs(tfl_type), DimsAre({2, 2}),
FilledWith(static_cast<TypeParam>(0))));
}
TYPED_TEST_P(SetItemWithTypeTest,
SetItemOnNonEmptyListAtEmptyIndex_ItemsSameShape) {
TfLiteType tfl_type = typeToTfLiteType<TypeParam>();
std::optional<TensorType> tensor_type = TflToTensorType(tfl_type);
ASSERT_TRUE(tensor_type.has_value());
ListSetItemModel m({tensor_type.value(), {2, 2}});
m.PopulateTensor(m.index_input_, {1});
m.PopulateListTensor(m.list_input_, {}, kNumElements, tfl_type);
TypeParam init_item_data[4] = {1, 1, 1, 1};
m.ListSetItem(m.list_input_, 0, {2, 2}, tfl_type, init_item_data);
m.PopulateTensor<TypeParam>(m.tensor_input_, {0, 0, 0, 0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* arr = m.GetOutputTensorArray(m.list_output_);
ASSERT_EQ(arr->NumElements(), kNumElements);
ASSERT_EQ(arr->ElementType(), tfl_type);
for (int i = 2; i < arr->NumElements(); ++i) {
EXPECT_EQ(arr->At(i), nullptr);
}
EXPECT_THAT(arr->At(0), AllOf(IsAllocatedAs(tfl_type), DimsAre({2, 2}),
FilledWith(static_cast<TypeParam>(1))));
EXPECT_THAT(arr->At(1), AllOf(IsAllocatedAs(tfl_type), DimsAre({2, 2}),
FilledWith(static_cast<TypeParam>(0))));
}
TYPED_TEST_P(SetItemWithTypeTest, OverwriteSetItem_ItemsDifferentShape) {
TfLiteType tfl_type = typeToTfLiteType<TypeParam>();
std::optional<TensorType> tensor_type = TflToTensorType(tfl_type);
ASSERT_TRUE(tensor_type.has_value());
ListSetItemModel m({tensor_type.value(), {2}});
m.PopulateTensor(m.index_input_, {0});
m.PopulateListTensor(m.list_input_, {}, kNumElements, tfl_type);
TypeParam init_item_data[4] = {1, 1, 1, 1};
m.ListSetItem(m.list_input_, 0, {2, 2}, tfl_type, init_item_data);
m.PopulateTensor<TypeParam>(m.tensor_input_, {0, 0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* arr = m.GetOutputTensorArray(m.list_output_);
ASSERT_EQ(arr->NumElements(), kNumElements);
ASSERT_EQ(arr->ElementType(), tfl_type);
for (int i = 1; i < arr->NumElements(); ++i) {
EXPECT_EQ(arr->At(i), nullptr);
}
EXPECT_THAT(arr->At(0), AllOf(IsAllocatedAs(tfl_type), DimsAre({2}),
FilledWith(static_cast<TypeParam>(0))));
}
REGISTER_TYPED_TEST_SUITE_P(SetItemWithTypeTest,
SetItemOnEmptyTensorList_ListShapeDefined,
SetItemOnEmptyTensorList_ListShapeUnranked,
OverwriteSetItem_ItemsSameShape,
SetItemOnNonEmptyListAtEmptyIndex_ItemsSameShape,
OverwriteSetItem_ItemsDifferentShape);
using ValidTypes = ::testing::Types<int, int64_t, bool, float>;
INSTANTIATE_TYPED_TEST_SUITE_P(SetItemTests, SetItemWithTypeTest, ValidTypes);
TEST(ListSetItemTest, ItemNotSameTypeAsList_Fails) {
ListSetItemModel m{{TensorType_INT32, {2, 2}}};
m.PopulateTensor(m.index_input_, {0});
m.PopulateListTensor(m.list_input_, {}, kNumElements, kTfLiteInt64);
ASSERT_EQ(m.Invoke(), kTfLiteError);
}
TEST(ListSetItemTest, IndexLessThanZero_Fails) {
ListSetItemModel m{{TensorType_INT32, {2, 2}}};
m.PopulateTensor(m.index_input_, {-1});
m.PopulateListTensor(m.list_input_, {}, kNumElements, kTfLiteInt32);
ASSERT_EQ(m.Invoke(), kTfLiteError);
}
TEST(ListSetItemTest, IndexLessGreaterThanListLen_ResizesList) {
ListSetItemModel m{{TensorType_INT32, {2, 2}}};
m.PopulateTensor(m.index_input_, {2});
m.PopulateListTensor(m.list_input_, {}, 2, kTfLiteInt32);
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* arr = m.GetOutputTensorArray(m.list_output_);
ASSERT_EQ(arr->NumElements(), 3);
}
} // namespace
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/list_stack.cc
+160
View File
@@ -0,0 +1,160 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstring>
#include <utility>
#include "tensorflow/lite/array.h"
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_util.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
namespace tflite {
namespace variants {
namespace ops {
namespace {
constexpr int kListInput = 0;
constexpr int kShapeInput = 1;
constexpr int kTensorOutput = 0;
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
const TfLiteTensor* list_input;
TF_LITE_ENSURE_OK(context,
GetInputSafe(context, node, kListInput, &list_input));
TF_LITE_ENSURE_TYPES_EQ(context, list_input->type, kTfLiteVariant);
const TfLiteTensor* shape_input;
TF_LITE_ENSURE_OK(context,
GetInputSafe(context, node, kShapeInput, &shape_input));
TF_LITE_ENSURE_TYPES_EQ(context, shape_input->type, kTfLiteInt32);
TfLiteTensor* output;
TF_LITE_ENSURE_OK(context,
GetOutputSafe(context, node, kTensorOutput, &output));
// TODO(b/257472333) Consider leveraging arena when the shape is defined
// at compile time.
SetTensorToDynamic(output);
return kTfLiteOk;
}
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
const TfLiteTensor* list_input;
TF_LITE_ENSURE_OK(context,
GetInputSafe(context, node, kListInput, &list_input));
TF_LITE_ENSURE_EQ(context, list_input->allocation_type, kTfLiteVariantObject);
TF_LITE_ENSURE(context, list_input->data.data != nullptr);
TensorArray* arr = static_cast<TensorArray*>(
static_cast<VariantData*>(list_input->data.data));
const TfLiteTensor* shape_input;
TF_LITE_ENSURE_OK(context,
GetInputSafe(context, node, kShapeInput, &shape_input));
TfLiteTensor* output;
TF_LITE_ENSURE_OK(context,
GetOutputSafe(context, node, kTensorOutput, &output));
TF_LITE_ENSURE_TYPES_EQ(context, output->type, arr->ElementType());
IntArrayUniquePtr cur_shape_suffix;
// If succeeds and result not nullptr, guaranteed to be fully defined.
TF_LITE_ENSURE_OK(context, GetShapeIfAllEqual(*arr, cur_shape_suffix));
IntArrayUniquePtr shape_input_arr;
TF_LITE_ENSURE_OK(context,
TensorAsShape(context, *shape_input, shape_input_arr));
TF_LITE_ENSURE(context, shape_input_arr != nullptr);
// Confirm that input shape, shape of elements and list shape are all
// compatible.
cur_shape_suffix = MergeShapesOrNull(
MergeShapesOrNull(std::move(shape_input_arr),
BuildTfLiteArray(*arr->ElementShape())),
std::move(cur_shape_suffix));
TF_LITE_ENSURE_MSG(
context,
cur_shape_suffix != nullptr && IsShapeFullyDefined(*cur_shape_suffix),
"Shapes from input, list and elements are not compatible "
"or do not resolve to fully defined shape.");
// Now compute the first dimension and concat with computed suffix.
IntArrayUniquePtr final_output_shape;
const bool suffix_is_not_scalar =
!(cur_shape_suffix->size == 0 ||
(cur_shape_suffix->size == 1 && cur_shape_suffix->data[0] == 1));
if (suffix_is_not_scalar) {
final_output_shape = BuildTfLiteArray(cur_shape_suffix->size + 1);
memcpy(final_output_shape->data + 1, cur_shape_suffix->data,
cur_shape_suffix->size * sizeof(int));
final_output_shape->data[0] = arr->NumElements();
} else {
final_output_shape = BuildTfLiteArray({arr->NumElements()});
}
TF_LITE_ENSURE_OK(
context,
context->ResizeTensor(context, output, final_output_shape.release()));
const auto num_elements = static_cast<int>(NumElements(output));
if (num_elements == 0) {
TfLiteTensorDataFree(output);
return kTfLiteOk;
}
// This has to be an int and we would have returned already if divisor == 0.
const int element_num_elements = num_elements / output->dims->data[0];
const size_t bytes_per_element =
element_num_elements * TfLiteTypeGetSize(output->type);
// Copy buffer of constituent element tensors to output if they are present.
// Otherwise, zero that chunk of memory.
char* raw_data_offset = output->data.raw;
if (bytes_per_element > 0) {
TF_LITE_ENSURE(context, raw_data_offset != nullptr);
for (int i = 0; i < arr->NumElements(); ++i) {
if (arr->At(i) == nullptr) {
memset(raw_data_offset, 0, bytes_per_element);
} else {
TF_LITE_ENSURE(context, arr->At(i)->data.data != nullptr);
memcpy(raw_data_offset, arr->At(i)->data.data, bytes_per_element);
}
raw_data_offset = raw_data_offset + bytes_per_element;
}
}
return kTfLiteOk;
}
} // namespace
TfLiteRegistration* Register_LIST_STACK() {
static TfLiteRegistration r = {nullptr, nullptr, Prepare, Eval};
return &r;
}
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/list_stack_test.cc
+324
View File
@@ -0,0 +1,324 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/schema/schema_generated.h"
using ::testing::ElementsAreArray;
using ::tflite::variants::ops::Register_LIST_STACK;
namespace tflite {
namespace {
class ListStackModel : public ListOpModel {
public:
explicit ListStackModel(TensorData output_data) {
tensor_id_ = AddOutput(output_data);
list_id_ = AddInput({TensorType_VARIANT, {}});
shape_id_ = AddInput({TensorType_INT32, {1}});
SetCustomOp("ListStack", {}, Register_LIST_STACK);
BuildInterpreter({{}, {1}});
}
ListStackModel(TensorData output_data, TensorData shape_input_data) {
tensor_id_ = AddOutput(output_data);
list_id_ = AddInput({TensorType_VARIANT, {}});
shape_id_ = AddInput(shape_input_data);
SetCustomOp("ListStack", {}, Register_LIST_STACK);
BuildInterpreter({{}, shape_input_data.shape});
}
const TfLiteTensor* GetOutputTensor(int tensor_id) {
return interpreter_->tensor(tensor_id);
}
int tensor_id_;
int shape_id_;
int list_id_;
};
TEST(ListStackTest, MismatchedListShapeInputShape_Fails) {
ListStackModel m({TensorType_INT32, {2, 2}});
m.PopulateListTensor(m.list_id_, {1}, 2, kTfLiteInt32);
m.PopulateTensor(m.shape_id_, {3});
ASSERT_EQ(m.Invoke(), kTfLiteError);
}
TEST(ListStackTest, MismatchedShapeOfElementsAndInput_Fails) {
ListStackModel m({TensorType_INT32, {2, 2}});
m.PopulateListTensor(m.list_id_, {}, 4, kTfLiteInt32);
m.PopulateTensor(m.shape_id_, {2});
m.ListSetItem(m.list_id_, 0, {1}, kTfLiteInt32, std::vector<int>{0}.data());
m.ListSetItem(m.list_id_, 1, {1}, kTfLiteInt32, std::vector<int>{1}.data());
ASSERT_EQ(m.Invoke(), kTfLiteError);
}
TEST(ListStackTest, ElementsNotSameShape_Fails) {
ListStackModel m({TensorType_INT32, {2, 2}});
m.PopulateListTensor(m.list_id_, {}, 2, kTfLiteInt32);
m.PopulateTensor(m.shape_id_, {2});
m.ListSetItem(m.list_id_, 0, {2}, kTfLiteInt32,
std::vector<int>{2, 2}.data());
m.ListSetItem(m.list_id_, 1, {1}, kTfLiteInt32, std::vector<int>{3}.data());
ASSERT_EQ(m.Invoke(), kTfLiteError);
}
TEST(ListStackTest, NoElementsNoShape_Fails) {
ListStackModel m({TensorType_INT32, {4}});
m.PopulateListTensor(m.list_id_, {}, 2, kTfLiteInt32);
m.PopulateTensor<int>(m.shape_id_, {-1});
EXPECT_EQ(m.Invoke(), kTfLiteError);
}
TEST(ListStackTest, ListElementTypeNotEqualOutputType_Fails) {
ListStackModel m({TensorType_INT32, {4}});
m.PopulateListTensor(m.list_id_, {}, 0, kTfLiteInt64);
m.PopulateTensor<int>(m.shape_id_, {-1});
EXPECT_EQ(m.Invoke(), kTfLiteError);
}
TEST(ListStackTest, ScalarElementShape_FullList_Returns1D) {
ListStackModel m({TensorType_INT32, {2}});
m.PopulateListTensor(m.list_id_, {}, 2, kTfLiteInt32);
m.PopulateTensor(m.shape_id_, {1});
m.ListSetItem(m.list_id_, 0, {1}, kTfLiteInt32, std::vector<int>{2}.data());
m.ListSetItem(m.list_id_, 1, {1}, kTfLiteInt32, std::vector<int>{3}.data());
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutputTensor(m.tensor_id_);
ASSERT_THAT(output, DimsAre({2}));
ASSERT_THAT(output->type, kTfLiteInt32);
EXPECT_THAT(std::vector<int>(output->data.i32, output->data.i32 + 2),
ElementsAreArray({2, 3}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
}
TEST(ListStackTest, ScalarElementShape_PartialFilledList_Returns1DWithZeroed) {
ListStackModel m({TensorType_INT32, {2}});
m.PopulateListTensor(m.list_id_, {}, 2, kTfLiteInt32);
m.PopulateTensor(m.shape_id_, {1});
m.ListSetItem(m.list_id_, 0, {1}, kTfLiteInt32, std::vector<int>{2}.data());
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutputTensor(m.tensor_id_);
ASSERT_THAT(output, DimsAre({2}));
ASSERT_THAT(output->type, kTfLiteInt32);
EXPECT_THAT(std::vector<int>(output->data.i32, output->data.i32 + 2),
ElementsAreArray({2, 0}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
}
TEST(ListStackTest, ScalarElementShape_EmptyList_Returns1DAllZeroed) {
ListStackModel m({TensorType_INT32, {2}});
m.PopulateListTensor(m.list_id_, {}, 2, kTfLiteInt32);
m.PopulateTensor(m.shape_id_, {1});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutputTensor(m.tensor_id_);
ASSERT_THAT(output, DimsAre({2}));
ASSERT_THAT(output->type, kTfLiteInt32);
EXPECT_THAT(std::vector<int>(output->data.i32, output->data.i32 + 2),
ElementsAreArray({0, 0}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
}
TEST(ListStackTest, VectorElementShape_FilledList_Returns2D) {
ListStackModel m({TensorType_INT32, {2, 2}});
m.PopulateListTensor(m.list_id_, {}, 2, kTfLiteInt32);
m.PopulateTensor<int>(m.shape_id_, {2});
m.ListSetItem(m.list_id_, 0, {2}, kTfLiteInt32,
std::vector<int>{2, 2}.data());
m.ListSetItem(m.list_id_, 1, {2}, kTfLiteInt32,
std::vector<int>{3, 3}.data());
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutputTensor(m.tensor_id_);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_THAT(std::vector<int>(output->data.i32, output->data.i32 + 4),
ElementsAreArray({2, 2, 3, 3}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
}
TEST(ListStackTest, VectorElementShape_PartialFilledList_Returns2DWithZeroed) {
ListStackModel m({TensorType_INT32, {2, 2}});
m.PopulateListTensor(m.list_id_, {}, 2, kTfLiteInt32);
m.PopulateTensor<int>(m.shape_id_, {2});
m.ListSetItem(m.list_id_, 0, {2}, kTfLiteInt32,
std::vector<int>{2, 2}.data());
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutputTensor(m.tensor_id_);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_THAT(std::vector<int>(output->data.i32, output->data.i32 + 4),
ElementsAreArray({2, 2, 0, 0}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
}
TEST(ListStackTest, VectorElementShape_EmptyList_Returns2DAllZeroed) {
ListStackModel m({TensorType_INT32, {2, 2}});
m.PopulateListTensor(m.list_id_, {}, 2, kTfLiteInt32);
m.PopulateTensor<int>(m.shape_id_, {2});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutputTensor(m.tensor_id_);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_THAT(std::vector<int>(output->data.i32, output->data.i32 + 4),
ElementsAreArray({0, 0, 0, 0}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
}
TEST(ListStackTest, NoShapeArguments_ZeroSizeList_InfersShapeFromElements) {
ListStackModel m({TensorType_INT32, {2, 2}});
m.PopulateListTensor(m.list_id_, {}, 2, kTfLiteInt32);
m.PopulateTensor<int>(m.shape_id_, {-1});
m.ListSetItem(m.list_id_, 0, {2}, kTfLiteInt32,
std::vector<int>{2, 2}.data());
EXPECT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutputTensor(m.tensor_id_);
ASSERT_THAT(output, DimsAre({2, 2}));
EXPECT_THAT(std::vector<int>(output->data.i32, output->data.i32 + 4),
ElementsAreArray({2, 2, 0, 0}));
EXPECT_EQ(output->allocation_type, kTfLiteDynamic);
}
TEST(ListStackTest, ListFirstDimZero_ReturnsEmptyTensor) {
ListStackModel m({TensorType_INT32, {0, 2}});
m.PopulateListTensor(m.list_id_, {}, 0, kTfLiteInt32);
m.PopulateTensor<int>(m.shape_id_, {2});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutputTensor(m.tensor_id_);
EXPECT_THAT(output, DimsAre({0, 2}));
}
TEST(ListStackTest, MismatchedOutput_ReturnsResizedOutput1D) {
ListStackModel m({TensorType_INT32, {2}});
m.PopulateListTensor(m.list_id_, {}, 4, kTfLiteInt32);
m.PopulateTensor<int>(m.shape_id_, {1});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutputTensor(m.tensor_id_);
EXPECT_THAT(output, DimsAre({4}));
}
TEST(ListStackTest, MismatchedOutput_ReturnsResizedOutput2D) {
ListStackModel m({TensorType_INT32, std::vector<int>{}});
m.PopulateListTensor(m.list_id_, {}, 2, kTfLiteInt32);
m.PopulateTensor<int>(m.shape_id_, {2});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutputTensor(m.tensor_id_);
EXPECT_THAT(output, DimsAre({2, 2}));
}
TEST(ListStackTest, Trailing0DimInElementShape1D_NonZeroLen_Returns2DNoData) {
ListStackModel m({TensorType_INT32, std::vector<int>{}});
m.PopulateListTensor(m.list_id_, {}, 2, kTfLiteInt32);
m.PopulateTensor<int>(m.shape_id_, {0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutputTensor(m.tensor_id_);
ASSERT_THAT(output, DimsAre({2, 0}));
EXPECT_EQ(output->bytes, 0);
}
TEST(ListStackTest, Trailing0DimInElementShape2D_NonZeroLen_Returns3DNoData) {
ListStackModel m({TensorType_INT32, {}}, {TensorType_INT32, {2}});
m.PopulateListTensor(m.list_id_, {}, 2, kTfLiteInt32);
m.PopulateTensor<int>(m.shape_id_, {2, 0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutputTensor(m.tensor_id_);
ASSERT_THAT(output, DimsAre({2, 2, 0}));
EXPECT_EQ(output->bytes, 0);
}
TEST(ListStackTest, Trailing0DimInElementShape1D_ZeroLen_Returns2DNoData) {
ListStackModel m({TensorType_INT32, {}}, {TensorType_INT32, {1}});
m.PopulateListTensor(m.list_id_, {}, 0, kTfLiteInt32);
m.PopulateTensor<int>(m.shape_id_, {0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutputTensor(m.tensor_id_);
ASSERT_THAT(output, DimsAre({0, 0}));
EXPECT_EQ(output->bytes, 0);
}
TEST(ListStackTest, Trailing0DimInElementShape2D_ZeroLen_Returns3DNoData) {
ListStackModel m({TensorType_INT32, {}}, {TensorType_INT32, {2}});
m.PopulateListTensor(m.list_id_, {}, 0, kTfLiteInt32);
m.PopulateTensor<int>(m.shape_id_, {2, 0});
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TfLiteTensor* output = m.GetOutputTensor(m.tensor_id_);
ASSERT_THAT(output, DimsAre({0, 2, 0}));
EXPECT_EQ(output->bytes, 0);
}
} // namespace
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/test_util.cc
+146
View File
@@ -0,0 +1,146 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/lite/kernels/variants/list_kernels/test_util.h"
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <optional>
#include <utility>
#include <vector>
#include "absl/types/span.h"
#include "tensorflow/lite/array.h"
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/op_macros.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/util.h"
using ::tflite::variants::TensorArray;
namespace tflite {
// Populate tensor at given index as a TensorList.
void ListOpModel::PopulateListTensor(int index,
absl::Span<const int> element_shape_data,
int num_elements,
TfLiteType element_type) {
TfLiteTensor* tensor = interpreter_->tensor(index);
TF_LITE_ASSERT_EQ(tensor->type, kTfLiteVariant);
tensor->allocation_type = kTfLiteVariantObject;
tensor->buffer_handle = kTfLiteNullBufferHandle;
tensor->quantization = {kTfLiteNoQuantization};
IntArrayUniquePtr element_shape =
BuildTfLiteArray(element_shape_data.size(), element_shape_data.data());
TfLiteStatus stat = TfLiteTensorVariantRealloc<TensorArray>(
tensor, element_type, std::move(element_shape));
TF_LITE_ASSERT_EQ(stat, kTfLiteOk);
TensorArray* arr =
static_cast<TensorArray*>(static_cast<VariantData*>(tensor->data.data));
arr->Resize(num_elements);
}
// Set a list element with given data.
void ListOpModel::ListSetItem(int index, int list_index,
absl::Span<const int> item_dims,
TfLiteType item_type, const void* item_data) {
TfLiteTensor* tensor = interpreter_->tensor(index);
TF_LITE_ASSERT_EQ(tensor->type, kTfLiteVariant);
TF_LITE_ASSERT_EQ(tensor->allocation_type, kTfLiteVariantObject);
TensorArray* arr =
static_cast<TensorArray*>(static_cast<VariantData*>(tensor->data.data));
TF_LITE_ASSERT_EQ(arr->ElementType(), item_type);
// Build tensor to set in list.
TensorUniquePtr item = BuildTfLiteTensor(
item_type, std::vector<int>(item_dims.begin(), item_dims.end()),
kTfLiteDynamic);
size_t item_type_size;
TF_LITE_ASSERT_EQ(GetSizeOfType(nullptr, item_type, &item_type_size),
kTfLiteOk);
// Write data to item tensor.
const size_t item_data_size = NumElements(item.get()) * item_type_size;
memcpy(item->data.data, item_data, item_data_size);
TF_LITE_ASSERT_EQ(arr->Set(list_index, std::move(item)), true);
}
std::optional<size_t> TfLiteTypeSizeOf(TfLiteType type) {
switch (type) {
case kTfLiteFloat16:
return sizeof(int16_t);
case kTfLiteFloat32:
return sizeof(float);
case kTfLiteFloat64:
return sizeof(double);
case kTfLiteInt16:
return sizeof(int16_t);
case kTfLiteInt32:
return sizeof(int32_t);
case kTfLiteUInt32:
return sizeof(uint32_t);
case kTfLiteUInt8:
return sizeof(uint8_t);
case kTfLiteInt8:
return sizeof(int8_t);
case kTfLiteInt64:
return sizeof(int64_t);
case kTfLiteUInt64:
return sizeof(uint64_t);
case kTfLiteBool:
return sizeof(bool);
case kTfLiteResource:
return sizeof(int32_t);
case kTfLiteComplex64:
return sizeof(float) * 2;
case kTfLiteComplex128:
return sizeof(double) * 2;
case kTfLiteInt4:
return sizeof(int8_t);
case kTfLiteUInt16:
return sizeof(uint16_t);
default:
return std::nullopt;
}
}
std::optional<TensorType> TflToTensorType(TfLiteType tfl_type) {
switch (tfl_type) {
case kTfLiteInt32:
return TensorType_INT32;
case kTfLiteFloat32:
return TensorType_FLOAT32;
case kTfLiteInt64:
return TensorType_INT64;
case kTfLiteBool:
return TensorType_BOOL;
default:
return std::nullopt;
}
}
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/test_util.h
+104
View File
@@ -0,0 +1,104 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_VARIANTS_LIST_KERNELS_TEST_UTIL_H_
#define TENSORFLOW_LITE_KERNELS_VARIANTS_LIST_KERNELS_TEST_UTIL_H_
#include <cstring>
#include <optional>
#include <string>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include "absl/types/span.h"
#include "tensorflow/lite/array.h"
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/portable_type_to_tflitetype.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/util.h"
namespace tflite {
class ListOpModel : public MultiOpModel {
public:
// Populate tensor at given index as a TensorList.
void PopulateListTensor(int index, absl::Span<const int> element_shape_data,
int num_elements, TfLiteType element_type);
// Set a list element with given data.
void ListSetItem(int index, int list_index, absl::Span<const int> item_dims,
TfLiteType item_type, const void* item_data);
};
// Gets the number of bytes required for a single element of `TfLiteType`,
// or `nullopt` if type does not have a fixed size.
std::optional<size_t> TfLiteTypeSizeOf(TfLiteType type);
// Converts `TfLiteType` to `TensorType`.
std::optional<TensorType> TflToTensorType(TfLiteType tfl_type);
// Matches tensors on having the given type and buffer of proper size.
MATCHER_P(IsAllocatedAs, tfl_type,
std::string(negation ? "is not" : "is") +
" of given tensor data type") {
if (arg == nullptr) {
*result_listener << "Tensor is null";
return false;
}
if (arg->type != tfl_type) {
*result_listener << "Data type is " << arg->type << " expected "
<< tfl_type;
return false;
}
std::optional<size_t> type_size = TfLiteTypeSizeOf(tfl_type);
if (!type_size.has_value()) {
*result_listener << "Expected type size is unknown.";
return false;
}
const size_t expected_bytes = type_size.value() * NumElements(arg);
if (arg->bytes != expected_bytes) {
*result_listener << "Allocated size is " << arg->bytes << " expected "
<< expected_bytes;
return false;
}
return true;
}
// Matches tensors on having data filled with given value.
MATCHER_P(FilledWith, value,
std::string(negation ? "is not" : "is") + " filled with " +
std::to_string(value)) {
if (arg == nullptr) {
*result_listener << "Tensor is null";
return false;
}
const value_type* const element_data =
reinterpret_cast<value_type*>(arg->data.data);
if (element_data == nullptr) {
*result_listener << "Tensor data is null";
return false;
}
return ::testing::Value(
std::vector<value_type>(element_data, element_data + NumElements(arg)),
::testing::Each(value));
}
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_VARIANTS_LIST_KERNELS_TEST_UTIL_H_
tensorflow/lite/kernels/variants/list_kernels/variant_add_n.cc
+257
View File
@@ -0,0 +1,257 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <algorithm>
#include <cstring>
#include <utility>
#include <vector>
#include "tensorflow/lite/array.h"
#include "tensorflow/lite/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/cpu_backend_context.h"
#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
#include "tensorflow/lite/kernels/internal/portable_tensor.h"
#include "tensorflow/lite/kernels/internal/runtime_shape.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_util.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/util.h"
namespace tflite {
namespace variants {
namespace ops {
namespace add_n {
namespace {
using ::tflite::variants::TensorArray;
constexpr int kInputTensor1 = 0;
constexpr int kOutputTensor = 0;
struct OpData {
// The index of the temporary tensor where temporary accumulations are kept.
int scratch_tensor_index;
};
void* Init(TfLiteContext* context, const char* /*buffer*/, size_t /*length*/) {
auto* op_data = new OpData();
if (context->AddTensors(context, 1, &op_data->scratch_tensor_index) !=
kTfLiteOk) {
delete op_data;
return nullptr;
}
return op_data;
}
void Free(TfLiteContext* /*context*/, void* buffer) {
delete reinterpret_cast<OpData*>(buffer);
}
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
TF_LITE_ENSURE(context, NumInputs(node) >= 2);
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
TfLiteTensor* output;
TF_LITE_ENSURE_OK(context,
GetOutputSafe(context, node, kOutputTensor, &output));
OpData* op_data = reinterpret_cast<OpData*>(node->user_data);
TfLiteIntArrayFree(node->temporaries);
node->temporaries = TfLiteIntArrayCreate(1);
TF_LITE_ENSURE(context, node->temporaries != nullptr);
node->temporaries->data[0] = op_data->scratch_tensor_index;
TfLiteTensor* scratch_tensor;
TF_LITE_ENSURE_OK(
context, GetTemporarySafe(context, node, /*index=*/0, &scratch_tensor));
scratch_tensor->type = kTfLiteNoType;
scratch_tensor->allocation_type = kTfLiteDynamic;
for (int i = kInputTensor1; i < NumInputs(node); ++i) {
const TfLiteTensor* input;
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, i, &input));
TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteVariant);
}
output->type = kTfLiteVariant;
output->allocation_type = kTfLiteVariantObject;
return kTfLiteOk;
}
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
///
// Prepare Inputs/Outputs
// ---------------------------------------------------------------------------
// Fetch backend cpu context.
CpuBackendContext* cpu_backend_context =
CpuBackendContext::GetFromContext(context);
// Parse Tensors.
const TfLiteTensor* input1;
TF_LITE_ENSURE_OK(context,
GetInputSafe(context, node, kInputTensor1, &input1));
TfLiteTensor* output;
TF_LITE_ENSURE_OK(context,
GetOutputSafe(context, node, kOutputTensor, &output));
TfLiteTensor* scratch_tensor;
TF_LITE_ENSURE_OK(context,
GetTemporarySafe(context, node, 0, &scratch_tensor));
// Collect vector of input list pointers, validate num_elements and type,
// compute merged shape.
TF_LITE_ENSURE_EQ(context, input1->allocation_type, kTfLiteVariantObject);
TF_LITE_ENSURE(context, input1->data.data != nullptr);
const TensorArray* const arr = static_cast<const TensorArray*>(
static_cast<const VariantData*>(input1->data.data));
TF_LITE_ENSURE(context, arr != nullptr);
const int num_elements = arr->NumElements();
const TfLiteType t = arr->ElementType();
const int num_inputs = NumInputs(node);
TF_LITE_ENSURE(context, arr->ElementShape() != nullptr);
IntArrayUniquePtr merged_shape = BuildTfLiteArray(*arr->ElementShape());
TF_LITE_ENSURE(context, merged_shape != nullptr);
std::vector<const TensorArray*> input_arrs;
input_arrs.reserve(num_inputs);
input_arrs.push_back(arr);
for (int i = 1; i < num_inputs; ++i) {
const TfLiteTensor* input;
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, i, &input));
TF_LITE_ENSURE_EQ(context, input->allocation_type, kTfLiteVariantObject);
TF_LITE_ENSURE(context, input->data.data != nullptr);
const TensorArray* const arr_i = static_cast<const TensorArray*>(
static_cast<const VariantData*>(input->data.data));
TF_LITE_ENSURE(context, arr_i != nullptr);
TF_LITE_ENSURE_EQ(context, num_elements, arr_i->NumElements());
TF_LITE_ENSURE_TYPES_EQ(context, t, arr_i->ElementType());
TF_LITE_ENSURE(context, arr_i->ElementShape() != nullptr);
IntArrayUniquePtr cur_shape = BuildTfLiteArray(*arr_i->ElementShape());
TF_LITE_ENSURE(context, cur_shape != nullptr);
merged_shape = variants::MergeShapesOrNull(std::move(merged_shape),
std::move(cur_shape));
TF_LITE_ENSURE(context, merged_shape != nullptr);
input_arrs.push_back(arr_i);
}
// Allocate output list with same length and type as all inputs.
IntArrayUniquePtr copied_merged_shape = BuildTfLiteArray(*merged_shape.get());
TF_LITE_ENSURE(context, copied_merged_shape != nullptr);
TF_LITE_ENSURE_OK(context, TfLiteTensorVariantRealloc<TensorArray>(
output, t, std::move(copied_merged_shape)));
TensorArray* const output_arr =
static_cast<TensorArray*>(static_cast<VariantData*>(output->data.data));
TF_LITE_ENSURE(context, output_arr->Resize(num_elements));
///
// Compute out_list[i] = Sum(in_list_0[i] + ...) for 0 < i < num_elements
// ---------------------------------------------------------------------------
for (int i = 0; i < num_elements; ++i) {
TfLiteIntArray* row_shape = nullptr;
std::vector<TfLiteTensor*> row_tensors;
// Compute resolved shape for this row. Save tensor* for all present
// elements at list_j`i`.
for (const TensorArray* array : input_arrs) {
const TfLiteTensor* at = array->At(i);
if (!at) continue;
if (!row_shape) {
row_shape = at->dims;
} else {
TF_LITE_ENSURE(context, TfLiteIntArrayEqual(row_shape, at->dims));
}
// We need const_cast to comply with optimized_ops::AddN type.
row_tensors.push_back(const_cast<TfLiteTensor*>(at));
}
if (row_shape == nullptr) {
// There exists no set elements in any list at position `i`. We can
// set output[i] to be a zeroed tensor with shape from input
// lists `ElementShape`.
TF_LITE_ENSURE(context,
variants::IsShapeFullyDefined(*merged_shape.get()));
IntArrayUniquePtr cur_shape = BuildTfLiteArray(*merged_shape.get());
TF_LITE_ENSURE(context, cur_shape != nullptr);
TensorUniquePtr row_output =
BuildTfLiteTensor(t, std::move(cur_shape), kTfLiteDynamic);
TF_LITE_ENSURE(context, row_output != nullptr);
if (row_output->bytes > 0) {
TF_LITE_ENSURE(context, row_output->data.data != nullptr);
memset(row_output->data.data, 0, row_output->bytes);
}
TF_LITE_ENSURE(context, output_arr->Set(i, std::move(row_output)));
continue;
}
// Allocate tensor for the sum of this row.
IntArrayUniquePtr cur_shape = BuildTfLiteArray(*row_shape);
TF_LITE_ENSURE(context, cur_shape != nullptr);
TensorUniquePtr row_output =
BuildTfLiteTensor(t, std::move(cur_shape), kTfLiteDynamic);
TF_LITE_ENSURE(context, row_output != nullptr);
if (row_output->bytes > 0) {
TF_LITE_ENSURE(context, row_output->data.data != nullptr);
}
if (row_tensors.size() < 2) {
// There is only one set item in all `input_j[i]`, so just use that.
TF_LITE_ENSURE_OK(context,
TfLiteTensorCopy(row_tensors[0], row_output.get()));
TF_LITE_ENSURE(context, output_arr->Set(i, std::move(row_output)));
continue;
}
// Resize scratch tensor so it can be used for each row.
const int num_inputs_for_row = static_cast<int>(row_tensors.size());
const int thread_count =
std::min(std::max(1, static_cast<int>(num_inputs_for_row) / 2),
cpu_backend_context->max_num_threads());
IntArrayUniquePtr scratch_shape = BuildTfLiteArray(
{thread_count * static_cast<int>(NumElements(row_tensors[0]))});
TF_LITE_ENSURE(context, scratch_shape != nullptr);
scratch_tensor->type = t;
TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scratch_tensor,
scratch_shape.release()));
const RuntimeShape row_runtime_shape(row_shape->size, row_shape->data);
// Compute sum of row.
if (t == kTfLiteInt32) {
VectorOfTensors<int> tensors(row_tensors);
optimized_ops::AddN<int>(
row_runtime_shape, num_inputs_for_row, tensors.data(),
GetTensorData<int>(row_output.get()),
GetTensorData<int>(scratch_tensor), cpu_backend_context);
} else if (t == kTfLiteFloat32) {
VectorOfTensors<float> tensors(row_tensors);
optimized_ops::AddN<float>(
row_runtime_shape, num_inputs_for_row, tensors.data(),
GetTensorData<float>(row_output.get()),
GetTensorData<float>(scratch_tensor), cpu_backend_context);
} else {
TF_LITE_KERNEL_LOG(context, "Subtype is not supported for variant addn.");
return kTfLiteError;
}
// Write row sum to appropriate element in output list.
TF_LITE_ENSURE(context, output_arr->Set(i, std::move(row_output)));
}
return kTfLiteOk;
}
} // namespace
} // namespace add_n
TfLiteRegistration* Register_VARIANT_ADD_N() {
static TfLiteRegistration r = {add_n::Init, add_n::Free, add_n::Prepare,
add_n::Eval};
return &r;
}
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/variant_add_n_test.cc
+207
View File
@@ -0,0 +1,207 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <optional>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/internal/compatibility.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/portable_type_to_tflitetype.h"
#include "tensorflow/lite/schema/schema_generated.h"
namespace tflite {
namespace variants {
namespace ops {
namespace {
using ::testing::AllOf;
class ListAddNModel : public ListOpModel {
public:
explicit ListAddNModel(int num_inputs) {
std::vector<std::vector<int>> input_shapes(num_inputs, std::vector<int>{});
for (int i = 0; i < num_inputs; ++i) {
input_inds_.push_back(AddInput({TensorType_VARIANT, {}}));
}
output_ind_ = AddOutput({TensorType_VARIANT, {}});
SetCustomOp("VariantAddN", {}, Register_VARIANT_ADD_N);
BuildInterpreter(input_shapes);
}
const TensorArray* GetOutput() {
TfLiteTensor* tensor = interpreter_->tensor(output_ind_);
TFLITE_CHECK(tensor != nullptr && tensor->type == kTfLiteVariant &&
tensor->allocation_type == kTfLiteVariantObject);
return static_cast<const TensorArray*>(
static_cast<const VariantData*>(tensor->data.data));
}
int GetIndOfInput(int input) { return input_inds_[input]; }
private:
std::vector<int> input_inds_;
int output_ind_;
};
template <typename T>
class ListAddNTest : public ::testing::Test {};
TYPED_TEST_SUITE_P(ListAddNTest);
TYPED_TEST_P(ListAddNTest, TwoInputs_AllItemsPresent_AllSameShape) {
TfLiteType tfl_type = typeToTfLiteType<TypeParam>();
std::optional<TensorType> tensor_type = TflToTensorType(tfl_type);
ASSERT_TRUE(tensor_type.has_value());
ListAddNModel m(2);
for (const int i : {0, 1}) {
const int list_ind = m.GetIndOfInput(i);
m.PopulateListTensor(list_ind, {}, 2, tfl_type);
m.ListSetItem(list_ind, 0, {2, 2}, tfl_type,
std::vector<TypeParam>(4, 1).data());
m.ListSetItem(list_ind, 1, {2, 2}, tfl_type,
std::vector<TypeParam>(4, 1).data());
}
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* const arr = m.GetOutput();
ASSERT_EQ(arr->NumElements(), 2);
ASSERT_EQ(arr->ElementType(), tfl_type);
EXPECT_THAT(arr->At(0), AllOf(IsAllocatedAs(tfl_type), DimsAre({2, 2}),
FilledWith<TypeParam>(2)));
EXPECT_THAT(arr->At(1), AllOf(IsAllocatedAs(tfl_type), DimsAre({2, 2}),
FilledWith<TypeParam>(2)));
}
TYPED_TEST_P(ListAddNTest, TwoInputs_AllItemsPresent_ListsContainMixedShapes) {
TfLiteType tfl_type = typeToTfLiteType<TypeParam>();
std::optional<TensorType> tensor_type = TflToTensorType(tfl_type);
ASSERT_TRUE(tensor_type.has_value());
ListAddNModel m(2);
for (const int i : {0, 1}) {
const int list_ind = m.GetIndOfInput(i);
m.PopulateListTensor(list_ind, {}, 2, tfl_type);
m.ListSetItem(list_ind, 0, {2, 2}, tfl_type,
std::vector<TypeParam>(4, 1).data());
m.ListSetItem(list_ind, 1, {3, 3}, tfl_type,
std::vector<TypeParam>(9, 1).data());
}
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* const arr = m.GetOutput();
ASSERT_EQ(arr->NumElements(), 2);
ASSERT_EQ(arr->ElementType(), tfl_type);
EXPECT_THAT(arr->At(0), AllOf(IsAllocatedAs(tfl_type), DimsAre({2, 2}),
FilledWith<TypeParam>(2)));
EXPECT_THAT(arr->At(1), AllOf(IsAllocatedAs(tfl_type), DimsAre({3, 3}),
FilledWith<TypeParam>(2)));
}
TYPED_TEST_P(ListAddNTest, TwoInputs_NoItemsPresent_ListShapesMerge) {
TfLiteType tfl_type = typeToTfLiteType<TypeParam>();
std::optional<TensorType> tensor_type = TflToTensorType(tfl_type);
ASSERT_TRUE(tensor_type.has_value());
ListAddNModel m(2);
m.PopulateListTensor(m.GetIndOfInput(0), {2, -1}, 1, tfl_type);
m.PopulateListTensor(m.GetIndOfInput(1), {-1, 2}, 1, tfl_type);
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* const arr = m.GetOutput();
ASSERT_EQ(arr->NumElements(), 1);
ASSERT_EQ(arr->ElementType(), tfl_type);
EXPECT_THAT(arr->At(0), AllOf(IsAllocatedAs(tfl_type), DimsAre({2, 2}),
FilledWith<TypeParam>(0)));
}
TYPED_TEST_P(ListAddNTest, TwoInputs_NoItemsPresent_ListShapesUndefinedError) {
TfLiteType tfl_type = typeToTfLiteType<TypeParam>();
std::optional<TensorType> tensor_type = TflToTensorType(tfl_type);
ASSERT_TRUE(tensor_type.has_value());
ListAddNModel m(2);
m.PopulateListTensor(m.GetIndOfInput(0), {2, -1}, 1, tfl_type);
m.PopulateListTensor(m.GetIndOfInput(1), {-1, -1}, 1, tfl_type);
ASSERT_EQ(m.Invoke(), kTfLiteError);
}
TYPED_TEST_P(ListAddNTest, TwoInputs_SomeItemsPresent_UsesElementShape) {
TfLiteType tfl_type = typeToTfLiteType<TypeParam>();
std::optional<TensorType> tensor_type = TflToTensorType(tfl_type);
ASSERT_TRUE(tensor_type.has_value());
ListAddNModel m(2);
m.PopulateListTensor(m.GetIndOfInput(0), {}, 1, tfl_type);
m.PopulateListTensor(m.GetIndOfInput(1), {}, 1, tfl_type);
m.ListSetItem(m.GetIndOfInput(0), 0, {3, 3}, tfl_type,
std::vector<TypeParam>(9, 1).data());
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* const arr = m.GetOutput();
ASSERT_EQ(arr->NumElements(), 1);
ASSERT_EQ(arr->ElementType(), tfl_type);
EXPECT_THAT(arr->At(0), AllOf(IsAllocatedAs(tfl_type), DimsAre({3, 3}),
FilledWith<TypeParam>(1)));
}
TYPED_TEST_P(ListAddNTest, MultipleInputsWithThreads) {
TfLiteType tfl_type = typeToTfLiteType<TypeParam>();
std::optional<TensorType> tensor_type = TflToTensorType(tfl_type);
ASSERT_TRUE(tensor_type.has_value());
ListAddNModel m(4);
m.SetNumThreads(4);
for (int i = 0; i < 4; ++i) {
const int list_ind = m.GetIndOfInput(i);
m.PopulateListTensor(list_ind, {}, 2, tfl_type);
m.ListSetItem(list_ind, 0, {2, 2}, tfl_type,
std::vector<TypeParam>(4, 1).data());
m.ListSetItem(list_ind, 1, {2, 2}, tfl_type,
std::vector<TypeParam>(4, 1).data());
}
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* const arr = m.GetOutput();
ASSERT_EQ(arr->NumElements(), 2);
ASSERT_EQ(arr->ElementType(), tfl_type);
EXPECT_THAT(arr->At(0), AllOf(IsAllocatedAs(tfl_type), DimsAre({2, 2}),
FilledWith<TypeParam>(4)));
EXPECT_THAT(arr->At(1), AllOf(IsAllocatedAs(tfl_type), DimsAre({2, 2}),
FilledWith<TypeParam>(4)));
}
REGISTER_TYPED_TEST_SUITE_P(ListAddNTest,
TwoInputs_AllItemsPresent_AllSameShape,
TwoInputs_AllItemsPresent_ListsContainMixedShapes,
TwoInputs_NoItemsPresent_ListShapesMerge,
TwoInputs_SomeItemsPresent_UsesElementShape,
TwoInputs_NoItemsPresent_ListShapesUndefinedError,
MultipleInputsWithThreads);
using ValidTypes = ::testing::Types<int, float>;
INSTANTIATE_TYPED_TEST_SUITE_P(ListAddNTests, ListAddNTest, ValidTypes);
} // namespace
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_kernels/variant_zeros_like_test.cc
+132
View File
@@ -0,0 +1,132 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <tuple>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/internal/compatibility.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/schema/schema_generated.h"
namespace tflite {
namespace variants {
namespace ops {
namespace {
using ::testing::AllOf;
using ::testing::Combine;
using ::testing::ValuesIn;
using ::tflite::variants::TensorArray;
class VariantZerosLikeModel : public ListOpModel {
public:
explicit VariantZerosLikeModel() {
list_input_ = AddInput({TensorType_VARIANT, {}});
list_output_ = AddOutput({TensorType_VARIANT, {}});
SetCustomOp("VariantZerosLike", {}, Register_VARIANT_ZEROS_LIKE);
BuildInterpreter({{}});
}
const TensorArray* GetOutputTensorArray() {
TfLiteTensor* tensor = interpreter_->tensor(list_output_);
TFLITE_CHECK(tensor != nullptr && tensor->type == kTfLiteVariant &&
tensor->allocation_type == kTfLiteVariantObject);
return static_cast<const TensorArray*>(
static_cast<const VariantData*>(tensor->data.data));
}
int list_input_;
int list_output_;
};
using VariantZerosLikeTestParam = std::tuple<std::vector<int>, TfLiteType, int>;
class VariantZerosLikeTest
: public testing::TestWithParam<VariantZerosLikeTestParam> {
public:
enum { kShape, kType, kLen };
};
TEST_P(VariantZerosLikeTest, OutputsEmptyListWithSameAttrs) {
const auto& param = GetParam();
const std::vector<int>& shape = std::get<kShape>(param);
const TfLiteType t = std::get<kType>(param);
const int len = std::get<kLen>(param);
VariantZerosLikeModel m;
m.PopulateListTensor(m.list_input_, shape, len, t);
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* const out = m.GetOutputTensorArray();
ASSERT_EQ(out->NumElements(), len);
ASSERT_EQ(out->ElementType(), t);
ASSERT_THAT(out->ElementShape(), DimsAre(shape));
for (int i = 0; i < len; ++i) {
EXPECT_EQ(out->At(i), nullptr);
}
}
using VariantZerosLikeItemTestParam = std::tuple<int, std::vector<int>>;
class VariantZerosLikeItemTest
: public testing::TestWithParam<VariantZerosLikeItemTestParam> {
public:
enum { kLen, kShape };
};
TEST_P(VariantZerosLikeItemTest, OutputsEmptyListContainsZeroedElement) {
const auto& param = GetParam();
const int len = std::get<kLen>(param);
const std::vector<int>& item_shape = std::get<kShape>(param);
VariantZerosLikeModel m;
m.PopulateListTensor(m.list_input_, {}, len, kTfLiteInt32);
const int num_elements = NumElements(item_shape.data(), item_shape.size());
m.ListSetItem(m.list_input_, 0, item_shape, kTfLiteInt32,
std::vector<int>(num_elements, 1).data());
ASSERT_EQ(m.Invoke(), kTfLiteOk);
const TensorArray* const out = m.GetOutputTensorArray();
ASSERT_EQ(out->NumElements(), len);
ASSERT_EQ(out->ElementType(), kTfLiteInt32);
ASSERT_THAT(out->ElementShape(), DimsAre({}));
const TfLiteTensor* const zero = out->At(0);
ASSERT_NE(zero, nullptr);
EXPECT_THAT(zero, AllOf(DimsAre(item_shape), IsAllocatedAs(kTfLiteInt32),
FilledWith<int>(0)));
for (int i = 1; i < len; ++i) {
EXPECT_EQ(out->At(i), nullptr);
}
}
INSTANTIATE_TEST_SUITE_P(VariantZerosLikeTests, VariantZerosLikeTest,
Combine(ValuesIn(std::vector<std::vector<int>>{
{}, {-1}, {2, 2}, {3, 3, 3}}),
ValuesIn({kTfLiteInt32, kTfLiteInt64,
kTfLiteFloat32, kTfLiteBool}),
ValuesIn({0, 2, 10})));
INSTANTIATE_TEST_SUITE_P(VariantZerosLikeTests, VariantZerosLikeItemTest,
Combine(ValuesIn({1, 2, 10}),
ValuesIn(std::vector<std::vector<int>>{
{1}, {2, 2}, {3, 3, 3}})));
} // namespace
} // namespace ops
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_ops_lib.h
+58
View File
@@ -0,0 +1,58 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_VARIANTS_LIST_OPS_LIB_H_
#define TENSORFLOW_LITE_KERNELS_VARIANTS_LIST_OPS_LIB_H_
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/schema/schema_generated.h"
namespace tflite {
namespace variants {
namespace detail {
struct ListReserveOptions {
TensorType element_type;
};
} // namespace detail
namespace ops {
TfLiteRegistration* Register_LIST_RESERVE();
TfLiteRegistration* Register_LIST_STACK();
TfLiteRegistration* Register_LIST_SET_ITEM();
TfLiteRegistration* Register_LIST_FROM_TENSOR();
TfLiteRegistration* Register_LIST_GET_ITEM();
TfLiteRegistration* Register_LIST_LENGTH();
TfLiteRegistration* Register_LIST_ELEMENT_SHAPE();
TfLiteRegistration* Register_LIST_POP_BACK();
TfLiteRegistration* Register_LIST_PUSH_BACK();
TfLiteRegistration* Register_VARIANT_ADD_N();
TfLiteRegistration* Register_VARIANT_ZEROS_LIKE();
} // namespace ops
} // namespace variants
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_VARIANTS_LIST_OPS_LIB_H_
tensorflow/lite/kernels/variants/list_ops_subgraph_test.cc
+329
View File
@@ -0,0 +1,329 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <algorithm>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/types/span.h"
#include "tensorflow/lite/array.h"
#include "tensorflow/lite/c/c_api_types.h"
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/interpreter_test_util.h"
#include "tensorflow/lite/kernels/internal/compatibility.h"
#include "tensorflow/lite/kernels/kernel_util.h"
#include "tensorflow/lite/kernels/op_macros.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_subgraph_test_util.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/schema/schema_generated.h"
using ::testing::ElementsAreArray;
using ::tflite::variants::TensorArray;
namespace tflite {
namespace {
using UtilsTest = ListOpsSubgraphTest;
// This test just validates the test fixture. It doesn't test any business
// logic.
TEST_F(UtilsTest, SimpleAddConst) {
builder_.BuildAddConstSubgraph(&interpreter_.primary_subgraph());
TfLiteTensor* cst1 = interpreter_.tensor(0);
ASSERT_THAT(cst1, DimsAre({2}));
EXPECT_EQ(cst1->data.i32[0], 2);
EXPECT_EQ(cst1->data.i32[1], 2);
TfLiteTensor* cst2 = interpreter_.tensor(1);
ASSERT_THAT(cst2, DimsAre({2}));
EXPECT_EQ(cst2->data.i32[0], 3);
EXPECT_EQ(cst2->data.i32[1], 3);
ASSERT_EQ(interpreter_.AllocateTensors(), kTfLiteOk);
ASSERT_EQ(interpreter_.Invoke(), kTfLiteOk);
TfLiteTensor* out = interpreter_.tensor(2);
ASSERT_THAT(out, DimsAre({2}));
EXPECT_EQ(out->data.i32[0], 5);
EXPECT_EQ(out->data.i32[1], 5);
}
struct ListReserveSubgraphTestParams {
const TensorType tensor_type;
const TfLiteType expected_type;
const std::vector<int> element_shape_shape;
const std::vector<int> element_shape_data;
const std::vector<int> expected_element_shape;
const int num_elements;
};
class ListReserveSubgraphTest
: public ListOpsSubgraphTest,
public ::testing::WithParamInterface<ListReserveSubgraphTestParams> {};
TEST_P(ListReserveSubgraphTest, InterpreterOutputsTensorArray) {
const ListReserveSubgraphTestParams& params = GetParam();
builder_.BuildReserveSubgraph(&interpreter_.primary_subgraph(),
params.tensor_type);
ASSERT_EQ(interpreter_.ResizeInputTensor(0, params.element_shape_shape),
kTfLiteOk);
ASSERT_EQ(interpreter_.ResizeInputTensor(1, {}), kTfLiteOk);
ASSERT_EQ(interpreter_.AllocateTensors(), kTfLiteOk);
TfLiteTensor* element_shape = interpreter_.input_tensor(0);
std::copy(params.element_shape_data.begin(), params.element_shape_data.end(),
element_shape->data.i32);
TfLiteTensor* num_elements = interpreter_.input_tensor(1);
num_elements->data.i32[0] = params.num_elements;
ASSERT_EQ(interpreter_.Invoke(), kTfLiteOk);
TfLiteTensor* output = interpreter_.output_tensor(0);
ASSERT_EQ(output->type, kTfLiteVariant);
ASSERT_EQ(output->allocation_type, kTfLiteVariantObject);
ASSERT_TRUE(output->data.data != nullptr);
TensorArray* result =
static_cast<TensorArray*>(static_cast<VariantData*>(output->data.data));
EXPECT_EQ(result->NumElements(), params.num_elements);
EXPECT_THAT(result->ElementShape(), DimsAre(params.expected_element_shape));
EXPECT_EQ(result->ElementType(), params.expected_type);
for (int i = 0; i < params.num_elements; ++i) {
EXPECT_EQ(result->At(i), nullptr);
}
}
INSTANTIATE_TEST_SUITE_P(
ListOpsSubgraphParamTests, ListReserveSubgraphTest,
testing::ValuesIn({
ListReserveSubgraphTestParams{
TensorType_INT32, kTfLiteInt32, {}, {-1}, {}, 2},
ListReserveSubgraphTestParams{
TensorType_FLOAT32, kTfLiteFloat32, {}, {-1}, {}, 2},
ListReserveSubgraphTestParams{
TensorType_FLOAT32, kTfLiteFloat32, {1}, {-1}, {-1}, 2},
ListReserveSubgraphTestParams{
TensorType_FLOAT32, kTfLiteFloat32, {2}, {2, 2}, {2, 2}, 0},
ListReserveSubgraphTestParams{
TensorType_FLOAT32, kTfLiteFloat32, {2}, {2, -1}, {2, -1}, 10},
}));
struct ListStackSubgraphDynamicTestParams {
// Reserve params.
const std::vector<int> element_shape_shape;
const std::vector<int> element_shape_data;
const int num_elements;
// Stack params.
const std::vector<int> stack_shape_shape;
const std::vector<int> stack_shape_data;
// Expected.
const std::vector<int> expected_shape;
};
class ListStackDynamicSubgraphTest
: public ListOpsSubgraphTest,
public ::testing::WithParamInterface<ListStackSubgraphDynamicTestParams> {
};
TEST_P(ListStackDynamicSubgraphTest,
InterpreterOutputsStackTensor_DynamicOutput) {
const ListStackSubgraphDynamicTestParams& params = GetParam();
builder_.BuildReserveStackSubgraph(&interpreter_.primary_subgraph());
ASSERT_EQ(interpreter_.ResizeInputTensor(0, params.element_shape_shape),
kTfLiteOk);
ASSERT_EQ(interpreter_.ResizeInputTensor(1, {}), kTfLiteOk);
ASSERT_EQ(interpreter_.ResizeInputTensor(2, params.stack_shape_shape),
kTfLiteOk);
interpreter_.output_tensor(0)->allocation_type = kTfLiteDynamic;
ASSERT_EQ(interpreter_.AllocateTensors(), kTfLiteOk);
TfLiteTensor* element_shape = interpreter_.input_tensor(0);
std::copy(params.element_shape_data.begin(), params.element_shape_data.end(),
element_shape->data.i32);
TfLiteTensor* num_elements = interpreter_.input_tensor(1);
num_elements->data.i32[0] = params.num_elements;
TfLiteTensor* stack_shape = interpreter_.input_tensor(2);
std::copy(params.stack_shape_data.begin(), params.stack_shape_data.end(),
stack_shape->data.i32);
ASSERT_EQ(interpreter_.Invoke(), kTfLiteOk);
TfLiteTensor* output = interpreter_.output_tensor(0);
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_EQ(output->allocation_type, kTfLiteDynamic);
const int output_num_elements = NumElements(output);
ASSERT_TRUE((output_num_elements > 0 && output->data.data != nullptr) ||
(output_num_elements == 0 && output->data.data == nullptr));
ASSERT_THAT(output, DimsAre(params.expected_shape));
for (int i = 0; i < NumElements(output); ++i) {
EXPECT_EQ(output->data.i32[i], 0);
}
}
INSTANTIATE_TEST_SUITE_P(
ListOpsSubgraphParamTests, ListStackDynamicSubgraphTest,
testing::ValuesIn({
ListStackSubgraphDynamicTestParams{{1}, {2}, 4, {}, {-1}, {4, 2}},
ListStackSubgraphDynamicTestParams{
{}, {-1}, 4, {3}, {2, 3, 4}, {4, 2, 3, 4}},
ListStackSubgraphDynamicTestParams{{1}, {2}, 4, {}, {-1}, {4, 2}},
ListStackSubgraphDynamicTestParams{{1}, {2}, 0, {}, {-1}, {0, 2}},
ListStackSubgraphDynamicTestParams{{1}, {1}, 2, {}, {-1}, {2}},
}));
// Fixture that constructs a model that uses a "While" op to
// populate each element in a `TensorArray` with a constant tensor.
// See documentation of `BuildLessThanSubgraph`,
// `BuildSetItemAndIncrementSubgraph` and `BuildWhileSubgraph` for more
// detail.
class WhileIncrementListOpsTest : public InterpreterTest {
public:
WhileIncrementListOpsTest() {
AddSubgraphs(2);
builder_.BuildLessThanSubgraph(interpreter_->subgraph(1));
builder_.BuildSetItemAndIncrementSubgraph(interpreter_->subgraph(2));
builder_.BuildWhileSubgraph(&interpreter_->primary_subgraph());
TFLITE_CHECK(interpreter_->AllocateTensors() == kTfLiteOk);
}
protected:
// Allocates a `TensorArray` behind the `kTfLiteVariant` tensor at given
// index.
void PopulateListTensor(int index, absl::Span<const int> element_shape_data,
int num_elements, TfLiteType element_type) {
TfLiteTensor* tensor = interpreter_->tensor(index);
TF_LITE_ASSERT_EQ(tensor->type, kTfLiteVariant);
tensor->allocation_type = kTfLiteVariantObject;
tensor->buffer_handle = kTfLiteNullBufferHandle;
tensor->quantization = {kTfLiteNoQuantization};
IntArrayUniquePtr element_shape =
BuildTfLiteArray(element_shape_data.size(), element_shape_data.data());
TfLiteStatus stat = TfLiteTensorVariantRealloc<TensorArray>(
tensor, element_type, std::move(element_shape));
TF_LITE_ASSERT_EQ(stat, kTfLiteOk);
TensorArray* arr =
static_cast<TensorArray*>(static_cast<VariantData*>(tensor->data.data));
arr->Resize(num_elements);
}
// Retreives a pointer to the `TensorArray` sitting behind the
// `kTfLiteVariant` tensor at given index.
const TensorArray* GetOutputTensorArray(int tensor_id) {
TfLiteTensor* tensor = interpreter_->tensor(tensor_id);
TFLITE_CHECK(tensor != nullptr && tensor->type == kTfLiteVariant &&
tensor->allocation_type == kTfLiteVariantObject);
return static_cast<const TensorArray*>(
static_cast<const VariantData*>(tensor->data.data));
}
ListOpsSubgraphBuilder builder_;
};
TEST_F(WhileIncrementListOpsTest, PopulateListWithWhile) {
interpreter_->tensor(interpreter_->inputs()[0])->data.i32[0] = 0;
PopulateListTensor(interpreter_->inputs()[1], {2, 2}, 3, kTfLiteInt32);
ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
const TensorArray* output = GetOutputTensorArray(interpreter_->outputs()[1]);
ASSERT_TRUE(output != nullptr);
ASSERT_EQ(output->NumElements(), 3);
for (int i = 0; i < 3; ++i) {
const TfLiteTensor* item = output->At(i);
ASSERT_TRUE(item != nullptr);
ASSERT_THAT(item, DimsAre({2}));
EXPECT_THAT(std::vector<int>(item->data.i32, item->data.i32 + 2),
ElementsAreArray({2, 2}));
}
}
TEST_F(WhileIncrementListOpsTest,
PartiallyPopulateListWithWhile_UnsetItemsZeroed) {
interpreter_->tensor(interpreter_->inputs()[0])->data.i32[0] = 1;
PopulateListTensor(interpreter_->inputs()[1], {2, 2}, 3, kTfLiteInt32);
ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
const TensorArray* output = GetOutputTensorArray(interpreter_->outputs()[1]);
ASSERT_TRUE(output != nullptr);
ASSERT_EQ(output->NumElements(), 3);
ASSERT_EQ(output->At(0), nullptr);
for (int i = 1; i < 3; ++i) {
const TfLiteTensor* item = output->At(i);
ASSERT_TRUE(item != nullptr);
ASSERT_THAT(item, DimsAre({2}));
EXPECT_THAT(std::vector<int>(item->data.i32, item->data.i32 + 2),
ElementsAreArray({2, 2}));
}
}
class ListReserveLengthSubgraphTest
: public ListOpsSubgraphTest,
public ::testing::WithParamInterface<int> {};
TEST_P(ListReserveLengthSubgraphTest, InterpreterOutputsListLength) {
const int length = GetParam();
builder_.BuildReserveLengthSubgraph(&interpreter_.primary_subgraph());
ASSERT_EQ(interpreter_.ResizeInputTensor(0, {1}), kTfLiteOk);
ASSERT_EQ(interpreter_.ResizeInputTensor(1, {}), kTfLiteOk);
ASSERT_EQ(interpreter_.AllocateTensors(), kTfLiteOk);
TfLiteTensor* element_shape = interpreter_.input_tensor(0);
element_shape->data.i32[0] = 2;
TfLiteTensor* num_elements = interpreter_.input_tensor(1);
num_elements->data.i32[0] = length;
ASSERT_EQ(interpreter_.Invoke(), kTfLiteOk);
TfLiteTensor* output = interpreter_.output_tensor(0);
ASSERT_EQ(output->type, kTfLiteInt32);
ASSERT_EQ(output->allocation_type, kTfLiteArenaRw);
ASSERT_THAT(output, DimsAre({}));
ASSERT_TRUE(output->data.data != nullptr);
ASSERT_EQ(output->data.i32[0], length);
}
INSTANTIATE_TEST_SUITE_P(ListOpsSubgraphParamTests,
ListReserveLengthSubgraphTest,
testing::Values(0, 1, 2, 5, 10));
} // namespace
} // namespace tflite
tensorflow/lite/kernels/variants/list_ops_subgraph_test_util.cc
+408
View File
@@ -0,0 +1,408 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/lite/kernels/variants/list_ops_subgraph_test_util.h"
#include <algorithm>
#include <climits>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <numeric>
#include <vector>
#include <gtest/gtest.h>
#include "absl/types/span.h"
#include "tensorflow/lite/builtin_ops.h"
#include "tensorflow/lite/c/c_api_types.h"
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/core/c/builtin_op_data.h"
#include "tensorflow/lite/core/subgraph.h"
#include "tensorflow/lite/kernels/builtin_op_kernels.h"
#include "tensorflow/lite/kernels/op_macros.h"
#include "tensorflow/lite/kernels/subgraph_test_util.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/schema/schema_generated.h"
using ::tflite::subgraph_test_util::SetupTensor;
using ::tflite::variants::detail::ListReserveOptions;
using ::tflite::variants::ops::Register_LIST_LENGTH;
using ::tflite::variants::ops::Register_LIST_RESERVE;
using ::tflite::variants::ops::Register_LIST_SET_ITEM;
using ::tflite::variants::ops::Register_LIST_STACK;
namespace tflite {
ListReserveOptions* ListOpsSubgraphBuilder::RequestReserveOptions(
TensorType element_type) {
list_reserve_opts_.push_back(ListReserveOptions{element_type});
return &list_reserve_opts_.back();
}
void ListOpsSubgraphBuilder::CreateConstantInt32Tensor(
Subgraph* subgraph, int tensor_index, absl::Span<const int> shape,
absl::Span<const int> data) {
const bool all_static_dimensions =
std::all_of(shape.begin(), shape.end(), [](int i) { return i >= 0; });
TF_LITE_ASSERT(all_static_dimensions);
TF_LITE_ASSERT(!shape.empty());
// tflite only supports tensors with at most rank 5
const bool will_not_overflow =
std::all_of(shape.begin(), shape.end(),
[](int i) { return i < (INT_MAX / 5); }) &&
shape.size() <= 5;
TF_LITE_ASSERT(will_not_overflow);
const int num_elements =
std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<int>());
TF_LITE_ASSERT_EQ(num_elements, data.size());
size_t bytes = sizeof(int32_t) * num_elements;
int_buffers_.push_back(std::vector<int32_t>(data.begin(), data.end()));
TfLiteStatus stat = subgraph->SetTensorParametersReadOnly(
tensor_index, kTfLiteInt32, /*name=*/"",
std::vector<int>(shape.begin(), shape.end()), /*quantization=*/{},
reinterpret_cast<const char*>(int_buffers_.back().data()), bytes);
TF_LITE_ASSERT_EQ(stat, kTfLiteOk);
}
void ListOpsSubgraphBuilder::BuildAddConstSubgraph(Subgraph* subgraph) {
constexpr int kLHS = 0;
constexpr int kRHS = 1;
constexpr int kOut = 2;
constexpr int kTensorCount = 3;
// kLHS(0) --> +-----------+
// | ADD | --> kOut(2)
// kRHS(1) --> +-----------+
int first_new_tensor_index;
TF_LITE_ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
kTfLiteOk);
TF_LITE_ASSERT_EQ(first_new_tensor_index, 0);
TF_LITE_ASSERT_EQ(subgraph->SetOutputs({kOut}), kTfLiteOk);
CreateConstantInt32Tensor(subgraph, kLHS, {2}, {2, 2});
CreateConstantInt32Tensor(subgraph, kRHS, {2}, {3, 3});
SetupTensor(subgraph, kOut, kTfLiteInt32);
TfLiteAddParams* params =
reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
params->activation = kTfLiteActNone;
TfLiteRegistration* add_reg = ops::builtin::Register_ADD();
add_reg->builtin_code = kTfLiteBuiltinAdd;
int node_index;
TfLiteStatus stat = subgraph->AddNodeWithParameters(
{kLHS, kRHS}, {kOut}, /*intermediates=*/{}, /*init_data=*/nullptr,
/*init_data_size=*/0, params, add_reg, &node_index);
TF_LITE_ASSERT_EQ(stat, kTfLiteOk);
}
void ListOpsSubgraphBuilder::BuildReserveSubgraph(Subgraph* subgraph,
TensorType element_type) {
constexpr int kElementShape = 0;
constexpr int kNumElements = 1;
constexpr int kListOut = 2;
constexpr int kTensorCount = 3;
// kElementShape(0) --> +-------------------+
// | ListReserve | --> kListOut(2)
// kNumElements(1) --> +-------------------+
int first_new_tensor_index;
TF_LITE_ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
kTfLiteOk);
TF_LITE_ASSERT_EQ(first_new_tensor_index, 0);
TF_LITE_ASSERT_EQ(subgraph->SetOutputs({kListOut}), kTfLiteOk);
SetupTensor(subgraph, kListOut, kTfLiteVariant);
TF_LITE_ASSERT_EQ(subgraph->SetInputs({kElementShape, kNumElements}),
kTfLiteOk);
SetupTensor(subgraph, kElementShape, kTfLiteInt32);
SetupTensor(subgraph, kNumElements, kTfLiteInt32);
TfLiteRegistration* reserve_reg = Register_LIST_RESERVE();
reserve_reg->builtin_code = BuiltinOperator_CUSTOM;
reserve_reg->custom_name = "ListReserve";
ListReserveOptions* options = RequestReserveOptions(element_type);
int node_index;
TfLiteStatus stat = subgraph->AddNodeWithParameters(
{kElementShape, kNumElements}, {kListOut},
/*intermediates=*/{}, reinterpret_cast<const char*>(options),
sizeof(ListReserveOptions),
/*builtin_data=*/nullptr, reserve_reg, &node_index);
TF_LITE_ASSERT_EQ(stat, kTfLiteOk);
}
void ListOpsSubgraphBuilder::BuildReserveStackSubgraph(Subgraph* subgraph) {
constexpr int kElementShape = 0;
constexpr int kNumElements = 1;
constexpr int kStackShape = 2;
constexpr int kReserveOut = 3;
constexpr int kStackOut = 4;
constexpr int kTensorCount = 5;
// kElementShape(0) --> +-------------+
// | ListReserve |
// kNumElements(1) --> +-------------+ --> kReserveOut(2)
// |
// +-----------+
// | ListStack |
// kStackShape(3) --> +-----------+ --> kStackOut(4)
int first_new_tensor_index;
TF_LITE_ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
kTfLiteOk);
TF_LITE_ASSERT_EQ(first_new_tensor_index, 0);
TF_LITE_ASSERT_EQ(subgraph->SetOutputs({kStackOut}), kTfLiteOk);
SetupTensor(subgraph, kStackOut, kTfLiteInt32);
TF_LITE_ASSERT_EQ(
subgraph->SetInputs({kElementShape, kNumElements, kStackShape}),
kTfLiteOk);
SetupTensor(subgraph, kElementShape, kTfLiteInt32);
SetupTensor(subgraph, kNumElements, kTfLiteInt32);
SetupTensor(subgraph, kStackShape, kTfLiteInt32);
SetupTensor(subgraph, kReserveOut, kTfLiteVariant);
TfLiteRegistration* reserve_reg = Register_LIST_RESERVE();
reserve_reg->builtin_code = BuiltinOperator_CUSTOM;
reserve_reg->custom_name = "ListReserve";
ListReserveOptions* options = RequestReserveOptions(TensorType_INT32);
TfLiteStatus reg_stat = subgraph->AddNodeWithParameters(
{kElementShape, kNumElements}, {kReserveOut}, {},
reinterpret_cast<const char*>(options), sizeof(ListReserveOptions),
nullptr, reserve_reg);
TF_LITE_ASSERT_EQ(reg_stat, kTfLiteOk);
TfLiteRegistration* stack_reg = Register_LIST_STACK();
stack_reg->builtin_code = BuiltinOperator_CUSTOM;
stack_reg->custom_name = "ListStack";
TfLiteStatus stack_stat = subgraph->AddNodeWithParameters(
{kReserveOut, kStackShape}, {kStackOut}, {},
/*init_data=*/nullptr, /*init_data_size=*/0, /*builtin_data=*/nullptr,
stack_reg);
TF_LITE_ASSERT_EQ(stack_stat, kTfLiteOk);
}
void ListOpsSubgraphBuilder::BuildLessThanSubgraph(Subgraph* subgraph) {
const int kCurIn = 0;
const int kListIn = 1;
const int kBoolOut = 2;
const int kConstMax = 3;
const int kTensorCount = 4;
// kCurIn(0) --> +-----------+
// | LESS_THAN | --> kBoolOut(2)
// kListIn(1) unused --> +-----------+
int first_new_tensor_index;
ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
kTfLiteOk);
ASSERT_EQ(first_new_tensor_index, 0);
ASSERT_EQ(subgraph->SetInputs({kCurIn, kListIn}), kTfLiteOk);
ASSERT_EQ(subgraph->SetOutputs({kBoolOut}), kTfLiteOk);
SetupTensor(subgraph, kCurIn, kTfLiteInt32);
SetupTensor(subgraph, kListIn, kTfLiteVariant);
SetupTensor(subgraph, kBoolOut, kTfLiteBool);
CreateConstantInt32Tensor(subgraph, kConstMax, {1}, {3});
auto* less_reg = ops::builtin::Register_LESS();
less_reg->builtin_code = kTfLiteBuiltinLessEqual;
int node_index;
TfLiteStatus stat = subgraph->AddNodeWithParameters(
{kCurIn, kConstMax}, {kBoolOut}, /*intermediates=*/{},
/*init_data=*/nullptr, /*init_data_size=*/0, /*builtin_data=*/nullptr,
less_reg, &node_index);
TF_LITE_ASSERT_EQ(stat, kTfLiteOk);
}
void ListOpsSubgraphBuilder::BuildWhileSubgraph(Subgraph* subgraph) {
const int kCurrentIndexIn = 0;
const int kCurrentIndexOut = 1;
const int kListIn = 2;
const int kListOut = 3;
const int kTensorCount = 4;
// |---------------|
// | |
// kCurrentIndexIn(0) --> +-------+ --> (CondSubgraph)
// | | --> kCurrentIndexOut
// | WHILE |
// | | --> kListOut
// kListIn(2) ---> +-------+ --> (BodySubgraph)
// | |
// |---------------|
int first_new_tensor_index;
ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
kTfLiteOk);
ASSERT_EQ(first_new_tensor_index, 0);
ASSERT_EQ(subgraph->SetInputs({kCurrentIndexIn, kListIn}), kTfLiteOk);
ASSERT_EQ(subgraph->SetOutputs({kCurrentIndexOut, kListOut}), kTfLiteOk);
SetupTensor(subgraph, kCurrentIndexIn, kTfLiteInt32);
SetupTensor(subgraph, kCurrentIndexOut, kTfLiteInt32);
SetupTensor(subgraph, kListIn, kTfLiteVariant);
SetupTensor(subgraph, kListOut, kTfLiteVariant);
// `subgraph` takes ownership of `TfLiteWhileParams` and will `free`,
// so `malloc` is required.
TfLiteWhileParams* params =
reinterpret_cast<TfLiteWhileParams*>(malloc(sizeof(TfLiteWhileParams)));
params->cond_subgraph_index = 1;
params->body_subgraph_index = 2;
auto* while_reg = ops::builtin::Register_WHILE();
while_reg->builtin_code = kTfLiteBuiltinWhile;
int node_index;
TfLiteStatus stat = subgraph->AddNodeWithParameters(
{kCurrentIndexIn, kListIn}, {kCurrentIndexOut, kListOut},
/*intermediates=*/{}, /*init_data=*/nullptr, /*init_data_size=*/0, params,
while_reg, &node_index);
TF_LITE_ASSERT_EQ(stat, kTfLiteOk);
}
void ListOpsSubgraphBuilder::BuildSetItemAndIncrementSubgraph(
Subgraph* subgraph) {
const int kCurIn = 0;
const int kCurOut = 1;
const int kListIn = 2;
const int kListOut = 3;
const int kConstIncrement = 4;
const int kConstTensor = 5;
const int kTensorCount = 6;
// kListIn(2) --> +----------+
// kConstTensor(5) --> | SET_ITEM | --> kListOut(3)
// kCurIn(0) --> +----------+
//
// kConstIncrement(4) --> +-----+
// | ADD | --> kCurOut(1)
// kCurIn(0) --> +-----+
int first_new_tensor_index;
ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
kTfLiteOk);
ASSERT_EQ(first_new_tensor_index, 0);
ASSERT_EQ(subgraph->SetInputs({kCurIn, kListIn}), kTfLiteOk);
ASSERT_EQ(subgraph->SetOutputs({kCurOut, kListOut}), kTfLiteOk);
SetupTensor(subgraph, kCurIn, kTfLiteInt32);
SetupTensor(subgraph, kCurOut, kTfLiteInt32);
SetupTensor(subgraph, kListIn, kTfLiteVariant);
SetupTensor(subgraph, kListOut, kTfLiteVariant);
CreateConstantInt32Tensor(subgraph, kConstIncrement, {1}, {1});
CreateConstantInt32Tensor(subgraph, kConstTensor, {2}, {2, 2});
auto* set_item_res = Register_LIST_SET_ITEM();
set_item_res->custom_name = "ListSetItem";
int node_index;
TfLiteStatus set_stat = subgraph->AddNodeWithParameters(
{kListIn, kCurIn, kConstTensor}, {kListOut}, {}, /*init_data=*/nullptr,
/*init_data_size=*/0, /*builtin_data=*/nullptr, set_item_res,
&node_index);
TF_LITE_ASSERT_EQ(set_stat, kTfLiteOk);
TfLiteAddParams* params =
reinterpret_cast<TfLiteAddParams*>(malloc(sizeof(TfLiteAddParams)));
params->activation = kTfLiteActNone;
auto* add_reg = ops::builtin::Register_ADD();
add_reg->builtin_code = kTfLiteBuiltinAdd;
int node_index_add;
TfLiteStatus add_stat = subgraph->AddNodeWithParameters(
{kConstIncrement, kCurIn}, {kCurOut}, {}, /*init_data=*/nullptr,
/*init_data_size=*/0, params, add_reg, &node_index_add);
TF_LITE_ASSERT_EQ(add_stat, kTfLiteOk);
}
void ListOpsSubgraphBuilder::BuildReserveLengthSubgraph(Subgraph* subgraph) {
constexpr int kElementShape = 0;
constexpr int kNumElements = 1;
constexpr int kReserveOut = 2;
constexpr int kLengthOut = 3;
constexpr int kTensorCount = 4;
// kElementShape(0) --> +-------------+
// | ListReserve |
// kNumElements(1) --> +-------------+ --> kReserveOut(2)
// |
// +------------+
// | ListLength |
// +------------+ --> kLengthOut(3)
int first_new_tensor_index;
TF_LITE_ASSERT_EQ(subgraph->AddTensors(kTensorCount, &first_new_tensor_index),
kTfLiteOk);
TF_LITE_ASSERT_EQ(first_new_tensor_index, 0);
TF_LITE_ASSERT_EQ(subgraph->SetOutputs({kLengthOut}), kTfLiteOk);
SetupTensor(subgraph, kLengthOut, kTfLiteInt32);
TF_LITE_ASSERT_EQ(subgraph->SetInputs({kElementShape, kNumElements}),
kTfLiteOk);
SetupTensor(subgraph, kElementShape, kTfLiteInt32);
SetupTensor(subgraph, kNumElements, kTfLiteInt32);
SetupTensor(subgraph, kReserveOut, kTfLiteVariant);
TfLiteRegistration* reserve_reg = Register_LIST_RESERVE();
reserve_reg->builtin_code = BuiltinOperator_CUSTOM;
reserve_reg->custom_name = "ListReserve";
ListReserveOptions* options = RequestReserveOptions(TensorType_INT32);
int reserve_node_index;
TfLiteStatus stat = subgraph->AddNodeWithParameters(
{kElementShape, kNumElements}, {kReserveOut},
/*intermediates=*/{}, reinterpret_cast<const char*>(options),
sizeof(ListReserveOptions),
/*builtin_data=*/nullptr, reserve_reg, &reserve_node_index);
TF_LITE_ASSERT_EQ(stat, kTfLiteOk);
TfLiteRegistration* length_reg = Register_LIST_LENGTH();
length_reg->builtin_code = BuiltinOperator_CUSTOM;
length_reg->custom_name = "ListLength";
int length_node_index;
stat = subgraph->AddNodeWithParameters(
{kReserveOut}, {kLengthOut},
/*intermediates=*/{}, /*init_data=*/nullptr,
/*init_data_size=*/0,
/*builtin_data=*/nullptr, length_reg, &length_node_index);
TF_LITE_ASSERT_EQ(stat, kTfLiteOk);
}
} // namespace tflite
tensorflow/lite/kernels/variants/list_ops_subgraph_test_util.h
+89
View File
@@ -0,0 +1,89 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_VARIANTS_LIST_OPS_SUBGRAPH_TEST_UTIL_H_
#define TENSORFLOW_LITE_KERNELS_VARIANTS_LIST_OPS_SUBGRAPH_TEST_UTIL_H_
#include <cstdint>
#include <vector>
#include <gtest/gtest.h>
#include "absl/types/span.h"
#include "tensorflow/lite/core/interpreter.h"
#include "tensorflow/lite/core/subgraph.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/schema/schema_generated.h"
namespace tflite {
// Helper class for constructing complicated subgraphs for testing.
class ListOpsSubgraphBuilder {
public:
// Populates the given Subgraph with ops to add the value of two constants
// created by `CreateConstantInt32Tensor`.
void BuildAddConstSubgraph(Subgraph* subgraph);
// Populates the given Subgraph with a "ListReserve" op whose elements
// have type `element_type`.
void BuildReserveSubgraph(Subgraph* subgraph, TensorType element_type);
// Populates the given Subgraph with "ListStack" op that takes in
// a "ListReserve". Element types are i32.
void BuildReserveStackSubgraph(Subgraph* subgraph);
// Populates the given Subgraph with a "While" op, whose cond and body
// subgraphs are located at index 1, 2 respectively. The input signatures
// of both subgraphs 1, 2 are expected to be (kTfLiteInt32, kTfLiteVariant).
void BuildWhileSubgraph(Subgraph* subgraph);
// Populates the given Subgraph with a single "Less" op which checks
// if a given int is less that the constant 3. Also takes a `kTfLiteVariant`
// tensor in order to be compliant with `BuildWhileSubgraph`.
void BuildLessThanSubgraph(Subgraph* subgraph);
// Populates the given Subgraph with a "ListSetItem" op which sets the element
// at given indice into given tensorlist. Additionally increment and return
// the given int by 1.
void BuildSetItemAndIncrementSubgraph(Subgraph* subgraph);
// Populates the given Subgraph with a "ListReserve" and "ListLength" op.
void BuildReserveLengthSubgraph(Subgraph* subgraph);
private:
// Creates a constant tensor in given Subgraphs at given indice with
// corresponding data.
void CreateConstantInt32Tensor(Subgraph* subgraph, int tensor_index,
absl::Span<const int> shape,
absl::Span<const int> data);
// Custom options usually live in the flatbuffer, so they won't
// be cleaned up by the `Interpreter`. When we create them in test
// we need to free when test is done. So we provide a factory function
// for construction.
variants::detail::ListReserveOptions* RequestReserveOptions(
TensorType element_type);
std::vector<variants::detail::ListReserveOptions> list_reserve_opts_;
std::vector<std::vector<int32_t>> int_buffers_;
};
class ListOpsSubgraphTest : public ::testing::Test {
protected:
Interpreter interpreter_;
ListOpsSubgraphBuilder builder_;
};
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_VARIANTS_LIST_OPS_SUBGRAPH_TEST_UTIL_H_
tensorflow/lite/kernels/variants/list_ops_util.cc
+112
View File
@@ -0,0 +1,112 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/lite/kernels/variants/list_ops_util.h"
#include <cstdint>
#include <vector>
#include "tensorflow/lite/array.h"
#include "tensorflow/lite/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
namespace tflite {
namespace variants {
// Creates a `TfLiteIntArray*` from tensor data that represents a shape.
TfLiteStatus TensorAsShape(TfLiteContext* context, const TfLiteTensor& shape,
IntArrayUniquePtr& result) {
if (shape.dims->size == 0) {
// `shape` tensor encode an unranked shape.
result = BuildTfLiteArray({});
return kTfLiteOk;
}
TF_LITE_ENSURE_MSG(context, shape.dims->size <= 1,
"Shape tensor must be 1D or unranked.");
const int rank = shape.dims->data[0];
TF_LITE_ENSURE_MSG(context,
rank >= 0 && shape.bytes >= rank * sizeof(int32_t) &&
(rank == 0 || shape.data.data != nullptr),
"Malformed shape tensor.");
const int* begin = reinterpret_cast<const int*>(shape.data.data);
const int* end = begin + rank;
result = BuildTfLiteArray(std::vector<int>(begin, end));
return kTfLiteOk;
}
IntArrayUniquePtr MergeShapesOrNull(IntArrayUniquePtr l, IntArrayUniquePtr r) {
if (l == nullptr) {
return r;
}
if (r == nullptr) {
return l;
}
if (l->size == 0) {
return r;
}
if (r->size == 0) {
return l;
}
if (l->size != r->size) {
return nullptr;
}
for (int i = 0; i < r->size; ++i) {
if (l->data[i] == -1) {
l->data[i] = r->data[i];
} else if (r->data[i] != -1 && l->data[i] != r->data[i]) {
return nullptr;
}
}
return l;
}
bool IsShapeFullyDefined(const TfLiteIntArray& shape) {
for (int i = 0; i < shape.size; ++i) {
if (shape.data[i] < 0) {
return false;
}
}
return true;
}
TfLiteStatus GetShapeIfAllEqual(const TensorArray& arr,
IntArrayUniquePtr& result) {
// All shapes are equal iff all shapes equal the first present shape.
const TfLiteIntArray* common_shape = nullptr;
for (int i = 0; i < arr.NumElements(); ++i) {
const TfLiteTensor* cur_element = arr.At(i);
if (cur_element == nullptr) {
continue;
}
if (common_shape == nullptr) {
common_shape = cur_element->dims;
continue;
}
if (!TfLiteIntArrayEqual(common_shape, cur_element->dims)) {
return kTfLiteError;
}
}
result = common_shape != nullptr ? BuildTfLiteArray(*common_shape) : nullptr;
return kTfLiteOk;
}
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/list_ops_util.h
+60
View File
@@ -0,0 +1,60 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_VARIANTS_LIST_OPS_UTIL_H_
#define TENSORFLOW_LITE_KERNELS_VARIANTS_LIST_OPS_UTIL_H_
#include "tensorflow/lite/array.h"
#include "tensorflow/lite/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/util.h"
namespace tflite {
namespace variants {
// Creates a dims array from tensor whose data represents a shape
// signature.
// TODO(b/288302706) `TfLiteIntArray` is ill-equiped to encode the semantics
// of something like a `tf.TensorShape`. In particular, there is no way
// to cleanly capture the difference between a concrete scalar shape, and an
// unranked shape-signature. The latter is defined to be compatible with any
// shape (like `tf.TensorShape(None)`). This causes the need for some extra
// checks. Consider wrapping in something like a `std::union` to differentiate
// between these cases.
TfLiteStatus TensorAsShape(TfLiteContext* context, const TfLiteTensor& shape,
IntArrayUniquePtr& result);
// "Merges" two shape signature arrays if possible, returns nullptr otherwise.
// Merging means to compute the most specified possible signature which
// is compatible with both inputs.
IntArrayUniquePtr MergeShapesOrNull(IntArrayUniquePtr l, IntArrayUniquePtr r);
// Checks if array encodes a fully defined shape.
bool IsShapeFullyDefined(const TfLiteIntArray& shape);
// Returns a status denoting whether all of the elements in the `arr`
// have the same shape. Write that shape to `result`.
// If the `arr` has no set elements, still succeed but set `result` to nullptr.
// TODO(b/288302706) This may be a performance bottleneck. We could potentially
// amortize this work by constraining `TensorArray::element_shape_` every
// time an element is added. This may cause divergence from tensorflow behavior
// however; further investigation is needed.
TfLiteStatus GetShapeIfAllEqual(const TensorArray& arr,
IntArrayUniquePtr& result);
} // namespace variants
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_VARIANTS_LIST_OPS_UTIL_H_
tensorflow/lite/kernels/variants/list_ops_util_test.cc
+277
View File
@@ -0,0 +1,277 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/lite/kernels/variants/list_ops_util.h"
#include <cstdarg>
#include <cstdio>
#include <string>
#include <utility>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/lite/array.h"
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include "tensorflow/lite/util.h"
namespace tflite {
namespace variants {
namespace {
struct TestContext : public TfLiteContext {
std::string error;
};
void ReportError(TfLiteContext* context, const char* format, ...) {
TestContext* c = static_cast<TestContext*>(context);
const size_t kBufferSize = 1024;
char temp_buffer[kBufferSize];
va_list args;
va_start(args, format);
vsnprintf(temp_buffer, kBufferSize, format, args);
va_end(args);
c->error = temp_buffer;
}
class TensorAsShapeTest : public ::testing::Test {
public:
TensorAsShapeTest() { context_.ReportError = ReportError; }
protected:
TestContext context_;
};
TEST_F(TensorAsShapeTest, ScalarTensor_ReturnsEmptyIntArray) {
TensorUniquePtr scalar_tensor =
BuildTfLiteTensor(kTfLiteInt32, BuildTfLiteArray({}), kTfLiteDynamic);
IntArrayUniquePtr shape_from_tensor;
ASSERT_EQ(TensorAsShape(&context_, *scalar_tensor, shape_from_tensor),
kTfLiteOk);
ASSERT_THAT(shape_from_tensor.get(), DimsAre({}));
}
TEST_F(TensorAsShapeTest, SingleElementTensor_ReturnsSize1Shape) {
TensorUniquePtr single_el_tensor =
BuildTfLiteTensor(kTfLiteInt32, BuildTfLiteArray({1}), kTfLiteDynamic);
single_el_tensor->data.i32[0] = 10;
IntArrayUniquePtr shape_from_tensor;
ASSERT_EQ(TensorAsShape(&context_, *single_el_tensor, shape_from_tensor),
kTfLiteOk);
ASSERT_THAT(shape_from_tensor.get(), DimsAre({10}));
}
TEST_F(TensorAsShapeTest, OneDMultipleElementShape_ReturnsHighRankedShape) {
TensorUniquePtr one_d_mul_el_tensor =
BuildTfLiteTensor(kTfLiteInt32, BuildTfLiteArray({3}), kTfLiteDynamic);
one_d_mul_el_tensor->data.i32[0] = 10;
one_d_mul_el_tensor->data.i32[1] = 9;
one_d_mul_el_tensor->data.i32[2] = 8;
IntArrayUniquePtr shape_from_tensor;
ASSERT_EQ(TensorAsShape(&context_, *one_d_mul_el_tensor, shape_from_tensor),
kTfLiteOk);
ASSERT_THAT(shape_from_tensor.get(), DimsAre({10, 9, 8}));
}
TEST_F(TensorAsShapeTest, MultidimensionalTensor_ReturnsNull) {
TensorUniquePtr multi_d_tensor =
BuildTfLiteTensor(kTfLiteInt32, BuildTfLiteArray({2, 3}), kTfLiteDynamic);
IntArrayUniquePtr shape_from_tensor;
EXPECT_EQ(TensorAsShape(&context_, *multi_d_tensor, shape_from_tensor),
kTfLiteError);
EXPECT_EQ(shape_from_tensor.get(), nullptr);
}
TEST_F(TensorAsShapeTest, BufferTooSmall_ReturnsNull) {
TensorUniquePtr small_buffer_tensor =
BuildTfLiteTensor(kTfLiteInt32, BuildTfLiteArray({1000}), kTfLiteDynamic);
// Artificially reduce bytes to simulate OOB read vulnerability
small_buffer_tensor->bytes = 0;
IntArrayUniquePtr shape_from_tensor;
EXPECT_EQ(TensorAsShape(&context_, *small_buffer_tensor, shape_from_tensor),
kTfLiteError);
EXPECT_EQ(shape_from_tensor.get(), nullptr);
}
TEST(MergeShapesOrNull, IncompatibleSameRank_ReturnsNull) {
IntArrayUniquePtr l = BuildTfLiteArray({2, 3});
IntArrayUniquePtr r = BuildTfLiteArray({3, 3});
EXPECT_EQ(MergeShapesOrNull(std::move(l), std::move(r)).get(), nullptr);
}
TEST(MergeShapesOrNull, NotSameRank_ReturnsNull) {
IntArrayUniquePtr l = BuildTfLiteArray({1});
IntArrayUniquePtr r = BuildTfLiteArray({1, 2});
EXPECT_EQ(MergeShapesOrNull(std::move(l), std::move(r)).get(), nullptr);
}
TEST(MergeShapesOrNull, IncompatibleSameRank1D_ReturnsNull) {
IntArrayUniquePtr l = BuildTfLiteArray({1});
IntArrayUniquePtr r = BuildTfLiteArray({2});
EXPECT_EQ(MergeShapesOrNull(std::move(l), std::move(r)).get(), nullptr);
}
TEST(MergeShapesOrNull, RankedUnknownLKnownR_ReturnsStatic) {
IntArrayUniquePtr l = BuildTfLiteArray({-1});
IntArrayUniquePtr r = BuildTfLiteArray({2});
EXPECT_THAT(MergeShapesOrNull(std::move(l), std::move(r)).get(),
DimsAre({2}));
}
TEST(MergeShapesOrNull, UnknownRKnownL_ReturnsStatic) {
IntArrayUniquePtr l = BuildTfLiteArray({2});
IntArrayUniquePtr r = BuildTfLiteArray({-1});
EXPECT_THAT(MergeShapesOrNull(std::move(l), std::move(r)).get(),
DimsAre({2}));
}
TEST(MergeShapesOrNull, UnknownBoth_ReturnsUnknown) {
IntArrayUniquePtr l = BuildTfLiteArray({-1});
IntArrayUniquePtr r = BuildTfLiteArray({-1});
EXPECT_THAT(MergeShapesOrNull(std::move(l), std::move(r)).get(),
DimsAre({-1}));
}
TEST(MergeShapesOrNull, RankedUnknownDifferentDims_ConstrainsUnknownDims) {
IntArrayUniquePtr l = BuildTfLiteArray({-1, 2, 5});
IntArrayUniquePtr r = BuildTfLiteArray({1, -1, 5});
EXPECT_THAT(MergeShapesOrNull(std::move(l), std::move(r)).get(),
DimsAre({1, 2, 5}));
}
TEST(MergeShapesOrNull, BothUnranked_ReturnsUnranked) {
IntArrayUniquePtr l = BuildTfLiteArray({});
IntArrayUniquePtr r = BuildTfLiteArray({});
EXPECT_THAT(MergeShapesOrNull(std::move(l), std::move(r)).get(), DimsAre({}));
}
TEST(MergeShapesOrNull, UnrankedAndStatic1D_ReturnsStatic1D) {
IntArrayUniquePtr l = BuildTfLiteArray({});
IntArrayUniquePtr r = BuildTfLiteArray({1});
EXPECT_THAT(MergeShapesOrNull(std::move(l), std::move(r)).get(),
DimsAre({1}));
}
TEST(MergeShapesOrNull, UnrankedAndStaticND_ReturnsStaticND) {
IntArrayUniquePtr l = BuildTfLiteArray({});
IntArrayUniquePtr r = BuildTfLiteArray({2, 3});
EXPECT_THAT(MergeShapesOrNull(std::move(l), std::move(r)).get(),
DimsAre({2, 3}));
}
TEST(MergeShapesOrNull, UnrankedAndRankedUnknown_ReturnsRankedUnknown) {
IntArrayUniquePtr l = BuildTfLiteArray({});
IntArrayUniquePtr r = BuildTfLiteArray({-1});
EXPECT_THAT(MergeShapesOrNull(std::move(l), std::move(r)).get(),
DimsAre({-1}));
}
TEST(MergeShapesOrNull, NullInput_ReturnsOther) {
EXPECT_THAT(MergeShapesOrNull(BuildTfLiteArray({3}), nullptr).get(),
DimsAre({3}));
EXPECT_THAT(MergeShapesOrNull(nullptr, BuildTfLiteArray({2})).get(),
DimsAre({2}));
EXPECT_EQ(MergeShapesOrNull(nullptr, nullptr).get(), nullptr);
}
TEST(MergeShapesOrNull, NullInput_ReturnsUnrankedOther) {
EXPECT_THAT(MergeShapesOrNull(BuildTfLiteArray({}), nullptr).get(),
DimsAre({}));
EXPECT_THAT(MergeShapesOrNull(nullptr, BuildTfLiteArray({})).get(),
DimsAre({}));
}
TEST(ElementsSameShape, NoElements_SucceedsWithNullptr) {
TensorArray arr = {kTfLiteInt32, BuildTfLiteArray({})};
arr.Resize(2);
IntArrayUniquePtr res;
ASSERT_EQ(GetShapeIfAllEqual(arr, res), kTfLiteOk);
EXPECT_EQ(res.get(), nullptr);
}
TEST(ElementsSameShape, ZeroSize_SucceedsWithNullptr) {
TensorArray arr = {kTfLiteInt32, BuildTfLiteArray({})};
IntArrayUniquePtr res;
ASSERT_EQ(GetShapeIfAllEqual(arr, res), kTfLiteOk);
EXPECT_EQ(res.get(), nullptr);
}
TEST(ElementsSameShape, OneSize_SucceedsWithShape) {
TensorArray arr = {kTfLiteInt32, BuildTfLiteArray({})};
arr.Resize(1);
arr.Set(0, BuildTfLiteTensor(kTfLiteInt32, BuildTfLiteArray({2}),
kTfLiteDynamic));
IntArrayUniquePtr res;
ASSERT_EQ(GetShapeIfAllEqual(arr, res), kTfLiteOk);
EXPECT_THAT(res.get(), DimsAre({2}));
}
TEST(ElementsSameShape, MultipleElements_AllSet_SucceedsWithShape) {
TensorArray arr = {kTfLiteInt32, BuildTfLiteArray({})};
arr.Resize(2);
arr.Set(0, BuildTfLiteTensor(kTfLiteInt32, BuildTfLiteArray({2}),
kTfLiteDynamic));
arr.Set(1, BuildTfLiteTensor(kTfLiteInt32, BuildTfLiteArray({2}),
kTfLiteDynamic));
IntArrayUniquePtr res;
EXPECT_EQ(GetShapeIfAllEqual(arr, res), kTfLiteOk);
EXPECT_THAT(res.get(), DimsAre({2}));
}
TEST(ElementsSameShape, MultipleElements_SomeSet_SucceedsWithShape) {
TensorArray arr = {kTfLiteInt32, BuildTfLiteArray({})};
arr.Resize(3);
arr.Set(0, BuildTfLiteTensor(kTfLiteInt32, BuildTfLiteArray({2}),
kTfLiteDynamic));
arr.Set(2, BuildTfLiteTensor(kTfLiteInt32, BuildTfLiteArray({2}),
kTfLiteDynamic));
IntArrayUniquePtr res;
EXPECT_EQ(GetShapeIfAllEqual(arr, res), kTfLiteOk);
EXPECT_THAT(res.get(), DimsAre({2}));
}
TEST(ElementsSameShape, MultipleElements_SomeSetNotSameRank_Fails) {
TensorArray arr = {kTfLiteInt32, BuildTfLiteArray({})};
arr.Resize(3);
arr.Set(0, BuildTfLiteTensor(kTfLiteInt32, BuildTfLiteArray({2}),
kTfLiteDynamic));
arr.Set(2, BuildTfLiteTensor(kTfLiteInt32, BuildTfLiteArray({2, 3}),
kTfLiteDynamic));
IntArrayUniquePtr res;
EXPECT_EQ(GetShapeIfAllEqual(arr, res), kTfLiteError);
}
TEST(ElementsSameShape, MultipleElements_SomeSetNotSameDim_Fails) {
TensorArray arr = {kTfLiteInt32, BuildTfLiteArray({})};
arr.Resize(3);
arr.Set(0, BuildTfLiteTensor(kTfLiteInt32, BuildTfLiteArray({2, 2}),
kTfLiteDynamic));
arr.Set(2, BuildTfLiteTensor(kTfLiteInt32, BuildTfLiteArray({2, 3}),
kTfLiteDynamic));
IntArrayUniquePtr res;
EXPECT_EQ(GetShapeIfAllEqual(arr, res), kTfLiteError);
}
} // namespace
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/py/BUILD
+37
View File
@@ -0,0 +1,37 @@
load("//tensorflow:tensorflow.bzl", "py_test")
load("//tensorflow:tensorflow.default.bzl", "pybind_extension")
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:LICENSE"])
pybind_extension(
name = "register_list_ops_py",
srcs = ["register_list_ops_py.cc"],
enable_stub_generation = True,
pytype_srcs = [
"register_list_ops_py.pyi",
],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/lite:mutable_op_resolver",
"//tensorflow/lite/kernels/variants:list_ops_lib",
"@pybind11",
],
)
py_test(
name = "end_to_end_test",
srcs = ["end_to_end_test.py"],
strict_deps = True,
tags = ["nochromiumos_arm"],
deps = [
":register_list_ops_py",
"@absl_py//absl/testing:parameterized",
#internal proto upb dep
"//third_party/py/numpy",
"//tensorflow:tensorflow_py",
"//tensorflow/lite/python:interpreter",
"//tensorflow/lite/python:lite",
"//tensorflow/python/ops:list_ops",
"//tensorflow/python/platform:test",
],
)
tensorflow/lite/kernels/variants/py/end_to_end_test.py
+474
View File
@@ -0,0 +1,474 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from absl.testing import parameterized
import numpy as np
import tensorflow as tf
from tensorflow.lite.kernels.variants.py import register_list_ops_py
from tensorflow.lite.python import interpreter as _interpreter
from tensorflow.lite.python import lite
from tensorflow.python.ops import list_ops
from tensorflow.python.platform import googletest
class ListOpsTest(parameterized.TestCase):
"""Checks for tfl custom TensorList kernels.
* Validate conversion through python api with
`legalize_custom_tensor_list_ops = True`.
* Validate python bindings for registration through
`InterpreterWithCustomOps`.
* Check tensorflow and tensorflow lite output the same data.
"""
def _get_interpreter_from_c_func(self, func):
concrete_function_list = [func.get_concrete_function()]
converter = lite.TFLiteConverterV2.from_concrete_functions(
concrete_function_list
)
# Don't allow flex ops.
converter.target_spec.supported_ops = [
lite.OpsSet.TFLITE_BUILTINS,
]
converter.allow_custom_ops = True
converter.legalize_custom_tensor_list_ops = True
converter._experimental_lower_tensor_list_ops = False
tfl_model = converter.convert()
# Instantiate interpreter with custom tensor list ops.
interpreter = _interpreter.InterpreterWithCustomOps(
model_content=tfl_model,
custom_op_registerers=[register_list_ops_py.TFLRegisterListOps],
)
return interpreter
def _get_tfl_tf_outputs(self, tf_callable, *inp):
interpreter = self._get_interpreter_from_c_func(tf_callable)
get_inp_index = lambda i: interpreter.get_input_details()[i]["index"]
for i, arr in enumerate(inp):
interpreter.resize_tensor_input(get_inp_index(i), arr.shape)
interpreter.allocate_tensors()
for i, arr in enumerate(inp):
interpreter.set_tensor(get_inp_index(i), arr)
interpreter.invoke()
tfl_out = interpreter.get_tensor(
interpreter.get_output_details()[0]["index"]
)
tf_out = tf_callable(*inp)
return tfl_out, tf_out
@parameterized.named_parameters(
("1DFilledList", 2, 2, [2]),
("1DPartialList", 2, 1, [2]),
("2DFilledList", 2, 2, [2, 2]),
("2DPartialList", 2, 1, [2, 2]),
)
def test_reserve_set_stack(self, num_elements, num_set_ops, item_shape):
@tf.function(
input_signature=[
tf.TensorSpec(shape=tf.TensorShape(None), dtype=tf.float32)
]
)
def f(x: tf.Tensor) -> tf.Tensor:
l = list_ops.tensor_list_reserve(
element_shape=tf.TensorShape(None),
element_dtype=tf.float32,
num_elements=num_elements,
)
# This will be unrolled.
for i in range(num_set_ops):
l = list_ops.tensor_list_set_item(l, i, x)
return list_ops.tensor_list_stack(l, tf.float32)
inp = np.ndarray(shape=item_shape, dtype=np.float32)
inp.fill(1)
tfl_out, tf_out = self._get_tfl_tf_outputs(f, inp)
self.assertEqual(tf_out.dtype, tfl_out.dtype)
self.assertEqual(tf_out.shape, tfl_out.shape)
self.assertTrue((tf_out == tfl_out).numpy().all())
@parameterized.named_parameters(
("2DStatic", [2, 2], [2, 2]),
("3DStatic", [3, 3, 3], [3, 3, 3]),
("2DDynamic", None, [2, 2]),
("3DDynamic", None, [3, 3, 3]),
)
def test_from_tensor_stack(self, f_shape, t_shape):
@tf.function(
input_signature=[tf.TensorSpec(shape=f_shape, dtype=tf.float32)]
)
def f(x: tf.Tensor) -> tf.Tensor:
l = list_ops.tensor_list_from_tensor(x, element_shape=x.shape[1:])
return list_ops.tensor_list_stack(l, tf.float32)
inp = np.ndarray(shape=t_shape, dtype=np.float32)
inp.fill(1)
tfl_out, tf_out = self._get_tfl_tf_outputs(f, inp)
self.assertEqual(tf_out.dtype, tfl_out.dtype)
self.assertEqual(tf_out.shape, tfl_out.shape)
self.assertTrue((tf_out == tfl_out).numpy().all())
@parameterized.named_parameters(
("StaticReadPresent", [2, 2], [2, 2], 0, 0),
("StaticReadEmpty", [2, 2], [2, 2], 0, 1),
("DynamicReadPresent", None, [2, 2], 0, 0),
("DynamicReadEmpty", None, [2, 2], 0, 1),
)
def test_reserve_set_get(self, f_shape, t_shape, i_write, i_read):
@tf.function(
input_signature=[tf.TensorSpec(shape=f_shape, dtype=tf.float32)]
)
def f(x: tf.Tensor) -> tf.Tensor:
l = list_ops.tensor_list_reserve(
element_shape=tf.TensorShape(None),
element_dtype=tf.float32,
num_elements=2,
)
l = list_ops.tensor_list_set_item(l, i_write, x)
return list_ops.tensor_list_get_item(l, i_read, tf.float32)
inp = np.ndarray(shape=t_shape, dtype=np.float32)
inp.fill(1)
tfl_out, tf_out = self._get_tfl_tf_outputs(f, inp)
self.assertEqual(tf_out.dtype, tfl_out.dtype)
self.assertEqual(tf_out.shape, tfl_out.shape)
self.assertTrue((tf_out == tfl_out).numpy().all())
@parameterized.named_parameters(
("ZeroElements_ScalarStackShape", [], 0),
("NonZeroElements_ScalarStackShape", [], 2),
("NonZeroElements_ZeroStackShape", [0], 2),
("ZeroElements_ZeroStackShape", [0], 0),
("ZeroElements_2DZeroStackShape", [0, 2], 0),
("NonZeroElements_2DZeroStackShape", [0, 2], 2),
)
def test_stack_empty_list(
self, stack_element_shape: list[int], num_elements: int
):
@tf.function(
input_signature=[
tf.TensorSpec(shape=tf.TensorShape(None), dtype=tf.int32)
]
)
def f(stack_element_shape) -> tf.Tensor:
l = list_ops.tensor_list_reserve(
element_shape=tf.TensorShape(None),
element_dtype=tf.float32,
num_elements=num_elements,
)
return list_ops.tensor_list_stack(
l, element_shape=stack_element_shape, element_dtype=tf.float32
)
inp = np.array(stack_element_shape, dtype=np.int32)
tfl_out, tf_out = self._get_tfl_tf_outputs(f, inp)
self.assertEqual(tf_out.dtype, tfl_out.dtype)
self.assertEqual(tf_out.shape, tfl_out.shape)
self.assertTrue((tf_out == tfl_out).numpy().all())
@parameterized.named_parameters(
("OneD", [4]),
("TwoD", [2, 2]),
("ThreeD", [2, 3, 4]),
)
def test_from_pop_back_unrolled(
self,
tensor_shape: list[int],
):
@tf.function(
input_signature=[
tf.TensorSpec(shape=tf.TensorShape(None), dtype=tf.int32)
]
)
def f(x: tf.Tensor) -> tf.Tensor:
l = list_ops.tensor_list_from_tensor(x, element_shape=tensor_shape[1:])
a = tf.constant(0, tf.int32, shape=tensor_shape[1:])
for unused_i in range(tensor_shape[0]):
l, e = list_ops.tensor_list_pop_back(l, tf.int32)
a += e
return a
input_tensor = np.ndarray(shape=tensor_shape, dtype=np.int32)
input_tensor.fill(1)
tfl_out, tf_out = self._get_tfl_tf_outputs(f, input_tensor)
self.assertEqual(tf_out.dtype, tfl_out.dtype)
self.assertEqual(tf_out.shape, tfl_out.shape)
self.assertTrue((tf_out == tfl_out).numpy().all())
@parameterized.named_parameters(
("OneD", [4]),
("TwoD", [2, 2]),
("ThreeD", [2, 3, 4]),
)
def test_from_pop_back_loop(
self,
tensor_shape: list[int],
):
@tf.function(
input_signature=[
tf.TensorSpec(shape=tf.TensorShape(None), dtype=tf.int32)
]
)
def f(x: tf.Tensor) -> tf.Tensor:
l = list_ops.tensor_list_from_tensor(x, element_shape=tensor_shape[1:])
a = tf.constant(0, tf.int32, shape=tensor_shape[1:])
while len(l) > 0: # pylint: disable=g-explicit-length-test
l, e = list_ops.tensor_list_pop_back(l, tf.int32)
a += e
return a
input_tensor = np.ndarray(shape=tensor_shape, dtype=np.int32)
input_tensor.fill(1)
tfl_out, tf_out = self._get_tfl_tf_outputs(f, input_tensor)
self.assertEqual(tf_out.dtype, tfl_out.dtype)
self.assertEqual(tf_out.shape, tfl_out.shape)
self.assertTrue((tf_out == tfl_out).numpy().all())
def test_empty_tensorlist_set_stack(self):
@tf.function(
input_signature=[
tf.TensorSpec(shape=tf.TensorShape(None), dtype=tf.int32)
]
)
def f(x) -> tf.Tensor:
l = list_ops.empty_tensor_list(tf.TensorShape(None), tf.int32)
l2 = list_ops.tensor_list_set_item(l, 0, x, True)
return list_ops.tensor_list_stack(l2, tf.int32)
inp = np.ndarray(shape=[2, 2], dtype=np.int32)
inp.fill(0)
tfl_out, tf_out = self._get_tfl_tf_outputs(f, inp)
self.assertEqual(tf_out.dtype, tfl_out.dtype)
self.assertEqual(tf_out.shape, tfl_out.shape)
self.assertTrue((tf_out == tfl_out).numpy().all())
@parameterized.named_parameters(
("Unranked", None),
("DynDim", [None]),
("DynMultiDim", [None, 2]),
("AllStatic", [2, 2]),
)
def test_reserve_element_shape(self, element_shape):
@tf.function
def f() -> tf.Tensor:
l = list_ops.tensor_list_reserve(
element_shape=tf.TensorShape(element_shape),
element_dtype=tf.int32,
num_elements=10,
)
return list_ops.tensor_list_element_shape(l, tf.int32)
tfl_out, tf_out = self._get_tfl_tf_outputs(f)
self.assertEqual(tf_out.dtype, tfl_out.dtype)
self.assertEqual(tf_out.shape, tfl_out.shape)
self.assertTrue((tf_out == tfl_out).numpy().all())
@parameterized.named_parameters(
("ReserveLength0", 0),
("ReserveLength1", 1),
("ReserveLength5", 5),
)
def test_reserve_length(self, num_elements: int):
@tf.function(
input_signature=[
tf.TensorSpec(shape=tf.TensorShape([]), dtype=tf.int32)
]
)
def f(num_elements) -> tf.Tensor:
l = list_ops.tensor_list_reserve(
element_shape=tf.TensorShape(None),
element_dtype=tf.int32,
num_elements=num_elements,
)
return list_ops.tensor_list_length(l)
inp = np.array(num_elements, dtype=np.int32)
tfl_out, tf_out = self._get_tfl_tf_outputs(f, inp)
self.assertEqual(tf_out.dtype, tfl_out.dtype)
self.assertEqual(tf_out.shape, tfl_out.shape)
self.assertTrue((tf_out == tfl_out).numpy().all())
@parameterized.named_parameters(
("OneD", [2]),
("TwoD", [2, 2]),
("3D", [2, 3, 4]),
("Scalar", []),
)
def test_push_back_stack(self, element_shape):
@tf.function(
input_signature=[
tf.TensorSpec(shape=tf.TensorShape(element_shape), dtype=tf.int32)
]
)
def f(x) -> tf.Tensor:
l = list_ops.empty_tensor_list(tf.TensorShape(None), tf.int32)
for unused_i in range(10):
l = list_ops.tensor_list_push_back(l, x)
return list_ops.tensor_list_stack(l, tf.int32)
input_tensor = np.ndarray(shape=element_shape, dtype=np.int32)
input_tensor.fill(0)
tfl_out, tf_out = self._get_tfl_tf_outputs(f, input_tensor)
self.assertEqual(tf_out.dtype, tfl_out.dtype)
self.assertEqual(tf_out.shape, tfl_out.shape)
self.assertTrue((tf_out == tfl_out).numpy().all())
@parameterized.named_parameters(
("1D", [4]),
("2D", [2, 2]),
("3D", [2, 3, 4]),
)
def test_from_tensor_addn(self, element_shape):
@tf.function(
input_signature=[
tf.TensorSpec(shape=tf.TensorShape(None), dtype=tf.int32)
]
)
def f(x) -> tf.Tensor:
l = list_ops.tensor_list_from_tensor(x, element_shape[1:])
l2 = list_ops.tensor_list_from_tensor(x, element_shape[1:])
l3 = tf.math.add_n([l, l2])
return list_ops.tensor_list_stack(l3, tf.int32)
inp = np.ndarray(shape=element_shape, dtype=np.int32)
inp.fill(1)
tfl_out, tf_out = self._get_tfl_tf_outputs(f, inp)
self.assertEqual(tf_out.dtype, tfl_out.dtype)
self.assertEqual(tf_out.shape, tfl_out.shape)
self.assertTrue((tf_out == tfl_out).numpy().all())
@parameterized.named_parameters(
("StaticSameShape", [2, 2], [2, 2]),
("FullyDefinedAndUnranked", [2, 2], None),
("MergeToFullyDefined", [2, None], [None, 2]),
)
def test_reserve_addn(self, list1_shape, list2_shape):
@tf.function
def f() -> tf.Tensor:
l = list_ops.tensor_list_reserve(
element_shape=tf.TensorShape(list1_shape),
element_dtype=tf.int32,
num_elements=10,
)
l2 = list_ops.tensor_list_reserve(
element_shape=tf.TensorShape(list2_shape),
element_dtype=tf.int32,
num_elements=10,
)
l3 = tf.math.add_n([l, l2])
return list_ops.tensor_list_stack(l3, tf.int32)
tfl_out, tf_out = self._get_tfl_tf_outputs(f)
self.assertEqual(tf_out.dtype, tfl_out.dtype)
self.assertEqual(tf_out.shape, tfl_out.shape)
self.assertTrue((tf_out == tfl_out).numpy().all())
@parameterized.named_parameters(
("AllStaticShapesEqual", [2, 2], [2, 2], [2, 2]),
("ListShapesUnranked", None, None, [3, 3]),
)
def test_reserve_set_single_item_addn(
self, list1_shape, list2_shape, input1_shape
):
@tf.function(
input_signature=[
tf.TensorSpec(shape=tf.TensorShape(None), dtype=tf.int32),
]
)
def f(x) -> tf.Tensor:
l = list_ops.tensor_list_reserve(
element_shape=tf.TensorShape(list1_shape),
element_dtype=tf.int32,
num_elements=1,
)
l = list_ops.tensor_list_set_item(l, 0, x)
l2 = list_ops.tensor_list_reserve(
element_shape=tf.TensorShape(list2_shape),
element_dtype=tf.int32,
num_elements=1,
)
m = tf.math.add_n([l, l2])
return list_ops.tensor_list_get_item(m, 0, tf.int32)
inp1 = np.ndarray(shape=input1_shape, dtype=np.int32)
inp1.fill(1)
tfl_out, tf_out = self._get_tfl_tf_outputs(f, inp1)
self.assertEqual(tf_out.dtype, tfl_out.dtype)
self.assertEqual(tf_out.shape, tfl_out.shape)
self.assertTrue((tf_out == tfl_out).numpy().all())
@parameterized.named_parameters(
("AllStaticShapesEqual", [2, 2], [2, 2], [2, 2], [2, 2]),
("UrankedListShapes", None, None, [2, 2], [2, 2]),
)
def test_reserve_set_both_item_addn(
self, list1_shape, list2_shape, input1_shape, input2_shape
):
@tf.function(
input_signature=[
tf.TensorSpec(shape=tf.TensorShape(None), dtype=tf.int32),
tf.TensorSpec(shape=tf.TensorShape(None), dtype=tf.int32),
]
)
def f(x, y) -> tf.Tensor:
l = list_ops.tensor_list_reserve(
element_shape=tf.TensorShape(list1_shape),
element_dtype=tf.int32,
num_elements=1,
)
l = list_ops.tensor_list_set_item(l, 0, x)
l2 = list_ops.tensor_list_reserve(
element_shape=tf.TensorShape(list2_shape),
element_dtype=tf.int32,
num_elements=1,
)
l2 = list_ops.tensor_list_set_item(l2, 0, y)
m = tf.math.add_n([l, l2])
return list_ops.tensor_list_get_item(m, 0, tf.int32)
inp1 = np.ndarray(shape=input1_shape, dtype=np.int32)
inp1.fill(1)
inp2 = np.ndarray(shape=input2_shape, dtype=np.int32)
inp2.fill(1)
tfl_out, tf_out = self._get_tfl_tf_outputs(f, inp1, inp2)
self.assertEqual(tf_out.dtype, tfl_out.dtype)
self.assertEqual(tf_out.shape, tfl_out.shape)
self.assertTrue((tf_out == tfl_out).numpy().all())
if __name__ == "__main__":
googletest.main()
tensorflow/lite/kernels/variants/py/register_list_ops_py.cc
+32
View File
@@ -0,0 +1,32 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstdint>
#include "pybind11/pybind11.h" // from @pybind11
#include "pybind11/pytypes.h" // from @pybind11
#include "tensorflow/lite/kernels/variants/register_list_ops.h"
#include "tensorflow/lite/mutable_op_resolver.h"
PYBIND11_MODULE(register_list_ops_py, m) {
m.doc() = R"pbdoc(
Bindings to register list ops with python interpreter.
)pbdoc";
m.def(
"TFLRegisterListOps",
[](uintptr_t resolver) {
::tflite::variants::ops::RegisterListOps(
reinterpret_cast<::tflite::MutableOpResolver*>(resolver));
},
R"pbdoc(
Register all custom list ops.
)pbdoc");
}
tensorflow/lite/kernels/variants/py/register_list_ops_py.pyi
+16
View File
@@ -0,0 +1,16 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
def TFLRegisterListOps(arg0: int) -> None: ...
tensorflow/lite/kernels/variants/register_list_ops.cc
+44
View File
@@ -0,0 +1,44 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/lite/kernels/variants/register_list_ops.h"
#include "tensorflow/lite/kernels/variants/list_ops_lib.h"
#include "tensorflow/lite/mutable_op_resolver.h"
namespace tflite {
namespace variants {
namespace ops {
void RegisterListOps(MutableOpResolver* resolver) {
resolver->AddCustom("TensorListReserve", Register_LIST_RESERVE());
resolver->AddCustom("TensorListStack", Register_LIST_STACK());
resolver->AddCustom("TensorListSetItem", Register_LIST_SET_ITEM());
resolver->AddCustom("TensorListFromTensor", Register_LIST_FROM_TENSOR());
resolver->AddCustom("TensorListGetItem", Register_LIST_GET_ITEM());
resolver->AddCustom("TensorListLength", Register_LIST_LENGTH());
resolver->AddCustom("TensorListElementShape", Register_LIST_ELEMENT_SHAPE());
resolver->AddCustom("TensorListPopBack", Register_LIST_POP_BACK());
resolver->AddCustom("TensorListPushBack", Register_LIST_PUSH_BACK());
resolver->AddCustom("VariantAddN", Register_VARIANT_ADD_N());
resolver->AddCustom("VariantZerosLike", Register_VARIANT_ZEROS_LIKE());
}
} // namespace ops
} // namespace variants
} // namespace tflite
void RegisterSelectedOps(tflite::MutableOpResolver* resolver) {
tflite::variants::ops::RegisterListOps(resolver);
}
tensorflow/lite/kernels/variants/register_list_ops.h
+31
View File
@@ -0,0 +1,31 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_VARIANTS_REGISTER_LIST_OPS_H_
#define TENSORFLOW_LITE_KERNELS_VARIANTS_REGISTER_LIST_OPS_H_
#include "tensorflow/lite/mutable_op_resolver.h"
namespace tflite {
namespace variants {
namespace ops {
// Register all tflite TensorList kernels as custom ops with the given resolver.
void RegisterListOps(MutableOpResolver* resolver);
} // namespace ops
} // namespace variants
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_VARIANTS_REGISTER_LIST_OPS_H_
tensorflow/lite/kernels/variants/tensor_array.cc
+206
View File
@@ -0,0 +1,206 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include <cstdlib>
#include <cstring>
#include <utility>
#include "tensorflow/lite/array.h"
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/util.h"
namespace tflite {
namespace variants {
TensorArray::TensorArray(const TensorArray& other) {
TfLiteIntArray* copied_shape = TfLiteIntArrayCopy(other.element_shape_.get());
if (copied_shape == nullptr && other.element_shape_ != nullptr) {
element_shape_ = nullptr;
num_elements_ = 0;
elements_ = nullptr;
return;
}
element_shape_ = IntArrayUniquePtr(copied_shape);
element_type_ = other.element_type_;
if (other.num_elements_ == 0) {
num_elements_ = 0;
elements_ = nullptr;
return;
}
size_t bytes;
if (__builtin_mul_overflow(other.num_elements_, sizeof(RefCountedTensor),
&bytes)) {
num_elements_ = 0;
elements_ = nullptr;
return;
}
elements_ = static_cast<RefCountedTensor*>(malloc(bytes));
if (elements_ == nullptr) {
num_elements_ = 0;
return;
}
num_elements_ = other.num_elements_;
other.AssignBuffer(elements_);
}
TensorArray::TensorArray(TensorArray&& other) noexcept
: elements_(other.elements_),
num_elements_(other.num_elements_),
element_shape_(std::move(other.element_shape_)),
element_type_(other.element_type_) {
other.elements_ = nullptr;
other.num_elements_ = 0;
}
TensorArray& TensorArray::operator=(const TensorArray& other) {
if (this == &other) return *this;
TensorArray temp(other);
if ((other.num_elements_ > 0 && temp.elements_ == nullptr) ||
(other.element_shape_ != nullptr && temp.element_shape_ == nullptr)) {
return *this;
}
std::swap(element_shape_, temp.element_shape_);
std::swap(element_type_, temp.element_type_);
std::swap(elements_, temp.elements_);
std::swap(num_elements_, temp.num_elements_);
return *this;
}
TensorArray& TensorArray::operator=(TensorArray&& other) noexcept {
if (this == &other) return *this;
Clear();
free(elements_);
elements_ = other.elements_;
num_elements_ = other.num_elements_;
element_shape_ = std::move(other.element_shape_);
element_type_ = other.element_type_;
other.elements_ = nullptr;
other.num_elements_ = 0;
return *this;
}
bool TensorArray::Resize(int num_elements) {
if (num_elements == NumElements()) return true;
if (num_elements < 0) return false;
size_t bytes;
if (__builtin_mul_overflow(num_elements, sizeof(RefCountedTensor), &bytes)) {
return false;
}
if (num_elements > NumElements()) {
// The length of the array is being increased. Reallocate the buffer
// to the appropriate size and setup the new `RefCountedTensors`.
RefCountedTensor* new_elements =
static_cast<RefCountedTensor*>(realloc(elements_, bytes));
if (new_elements == nullptr && num_elements > 0) {
return false;
}
elements_ = new_elements;
for (int i = NumElements(); i < num_elements; ++i) {
elements_[i].count = nullptr;
elements_[i].tensor = nullptr;
}
} else {
// The length of the array is being decreased, update each of the
// references that we will no longer keep before reallocating.
for (int i = num_elements; i < NumElements(); ++i) {
Drop(i);
}
if (num_elements == 0) {
free(elements_);
elements_ = nullptr;
} else {
RefCountedTensor* new_elements =
static_cast<RefCountedTensor*>(realloc(elements_, bytes));
if (new_elements != nullptr) {
elements_ = new_elements;
}
}
}
num_elements_ = num_elements;
return true;
}
const TfLiteTensor* TensorArray::At(int index) const {
if (index < 0 || index >= NumElements()) {
return nullptr;
}
return elements_[index].tensor;
}
bool TensorArray::Set(int index, TensorUniquePtr tensor) {
if (index < 0 || index >= NumElements()) {
return false;
}
int* c = static_cast<int*>(malloc(sizeof(int)));
if (c == nullptr) {
return false;
}
// Drop element if it exists.
Drop(index);
// Setup the `RefCountedTensor` at given index to wrap the given tensor.
*c = 1;
elements_[index].tensor = tensor.release();
elements_[index].count = c;
return true;
}
TensorArray::~TensorArray() {
Clear();
free(elements_);
elements_ = nullptr;
}
void TensorArray::Drop(int i) {
RefCountedTensor* t = elements_ + i;
int* count = t->count;
if (count == nullptr) {
return;
}
if (*count == 1) {
if (t->tensor != nullptr) {
TfLiteTensorFree(t->tensor);
free(t->tensor);
}
free(t->count);
} else {
(*count)--;
}
t->tensor = nullptr;
t->count = nullptr;
}
// `Drop`s each element in the list.
void TensorArray::Clear() {
for (int i = 0; i < num_elements_; ++i) {
Drop(i);
}
}
void TensorArray::AssignBuffer(RefCountedTensor* dst) const {
if (num_elements_ == 0) return;
// Copy `this` underlying buffer.
std::memcpy(dst, elements_, sizeof(RefCountedTensor) * num_elements_);
// Increment the reference count for each copied tensor.
for (int i = 0; i < num_elements_; ++i) {
if (dst[i].count == nullptr) {
continue;
}
(*dst[i].count)++;
}
}
} // namespace variants
} // namespace tflite
tensorflow/lite/kernels/variants/tensor_array.h
+112
View File
@@ -0,0 +1,112 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_LITE_KERNELS_VARIANTS_TENSOR_ARRAY_H_
#define TENSORFLOW_LITE_KERNELS_VARIANTS_TENSOR_ARRAY_H_
#include <utility>
#include "tensorflow/lite/array.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/util.h"
namespace tflite {
namespace variants {
// `VariantData` implementation for a dynamically sized array of `TfLiteTensor`.
// Each element of the array is a lightweight `RefCountedTensor`.
// --- WARNING ---
// This is intended to be used in a single-threaded manner
// and users must take care when calling non-const methods, even on different
// instances. Different instances may share underlying control structures (when
// using the copy constructor to initialize them), and in such cases function
// calls across all affected instances must be properly synchronized. Calling
// non-const functions on any of the linked objects requires exclusive access to
// all of them.
//
// TODO(b/288302706) Implement standard container methods.
class TensorArray : public AbstractVariantData<TensorArray> {
public:
// Takes ownership of `element_shape` input.
TensorArray(TfLiteType element_type, IntArrayUniquePtr element_shape)
: element_shape_(std::move(element_shape)), element_type_(element_type) {}
// Copying a `TensorArray` copies the sources underlying array of
// `RefCountedTensor` in one `memcpy` and increments each of the ref counts.
TensorArray(const TensorArray& other);
TensorArray(TensorArray&& other) noexcept;
// Drops the references of `this` and assigns members in the same way
// as the copy constructor.
TensorArray& operator=(const TensorArray& other);
TensorArray& operator=(TensorArray&& other) noexcept;
const TfLiteIntArray* ElementShape() const { return element_shape_.get(); }
int NumElements() const { return num_elements_; }
// Resizes the array for given number of elements. If the length of the array
// is being decreased, `Drop` the reference to the elements that will no
// longer be in the array. If index is out of bounds, no effect.
bool Resize(int num_elements);
// Retrieve the tensor at the given index.
const TfLiteTensor* At(int index) const;
TfLiteType ElementType() const { return element_type_; }
// Set the item at the given index with the given tensor. Takes ownership
// of the given tensor. If there exists an element at the given index,
// `Drop` this array's reference to it.
bool Set(int index, TensorUniquePtr tensor);
// `Drop`s each reference that exists in the array.
~TensorArray() override;
private:
// Simple structure to hold tensor pointer and ref count. Only to be used
// as elements within `TensorArray`.
struct RefCountedTensor {
TfLiteTensor* tensor = nullptr;
int* count = nullptr;
};
// "Drops" the reference at the given index because it will no longer be held
// in this array. Decrements the reference count, if this array holds the only
// reference than free the underlying tensor.
void Drop(int i);
// `Drop`s each element in the list.
void Clear();
// Assigns this `elements_` buffer to `dst`. Requires that the size
// of this elements buffer be the same as `dst` and that `dst` has
// been `Clear`ed. Like the rest of this class, copying `this` buffer
// needs to increment references of `const this`, so beware.
void AssignBuffer(RefCountedTensor* dst) const;
// elements_ is nullptr iff num_elements is 0.
RefCountedTensor* elements_ = nullptr;
int num_elements_ = 0;
IntArrayUniquePtr element_shape_;
TfLiteType element_type_;
};
} // namespace variants
} // namespace tflite
#endif // TENSORFLOW_LITE_KERNELS_VARIANTS_TENSOR_ARRAY_H_
tensorflow/lite/kernels/variants/tensor_array_test.cc
+235
View File
@@ -0,0 +1,235 @@
/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/lite/kernels/variants/tensor_array.h"
#include <functional>
#include <memory>
#include <numeric>
#include <optional>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "tensorflow/lite/array.h"
#include "tensorflow/lite/c/c_api_types.h"
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/portable_type_to_tflitetype.h"
#include "tensorflow/lite/util.h"
namespace tflite {
namespace variants {
namespace {
template <typename T>
TensorUniquePtr MakeTensorWithData(std::vector<int> dims,
const std::vector<T>& data) {
TensorUniquePtr tensor =
BuildTfLiteTensor(typeToTfLiteType<T>(), dims, kTfLiteDynamic);
const int num_elements =
std::accumulate(dims.begin(), dims.end(), 1, std::multiplies<int>());
T* data_start = (T*)tensor->data.data;
// TODO(b/257472333) Investigate vector alignment and if this
// can be replaced with memcpy.
for (int i = 0; i < num_elements; ++i) {
data_start[i] = data[i];
}
// For these tests we want to give `TensorArray`s ownership of their
// constituent tensors, so we release here.
return tensor;
}
TensorArray MakeTensorArrayForTest(const std::vector<int>& dims) {
return TensorArray(kTfLiteInt32, BuildTfLiteArray(dims));
}
TEST(TensorArrayTest, InsertSingleElement) {
auto arr = MakeTensorArrayForTest({});
arr.Resize(2);
ASSERT_TRUE(arr.Set(0, MakeTensorWithData<int>({2}, {3, 4})));
const TfLiteTensor* added_tensor = arr.At(0);
ASSERT_TRUE(added_tensor != nullptr);
ASSERT_THAT(added_tensor, DimsAre({2}));
EXPECT_EQ(added_tensor->data.i32[0], 3);
EXPECT_EQ(added_tensor->data.i32[1], 4);
}
TEST(TensorArrayTest, ResizeToZero) {
auto arr = MakeTensorArrayForTest({});
arr.Resize(2);
EXPECT_EQ(arr.NumElements(), 2);
arr.Resize(0);
EXPECT_EQ(arr.NumElements(), 0);
}
TEST(TensorArrayTest, InsertOOB) {
auto arr = MakeTensorArrayForTest({});
TensorUniquePtr tensor = MakeTensorWithData<int>({2}, {3, 4});
arr.Resize(1);
ASSERT_FALSE(arr.Set(-1, std::move(tensor)));
EXPECT_FALSE(arr.At(0));
}
TEST(TensorArrayTest, InsertMultipleElements) {
auto arr = MakeTensorArrayForTest({});
arr.Resize(2);
EXPECT_TRUE(arr.Set(0, MakeTensorWithData<int>({2}, {3, 4})));
EXPECT_TRUE(arr.Set(1, MakeTensorWithData<int>({3}, {3, 4, 5})));
EXPECT_THAT(arr.At(0), DimsAre({2}));
EXPECT_THAT(arr.At(1), DimsAre({3}));
}
TEST(TensorArrayTest, InsertSameIndexTwiceDeletes) {
auto arr = MakeTensorArrayForTest({});
arr.Resize(2);
EXPECT_TRUE(arr.Set(0, MakeTensorWithData<int>({2}, {3, 2})));
EXPECT_TRUE(arr.Set(0, MakeTensorWithData<int>({3}, {3, 4, 5})));
EXPECT_THAT(arr.At(0), DimsAre({3}));
}
TEST(TensorArrayTest, ResizeUpWithElements) {
auto arr = MakeTensorArrayForTest({});
arr.Resize(1);
ASSERT_TRUE(arr.Set(0, MakeTensorWithData<int>({2}, {3, 4})));
arr.Resize(2);
EXPECT_THAT(arr.At(0), DimsAre({2}));
EXPECT_FALSE(arr.At(1));
EXPECT_EQ(arr.NumElements(), 2);
}
// resize down delete elements.
TEST(TensorArrayTest, ResizeDownDeletesElements) {
auto arr = MakeTensorArrayForTest({});
arr.Resize(2);
ASSERT_TRUE(arr.Set(1, MakeTensorWithData<int>({2}, {3, 4})));
arr.Resize(1);
EXPECT_EQ(arr.NumElements(), 1);
EXPECT_FALSE(arr.At(0));
}
TEST(TensorArrayTest, CopyListWithZeroLength) {
auto arr = MakeTensorArrayForTest({});
TensorArray arr2{arr};
EXPECT_EQ(arr.NumElements(), arr2.NumElements());
EXPECT_EQ(arr.NumElements(), 0);
}
TEST(TensorArrayTest, CopyAssignListWithZeroLength) {
auto arr = MakeTensorArrayForTest({});
arr = MakeTensorArrayForTest({2, 2});
EXPECT_EQ(arr.NumElements(), 0);
EXPECT_THAT(arr.ElementShape(), DimsAre({2, 2}));
}
TEST(TensorArrayTest, CopyEmptyList) {
auto arr = MakeTensorArrayForTest({});
arr.Resize(2);
TensorArray arr2{arr};
EXPECT_EQ(arr.NumElements(), arr2.NumElements());
EXPECT_EQ(arr.NumElements(), 2);
}
TEST(TensorArrayTest, CopyAssignToEmptyList) {
auto arr = MakeTensorArrayForTest({});
auto target_arr = MakeTensorArrayForTest({2, 2});
target_arr.Resize(2);
target_arr = arr;
EXPECT_EQ(target_arr.NumElements(), 0);
EXPECT_THAT(target_arr.ElementShape(), DimsAre({}));
}
TEST(TensorArrayTest, CopyListWithItem) {
std::optional<TensorArray> arr = TensorArray(kTfLiteInt32, {});
arr->Resize(1);
ASSERT_TRUE(arr->Set(0, MakeTensorWithData<int>({2}, {3, 4})));
TensorArray arr2{*arr};
EXPECT_EQ(arr->NumElements(), arr2.NumElements());
// Both point to the same tensor.
EXPECT_EQ(arr->At(0), arr2.At(0));
// They are ref counted so deleting one list doesn't effect the tensor
// in the other.
arr.reset();
EXPECT_THAT(arr2.At(0), DimsAre({2}));
}
TEST(TensorArrayTest, CopyAssignToListWithItem) {
auto target_arr = MakeTensorArrayForTest({});
target_arr.Resize(2);
ASSERT_TRUE(target_arr.Set(0, MakeTensorWithData<int>({2}, {3, 4})));
auto src_arr = MakeTensorArrayForTest({2, 2});
src_arr.Resize(1);
target_arr = src_arr;
EXPECT_EQ(target_arr.NumElements(), src_arr.NumElements());
EXPECT_EQ(target_arr.At(0), nullptr);
}
TEST(TensorArrayTest, CopyAssignFromListWithItem) {
auto target_arr = MakeTensorArrayForTest({2, 2});
target_arr.Resize(1);
auto src_arr = MakeTensorArrayForTest({});
src_arr.Resize(2);
ASSERT_TRUE(src_arr.Set(0, MakeTensorWithData<int>({2}, {3, 4})));
target_arr = src_arr;
EXPECT_EQ(target_arr.NumElements(), src_arr.NumElements());
EXPECT_EQ(src_arr.At(0), target_arr.At(0));
}
TEST(TensorArrayTest, DeleteEmptyTensorArray) {
TensorArray* arr = new TensorArray{kTfLiteInt32, {}};
delete arr;
}
TEST(TensorArrayTest, DeleteResizedEmptyTensorArray) {
TensorArray* arr = new TensorArray{kTfLiteInt32, {}};
arr->Resize(2);
delete arr;
}
// OpaqueVariantTensorArrayDataTest(s) test usage of the `TensorArray` through
// the generic interface methods defined in
// `third_party/tensorflow/lite/core/c/common.h`. While appearing slightly
// contrived in function, this test exemplifies proper casting protocol of
// `VariantData` and asserts that the derived methods are dispatched to.
TEST(OpaqueVariantTensorArrayDataTest, CastThroughVoidAndCopy) {
TensorArray* arr = new TensorArray{kTfLiteFloat32, {}};
arr->Resize(2);
ASSERT_TRUE(arr->Set(0, MakeTensorWithData<int>({2}, {3, 4})));
void* erased = static_cast<VariantData*>(arr);
VariantData* d = static_cast<VariantData*>(erased);
VariantData* copied_d = d->CloneTo(nullptr);
auto* copied_arr = static_cast<TensorArray*>(copied_d);
ASSERT_THAT(copied_arr->At(0), DimsAre({2}));
ASSERT_THAT(arr->At(0), DimsAre({2}));
ASSERT_EQ(arr->At(0), arr->At(0));
delete d;
delete copied_d;
}
} // namespace
} // namespace variants
} // namespace tflite