diff --git a/tensorflow/core/kernels/data/tensor_slice_dataset_op_test.cc b/tensorflow/core/kernels/data/tensor_slice_dataset_op_test.cc index 883440924f0..1f60cd3613d 100644 --- a/tensorflow/core/kernels/data/tensor_slice_dataset_op_test.cc +++ b/tensorflow/core/kernels/data/tensor_slice_dataset_op_test.cc @@ -46,10 +46,10 @@ class TensorSliceDatasetOpTest : public DatasetOpsTestBase { components.emplace_back(strings::StrCat("component_", i)); } - node_def_ = test::function::NDef( + NodeDef node_def = test::function::NDef( kNodeName, kOpName, components, {{"Toutput_types", dtypes}, {"output_shapes", shapes}}); - TF_RETURN_IF_ERROR(CreateOpKernel(node_def_, tensor_dataset_kernel)); + TF_RETURN_IF_ERROR(CreateOpKernel(node_def, tensor_dataset_kernel)); return Status::OK(); } @@ -63,9 +63,6 @@ class TensorSliceDatasetOpTest : public DatasetOpsTestBase { CreateOpKernelContext(tensor_dataset_kernel, inputs, context)); return Status::OK(); } - - private: - NodeDef node_def_; }; struct TestCase { @@ -74,142 +71,189 @@ struct TestCase { std::vector breakpoints; }; -std::vector TestCases() { +TestCase PlainTensorTestCase() { + return {/*components*/ + {DatasetOpsTestBase::CreateTensor(TensorShape({2}), {1, 2}), + DatasetOpsTestBase::CreateTensor(TensorShape({2, 2}), + {1, 2, 3, 4}), + DatasetOpsTestBase::CreateTensor(TensorShape({2, 1}), + {37.0, 38.0}), + DatasetOpsTestBase::CreateTensor(TensorShape({2, 1}), + {"a", "b"})}, + /*expected_outputs*/ + {DatasetOpsTestBase::CreateTensor(TensorShape({}), {1}), + DatasetOpsTestBase::CreateTensor(TensorShape({2}), {1, 2}), + DatasetOpsTestBase::CreateTensor(TensorShape({1}), {37.0}), + DatasetOpsTestBase::CreateTensor(TensorShape({1}), {"a"}), + DatasetOpsTestBase::CreateTensor(TensorShape({}), {2}), + DatasetOpsTestBase::CreateTensor(TensorShape({2}), {3, 4}), + DatasetOpsTestBase::CreateTensor(TensorShape({1}), {38.0}), + DatasetOpsTestBase::CreateTensor(TensorShape({1}), {"b"})}, + /*breakpoints*/ {0, 1, 3}}; +} + +TestCase NestedTensorTestCase() { return { - // A single tuple of tensors. - {{{DatasetOpsTestBase::CreateTensor(TensorShape({2}), {1, 2}), - DatasetOpsTestBase::CreateTensor(TensorShape({2, 2}), - {1, 2, 3, 4}), - DatasetOpsTestBase::CreateTensor(TensorShape({2, 1}), - {37.0, 38.0}), - DatasetOpsTestBase::CreateTensor(TensorShape({2, 1}), - {"a", "b"})}}, // components - {{DatasetOpsTestBase::CreateTensor(TensorShape({}), {1}), - DatasetOpsTestBase::CreateTensor(TensorShape({2}), {1, 2}), - DatasetOpsTestBase::CreateTensor(TensorShape({1}), {37.0}), - DatasetOpsTestBase::CreateTensor(TensorShape({1}), {"a"}), - DatasetOpsTestBase::CreateTensor(TensorShape({}), {2}), - DatasetOpsTestBase::CreateTensor(TensorShape({2}), {3, 4}), - DatasetOpsTestBase::CreateTensor(TensorShape({1}), {38.0}), - DatasetOpsTestBase::CreateTensor(TensorShape({1}), - {"b"})}}, // expected_outputs - {{0, 1, 3}}}, // breakpoints - // Nested tensors - {{{DatasetOpsTestBase::CreateTensor( - TensorShape({2, 1}), - {DatasetOpsTestBase::CreateTensor(TensorShape({2, 2}), - {1.0, 2.0, 3.0, 4.0}), - DatasetOpsTestBase::CreateTensor(TensorShape({2, 2}), - {5.0, 6.0, 7.0, 8.0})}), - DatasetOpsTestBase::CreateTensor( - TensorShape({2, 1}), {DatasetOpsTestBase::CreateTensor( - TensorShape({1, 2}), {"a", "b"}), - DatasetOpsTestBase::CreateTensor( - TensorShape({1, 2}), {"c", "d"})}), - DatasetOpsTestBase::CreateTensor( - TensorShape({2, 3}), {1, 2, 3, 4, 5, 6})}}, // components - {{DatasetOpsTestBase::CreateTensor( - TensorShape({1}), {DatasetOpsTestBase::CreateTensor( - TensorShape({2, 2}), {1.0, 2.0, 3.0, 4.0})}), - DatasetOpsTestBase::CreateTensor( - TensorShape({1}), {DatasetOpsTestBase::CreateTensor( - TensorShape({1, 2}), {"a", "b"})}), - DatasetOpsTestBase::CreateTensor(TensorShape({3}), {1, 2, 3}), - DatasetOpsTestBase::CreateTensor( - TensorShape({1}), {DatasetOpsTestBase::CreateTensor( - TensorShape({2, 2}), {5.0, 6.0, 7.0, 8.0})}), - DatasetOpsTestBase::CreateTensor( - TensorShape({1}), {DatasetOpsTestBase::CreateTensor( - TensorShape({1, 2}), {"c", "d"})}), - DatasetOpsTestBase::CreateTensor( - TensorShape({3}), {4, 5, 6})}}, // expected_outputs - {{0, 1, 2}}} // breakpoints - }; + /*components*/ + {DatasetOpsTestBase::CreateTensor( + TensorShape({2, 1}), + {DatasetOpsTestBase::CreateTensor(TensorShape({2, 2}), + {1.0, 2.0, 3.0, 4.0}), + DatasetOpsTestBase::CreateTensor(TensorShape({2, 2}), + {5.0, 6.0, 7.0, 8.0})}), + DatasetOpsTestBase::CreateTensor( + TensorShape({2, 1}), {DatasetOpsTestBase::CreateTensor( + TensorShape({1, 2}), {"a", "b"}), + DatasetOpsTestBase::CreateTensor( + TensorShape({1, 2}), {"c", "d"})}), + DatasetOpsTestBase::CreateTensor(TensorShape({2, 3}), + {1, 2, 3, 4, 5, 6})}, + /*expected_outputs*/ + {DatasetOpsTestBase::CreateTensor( + TensorShape({1}), {DatasetOpsTestBase::CreateTensor( + TensorShape({2, 2}), {1.0, 2.0, 3.0, 4.0})}), + DatasetOpsTestBase::CreateTensor( + TensorShape({1}), {DatasetOpsTestBase::CreateTensor( + TensorShape({1, 2}), {"a", "b"})}), + DatasetOpsTestBase::CreateTensor(TensorShape({3}), {1, 2, 3}), + DatasetOpsTestBase::CreateTensor( + TensorShape({1}), {DatasetOpsTestBase::CreateTensor( + TensorShape({2, 2}), {5.0, 6.0, 7.0, 8.0})}), + DatasetOpsTestBase::CreateTensor( + TensorShape({1}), {DatasetOpsTestBase::CreateTensor( + TensorShape({1, 2}), {"c", "d"})}), + DatasetOpsTestBase::CreateTensor(TensorShape({3}), {4, 5, 6})}, + /*breakpoints*/ {0, 1, 2}}; } -TEST_F(TensorSliceDatasetOpTest, GetNext) { - int thread_num = 2, cpu_num = 2; - for (auto &test_case : TestCases()) { - std::vector components = test_case.components; - std::vector expected_outputs = test_case.expected_outputs; - size_t num_tensors_per_slice = components.size(); +class ParameterizedTensorSliceDatasetOpTest + : public TensorSliceDatasetOpTest, + public ::testing::WithParamInterface {}; - TF_ASSERT_OK(InitThreadPool(thread_num)); - TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); - - DataTypeVector dtypes; - std::vector shapes; - gtl::InlinedVector inputs; - for (auto &component : components) { - inputs.push_back(&component); - dtypes.push_back(component.dtype()); - } - for (int i = 0; i < num_tensors_per_slice; ++i) { - shapes.emplace_back(expected_outputs[i].shape()); - } - - std::unique_ptr tensor_slice_dataset_kernel; - TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, - &tensor_slice_dataset_kernel)); - std::unique_ptr tensor_slice_dataset_context; - TF_ASSERT_OK(CreateTensorSliceDatasetContext( - tensor_slice_dataset_kernel.get(), &inputs, - &tensor_slice_dataset_context)); - DatasetBase *tensor_slice_dataset; - TF_ASSERT_OK(CreateDataset(tensor_slice_dataset_kernel.get(), - tensor_slice_dataset_context.get(), - &tensor_slice_dataset)); - core::ScopedUnref scored_unref(tensor_slice_dataset); - - std::unique_ptr iterator_context; - TF_ASSERT_OK(CreateIteratorContext(tensor_slice_dataset_context.get(), - &iterator_context)); - std::unique_ptr iterator; - TF_ASSERT_OK(tensor_slice_dataset->MakeIterator(iterator_context.get(), - "Iterator", &iterator)); - bool end_of_sequence = false; - std::vector out_tensors; - int cur_slice = 0; - - while (!end_of_sequence) { - TF_EXPECT_OK(iterator->GetNext(iterator_context.get(), &out_tensors, - &end_of_sequence)); - for (int i = 0; i < out_tensors.size(); ++i) { - EXPECT_LT(i + num_tensors_per_slice * cur_slice, - expected_outputs.size()); - if (out_tensors[i].dtype() == DT_VARIANT) { - // Currently `ExpectEqual()` does not support the variant tensor - // yet, so we manually cast the variant to numeric/string tensor. - const Tensor *output = - out_tensors[i].scalar()().get(); - const Tensor *expected_output = - expected_outputs[i + num_tensors_per_slice * cur_slice] - .scalar()() - .get(); - TF_EXPECT_OK(ExpectEqual(*output, *expected_output)); - } else { - TF_EXPECT_OK(ExpectEqual( - out_tensors[i], - expected_outputs[i + num_tensors_per_slice * cur_slice])); - } - } - out_tensors.clear(); - cur_slice++; - } - } -} - -TEST_F(TensorSliceDatasetOpTest, DatasetName) { +TEST_P(ParameterizedTensorSliceDatasetOpTest, GetNext) { int thread_num = 2, cpu_num = 2; TF_ASSERT_OK(InitThreadPool(thread_num)); TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); - Tensor t1 = CreateTensor(TensorShape({2, 2}), {1, 2, 3, 4}); - Tensor t2 = CreateTensor(TensorShape({2, 2}), {5, 6, 7, 8}); - gtl::InlinedVector inputs = {&t1, &t2}; - DataTypeVector dtypes({DT_INT64, DT_INT64}); - std::vector shapes = {PartialTensorShape({2}), - PartialTensorShape({2})}; + const TestCase &test_case = GetParam(); + const std::vector &expected_outputs = test_case.expected_outputs; + std::vector components = test_case.components; + DataTypeVector dtypes; + gtl::InlinedVector inputs; + for (auto &component : components) { + inputs.emplace_back(&component); + dtypes.emplace_back(component.dtype()); + } + size_t num_tensors_per_slice = components.size(); + std::vector shapes; + shapes.reserve(num_tensors_per_slice); + for (int i = 0; i < num_tensors_per_slice; ++i) { + shapes.emplace_back(expected_outputs[i].shape()); + } + std::unique_ptr tensor_slice_dataset_kernel; + TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, + &tensor_slice_dataset_kernel)); + std::unique_ptr tensor_slice_dataset_context; + TF_ASSERT_OK( + CreateTensorSliceDatasetContext(tensor_slice_dataset_kernel.get(), + &inputs, &tensor_slice_dataset_context)); + DatasetBase *tensor_slice_dataset; + TF_ASSERT_OK(CreateDataset(tensor_slice_dataset_kernel.get(), + tensor_slice_dataset_context.get(), + &tensor_slice_dataset)); + core::ScopedUnref scored_unref(tensor_slice_dataset); + + std::unique_ptr iterator_context; + TF_ASSERT_OK(CreateIteratorContext(tensor_slice_dataset_context.get(), + &iterator_context)); + std::unique_ptr iterator; + TF_ASSERT_OK(tensor_slice_dataset->MakeIterator(iterator_context.get(), + "Iterator", &iterator)); + bool end_of_sequence = false; + std::vector out_tensors; + int cur_slice = 0; + + while (!end_of_sequence) { + TF_EXPECT_OK(iterator->GetNext(iterator_context.get(), &out_tensors, + &end_of_sequence)); + for (int i = 0; i < out_tensors.size(); ++i) { + EXPECT_LT(i + num_tensors_per_slice * cur_slice, expected_outputs.size()); + if (out_tensors[i].dtype() == DT_VARIANT) { + // Currently `ExpectEqual()` does not support the variant tensor + // yet, so we manually cast the variant to numeric/string tensor. + const Tensor *output = out_tensors[i].scalar()().get(); + const Tensor *expected_output = + expected_outputs[i + num_tensors_per_slice * cur_slice] + .scalar()() + .get(); + TF_EXPECT_OK(ExpectEqual(*output, *expected_output)); + } else { + TF_EXPECT_OK(ExpectEqual( + out_tensors[i], + expected_outputs[i + num_tensors_per_slice * cur_slice])); + } + } + out_tensors.clear(); + cur_slice++; + } +} + +TEST_F(TensorSliceDatasetOpTest, DatasetNodeName) { + int thread_num = 2, cpu_num = 2; + TF_ASSERT_OK(InitThreadPool(thread_num)); + TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); + + const TestCase &test_case = PlainTensorTestCase(); + const std::vector &expected_outputs = test_case.expected_outputs; + std::vector components = test_case.components; + DataTypeVector dtypes; + gtl::InlinedVector inputs; + for (auto &component : components) { + inputs.emplace_back(&component); + dtypes.emplace_back(component.dtype()); + } + size_t num_tensors_per_slice = components.size(); + std::vector shapes; + shapes.reserve(num_tensors_per_slice); + for (int i = 0; i < num_tensors_per_slice; ++i) { + shapes.emplace_back(expected_outputs[i].shape()); + } + std::unique_ptr tensor_slice_dataset_kernel; + TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, + &tensor_slice_dataset_kernel)); + std::unique_ptr tensor_slice_dataset_context; + TF_ASSERT_OK( + CreateTensorSliceDatasetContext(tensor_slice_dataset_kernel.get(), + &inputs, &tensor_slice_dataset_context)); + DatasetBase *tensor_slice_dataset; + TF_ASSERT_OK(CreateDataset(tensor_slice_dataset_kernel.get(), + tensor_slice_dataset_context.get(), + &tensor_slice_dataset)); + core::ScopedUnref scored_unref(tensor_slice_dataset); + + EXPECT_EQ(tensor_slice_dataset->node_name(), kNodeName); +} + +TEST_F(TensorSliceDatasetOpTest, DatasetTypeString) { + int thread_num = 2, cpu_num = 2; + TF_ASSERT_OK(InitThreadPool(thread_num)); + TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); + + const TestCase &test_case = PlainTensorTestCase(); + const std::vector &expected_outputs = test_case.expected_outputs; + std::vector components = test_case.components; + DataTypeVector dtypes; + gtl::InlinedVector inputs; + for (auto &component : components) { + inputs.emplace_back(&component); + dtypes.emplace_back(component.dtype()); + } + size_t num_tensors_per_slice = components.size(); + std::vector shapes; + shapes.reserve(num_tensors_per_slice); + for (int i = 0; i < num_tensors_per_slice; ++i) { + shapes.emplace_back(expected_outputs[i].shape()); + } std::unique_ptr tensor_slice_dataset_kernel; TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, &tensor_slice_dataset_kernel)); @@ -226,129 +270,124 @@ TEST_F(TensorSliceDatasetOpTest, DatasetName) { EXPECT_EQ(tensor_slice_dataset->type_string(), kOpName); } -TEST_F(TensorSliceDatasetOpTest, DatasetOutputDtypes) { +TEST_P(ParameterizedTensorSliceDatasetOpTest, DatasetOutputDtypes) { int thread_num = 2, cpu_num = 2; - for (auto &test_case : TestCases()) { - std::vector components = test_case.components; - std::vector expected_outputs = test_case.expected_outputs; - size_t num_tensors_per_slice = components.size(); + TF_ASSERT_OK(InitThreadPool(thread_num)); + TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); - TF_ASSERT_OK(InitThreadPool(thread_num)); - TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); + const TestCase &test_case = GetParam(); + const std::vector &expected_outputs = test_case.expected_outputs; + std::vector components = test_case.components; + DataTypeVector dtypes; + gtl::InlinedVector inputs; + for (auto &component : components) { + inputs.emplace_back(&component); + dtypes.emplace_back(component.dtype()); + } + size_t num_tensors_per_slice = components.size(); + std::vector shapes; + shapes.reserve(num_tensors_per_slice); + for (int i = 0; i < num_tensors_per_slice; ++i) { + shapes.emplace_back(expected_outputs[i].shape()); + } + std::unique_ptr tensor_slice_dataset_kernel; + TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, + &tensor_slice_dataset_kernel)); + std::unique_ptr tensor_slice_dataset_context; + TF_ASSERT_OK( + CreateTensorSliceDatasetContext(tensor_slice_dataset_kernel.get(), + &inputs, &tensor_slice_dataset_context)); + DatasetBase *tensor_slice_dataset; + TF_ASSERT_OK(CreateDataset(tensor_slice_dataset_kernel.get(), + tensor_slice_dataset_context.get(), + &tensor_slice_dataset)); + core::ScopedUnref scored_unref(tensor_slice_dataset); - DataTypeVector dtypes; - std::vector shapes; - gtl::InlinedVector inputs; - for (auto &component : components) { - inputs.emplace_back(&component); - dtypes.emplace_back(component.dtype()); - } - for (int i = 0; i < num_tensors_per_slice; ++i) { - shapes.emplace_back(expected_outputs[i].shape()); - } - - std::unique_ptr tensor_slice_dataset_kernel; - TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, - &tensor_slice_dataset_kernel)); - std::unique_ptr tensor_slice_dataset_context; - TF_ASSERT_OK(CreateTensorSliceDatasetContext( - tensor_slice_dataset_kernel.get(), &inputs, - &tensor_slice_dataset_context)); - DatasetBase *tensor_slice_dataset; - TF_ASSERT_OK(CreateDataset(tensor_slice_dataset_kernel.get(), - tensor_slice_dataset_context.get(), - &tensor_slice_dataset)); - core::ScopedUnref scored_unref(tensor_slice_dataset); - - const DataTypeVector produced_output_dtypes = - tensor_slice_dataset->output_dtypes(); - EXPECT_EQ(produced_output_dtypes.size(), num_tensors_per_slice); - for (int i = 0; i < num_tensors_per_slice; ++i) { - EXPECT_EQ(produced_output_dtypes[i], expected_outputs[i].dtype()); - } + const DataTypeVector produced_output_dtypes = + tensor_slice_dataset->output_dtypes(); + EXPECT_EQ(produced_output_dtypes.size(), num_tensors_per_slice); + for (int i = 0; i < num_tensors_per_slice; ++i) { + EXPECT_EQ(produced_output_dtypes[i], expected_outputs[i].dtype()); } } -TEST_F(TensorSliceDatasetOpTest, DatasetOutputShapes) { +TEST_P(ParameterizedTensorSliceDatasetOpTest, DatasetOutputShapes) { int thread_num = 2, cpu_num = 2; - for (auto &test_case : TestCases()) { - std::vector components = test_case.components; - std::vector expected_outputs = test_case.expected_outputs; - size_t num_tensors_per_slice = components.size(); + TF_ASSERT_OK(InitThreadPool(thread_num)); + TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); - TF_ASSERT_OK(InitThreadPool(thread_num)); - TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); + const TestCase &test_case = GetParam(); + const std::vector &expected_outputs = test_case.expected_outputs; + std::vector components = test_case.components; + DataTypeVector dtypes; + gtl::InlinedVector inputs; + for (auto &component : components) { + inputs.emplace_back(&component); + dtypes.emplace_back(component.dtype()); + } + size_t num_tensors_per_slice = components.size(); + std::vector shapes; + shapes.reserve(num_tensors_per_slice); + for (int i = 0; i < num_tensors_per_slice; ++i) { + shapes.emplace_back(expected_outputs[i].shape()); + } + std::unique_ptr tensor_slice_dataset_kernel; + TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, + &tensor_slice_dataset_kernel)); + std::unique_ptr tensor_slice_dataset_context; + TF_ASSERT_OK( + CreateTensorSliceDatasetContext(tensor_slice_dataset_kernel.get(), + &inputs, &tensor_slice_dataset_context)); + DatasetBase *tensor_slice_dataset; + TF_ASSERT_OK(CreateDataset(tensor_slice_dataset_kernel.get(), + tensor_slice_dataset_context.get(), + &tensor_slice_dataset)); + core::ScopedUnref scored_unref(tensor_slice_dataset); - DataTypeVector dtypes; - std::vector shapes; - gtl::InlinedVector inputs; - for (auto &component : components) { - inputs.emplace_back(&component); - dtypes.emplace_back(component.dtype()); - } - for (int i = 0; i < num_tensors_per_slice; ++i) { - shapes.emplace_back(expected_outputs[i].shape()); - } - std::unique_ptr tensor_slice_dataset_kernel; - TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, - &tensor_slice_dataset_kernel)); - std::unique_ptr tensor_slice_dataset_context; - TF_ASSERT_OK(CreateTensorSliceDatasetContext( - tensor_slice_dataset_kernel.get(), &inputs, - &tensor_slice_dataset_context)); - DatasetBase *tensor_slice_dataset; - TF_ASSERT_OK(CreateDataset(tensor_slice_dataset_kernel.get(), - tensor_slice_dataset_context.get(), - &tensor_slice_dataset)); - core::ScopedUnref scored_unref(tensor_slice_dataset); - - const std::vector produced_output_shapes = - tensor_slice_dataset->output_shapes(); - std::vector expected_output_shapes; - EXPECT_EQ(produced_output_shapes.size(), num_tensors_per_slice); - for (int i = 0; i < num_tensors_per_slice; ++i) { - EXPECT_TRUE( - produced_output_shapes[i].IsIdenticalTo(expected_outputs[i].shape())); - } + const std::vector produced_output_shapes = + tensor_slice_dataset->output_shapes(); + std::vector expected_output_shapes; + EXPECT_EQ(produced_output_shapes.size(), num_tensors_per_slice); + for (int i = 0; i < num_tensors_per_slice; ++i) { + EXPECT_TRUE( + produced_output_shapes[i].IsIdenticalTo(expected_outputs[i].shape())); } } -TEST_F(TensorSliceDatasetOpTest, Cardinality) { +TEST_P(ParameterizedTensorSliceDatasetOpTest, Cardinality) { int thread_num = 2, cpu_num = 2; - for (auto &test_case : TestCases()) { - std::vector components = test_case.components; - std::vector expected_outputs = test_case.expected_outputs; - size_t num_tensors_per_slice = components.size(); + TF_ASSERT_OK(InitThreadPool(thread_num)); + TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); - TF_ASSERT_OK(InitThreadPool(thread_num)); - TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); - - DataTypeVector dtypes; - std::vector shapes; - gtl::InlinedVector inputs; - for (auto &component : components) { - inputs.emplace_back(&component); - dtypes.emplace_back(component.dtype()); - } - for (int i = 0; i < num_tensors_per_slice; ++i) { - shapes.emplace_back(expected_outputs[i].shape()); - } - std::unique_ptr tensor_slice_dataset_kernel; - TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, - &tensor_slice_dataset_kernel)); - std::unique_ptr tensor_slice_dataset_context; - TF_ASSERT_OK(CreateTensorSliceDatasetContext( - tensor_slice_dataset_kernel.get(), &inputs, - &tensor_slice_dataset_context)); - DatasetBase *tensor_slice_dataset; - TF_ASSERT_OK(CreateDataset(tensor_slice_dataset_kernel.get(), - tensor_slice_dataset_context.get(), - &tensor_slice_dataset)); - core::ScopedUnref scored_unref(tensor_slice_dataset); - - EXPECT_EQ(tensor_slice_dataset->Cardinality(), - inputs[0].tensor->dim_size(0)); + const TestCase &test_case = GetParam(); + const std::vector &expected_outputs = test_case.expected_outputs; + std::vector components = test_case.components; + DataTypeVector dtypes; + gtl::InlinedVector inputs; + for (auto &component : components) { + inputs.emplace_back(&component); + dtypes.emplace_back(component.dtype()); } + size_t num_tensors_per_slice = components.size(); + std::vector shapes; + shapes.reserve(num_tensors_per_slice); + for (int i = 0; i < num_tensors_per_slice; ++i) { + shapes.emplace_back(expected_outputs[i].shape()); + } + std::unique_ptr tensor_slice_dataset_kernel; + TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, + &tensor_slice_dataset_kernel)); + std::unique_ptr tensor_slice_dataset_context; + TF_ASSERT_OK( + CreateTensorSliceDatasetContext(tensor_slice_dataset_kernel.get(), + &inputs, &tensor_slice_dataset_context)); + DatasetBase *tensor_slice_dataset; + TF_ASSERT_OK(CreateDataset(tensor_slice_dataset_kernel.get(), + tensor_slice_dataset_context.get(), + &tensor_slice_dataset)); + core::ScopedUnref scored_unref(tensor_slice_dataset); + + EXPECT_EQ(tensor_slice_dataset->Cardinality(), inputs[0].tensor->dim_size(0)); } TEST_F(TensorSliceDatasetOpTest, DatasetSave) { @@ -356,12 +395,21 @@ TEST_F(TensorSliceDatasetOpTest, DatasetSave) { TF_ASSERT_OK(InitThreadPool(thread_num)); TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); - Tensor t1 = CreateTensor(TensorShape({2, 2}), {1, 2, 3, 4}); - Tensor t2 = CreateTensor(TensorShape({2, 2}), {5, 6, 7, 8}); - gtl::InlinedVector inputs = {&t1, &t2}; - DataTypeVector dtypes({DT_INT64, DT_INT64}); - std::vector shapes = {PartialTensorShape({2}), - PartialTensorShape({2})}; + const TestCase &test_case = PlainTensorTestCase(); + const std::vector &expected_outputs = test_case.expected_outputs; + std::vector components = test_case.components; + DataTypeVector dtypes; + gtl::InlinedVector inputs; + for (auto &component : components) { + inputs.emplace_back(&component); + dtypes.emplace_back(component.dtype()); + } + size_t num_tensors_per_slice = components.size(); + std::vector shapes; + shapes.reserve(num_tensors_per_slice); + for (int i = 0; i < num_tensors_per_slice; ++i) { + shapes.emplace_back(expected_outputs[i].shape()); + } std::unique_ptr tensor_slice_dataset_kernel; TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, &tensor_slice_dataset_kernel)); @@ -384,102 +432,98 @@ TEST_F(TensorSliceDatasetOpTest, DatasetSave) { TF_ASSERT_OK(writer.Flush()); } -TEST_F(TensorSliceDatasetOpTest, IteratorOutputDtypes) { +TEST_P(ParameterizedTensorSliceDatasetOpTest, IteratorOutputDtypes) { int thread_num = 2, cpu_num = 2; - for (auto &test_case : TestCases()) { - std::vector components = test_case.components; - std::vector expected_outputs = test_case.expected_outputs; - size_t num_tensors_per_slice = components.size(); + TF_ASSERT_OK(InitThreadPool(thread_num)); + TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); - TF_ASSERT_OK(InitThreadPool(thread_num)); - TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); + const TestCase &test_case = GetParam(); + const std::vector &expected_outputs = test_case.expected_outputs; + std::vector components = test_case.components; + DataTypeVector dtypes; + gtl::InlinedVector inputs; + for (auto &component : components) { + inputs.emplace_back(&component); + dtypes.emplace_back(component.dtype()); + } + size_t num_tensors_per_slice = components.size(); + std::vector shapes; + shapes.reserve(num_tensors_per_slice); + for (int i = 0; i < num_tensors_per_slice; ++i) { + shapes.emplace_back(expected_outputs[i].shape()); + } + std::unique_ptr tensor_slice_dataset_kernel; + TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, + &tensor_slice_dataset_kernel)); + std::unique_ptr tensor_slice_dataset_context; + TF_ASSERT_OK( + CreateTensorSliceDatasetContext(tensor_slice_dataset_kernel.get(), + &inputs, &tensor_slice_dataset_context)); + DatasetBase *tensor_slice_dataset; + TF_ASSERT_OK(CreateDataset(tensor_slice_dataset_kernel.get(), + tensor_slice_dataset_context.get(), + &tensor_slice_dataset)); + core::ScopedUnref scored_unref(tensor_slice_dataset); - DataTypeVector dtypes; - std::vector shapes; - gtl::InlinedVector inputs; - for (auto &component : components) { - inputs.emplace_back(&component); - dtypes.emplace_back(component.dtype()); - } - for (int i = 0; i < num_tensors_per_slice; ++i) { - shapes.emplace_back(expected_outputs[i].shape()); - } + std::unique_ptr iterator_context; + TF_ASSERT_OK(CreateIteratorContext(tensor_slice_dataset_context.get(), + &iterator_context)); + std::unique_ptr iterator; + TF_ASSERT_OK(tensor_slice_dataset->MakeIterator(iterator_context.get(), + "Iterator", &iterator)); + const DataTypeVector produced_output_dtypes = iterator->output_dtypes(); - std::unique_ptr tensor_slice_dataset_kernel; - TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, - &tensor_slice_dataset_kernel)); - std::unique_ptr tensor_slice_dataset_context; - TF_ASSERT_OK(CreateTensorSliceDatasetContext( - tensor_slice_dataset_kernel.get(), &inputs, - &tensor_slice_dataset_context)); - DatasetBase *tensor_slice_dataset; - TF_ASSERT_OK(CreateDataset(tensor_slice_dataset_kernel.get(), - tensor_slice_dataset_context.get(), - &tensor_slice_dataset)); - core::ScopedUnref scored_unref(tensor_slice_dataset); - - std::unique_ptr iterator_context; - TF_ASSERT_OK(CreateIteratorContext(tensor_slice_dataset_context.get(), - &iterator_context)); - std::unique_ptr iterator; - TF_ASSERT_OK(tensor_slice_dataset->MakeIterator(iterator_context.get(), - "Iterator", &iterator)); - const DataTypeVector produced_output_dtypes = iterator->output_dtypes(); - - EXPECT_EQ(produced_output_dtypes.size(), num_tensors_per_slice); - for (int i = 0; i < num_tensors_per_slice; ++i) { - EXPECT_EQ(produced_output_dtypes[i], expected_outputs[i].dtype()); - } + EXPECT_EQ(produced_output_dtypes.size(), num_tensors_per_slice); + for (int i = 0; i < num_tensors_per_slice; ++i) { + EXPECT_EQ(produced_output_dtypes[i], expected_outputs[i].dtype()); } } -TEST_F(TensorSliceDatasetOpTest, IteratorOutputShapes) { +TEST_P(ParameterizedTensorSliceDatasetOpTest, IteratorOutputShapes) { int thread_num = 2, cpu_num = 2; - for (auto &test_case : TestCases()) { - std::vector components = test_case.components; - std::vector expected_outputs = test_case.expected_outputs; - size_t num_tensors_per_slice = components.size(); + TF_ASSERT_OK(InitThreadPool(thread_num)); + TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); - TF_ASSERT_OK(InitThreadPool(thread_num)); - TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); + const TestCase &test_case = GetParam(); + const std::vector &expected_outputs = test_case.expected_outputs; + std::vector components = test_case.components; + DataTypeVector dtypes; + gtl::InlinedVector inputs; + for (auto &component : components) { + inputs.emplace_back(&component); + dtypes.emplace_back(component.dtype()); + } + size_t num_tensors_per_slice = components.size(); + std::vector shapes; + shapes.reserve(num_tensors_per_slice); + for (int i = 0; i < num_tensors_per_slice; ++i) { + shapes.emplace_back(expected_outputs[i].shape()); + } + std::unique_ptr tensor_slice_dataset_kernel; + TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, + &tensor_slice_dataset_kernel)); + std::unique_ptr tensor_slice_dataset_context; + TF_ASSERT_OK( + CreateTensorSliceDatasetContext(tensor_slice_dataset_kernel.get(), + &inputs, &tensor_slice_dataset_context)); + DatasetBase *tensor_slice_dataset; + TF_ASSERT_OK(CreateDataset(tensor_slice_dataset_kernel.get(), + tensor_slice_dataset_context.get(), + &tensor_slice_dataset)); + core::ScopedUnref scored_unref(tensor_slice_dataset); - DataTypeVector dtypes; - std::vector shapes; - gtl::InlinedVector inputs; - for (auto &component : components) { - inputs.emplace_back(&component); - dtypes.emplace_back(component.dtype()); - } - for (int i = 0; i < num_tensors_per_slice; ++i) { - shapes.emplace_back(expected_outputs[i].shape()); - } - - std::unique_ptr tensor_slice_dataset_kernel; - TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, - &tensor_slice_dataset_kernel)); - std::unique_ptr tensor_slice_dataset_context; - TF_ASSERT_OK(CreateTensorSliceDatasetContext( - tensor_slice_dataset_kernel.get(), &inputs, - &tensor_slice_dataset_context)); - DatasetBase *tensor_slice_dataset; - TF_ASSERT_OK(CreateDataset(tensor_slice_dataset_kernel.get(), - tensor_slice_dataset_context.get(), - &tensor_slice_dataset)); - core::ScopedUnref scored_unref(tensor_slice_dataset); - - std::unique_ptr iterator_context; - TF_ASSERT_OK(CreateIteratorContext(tensor_slice_dataset_context.get(), - &iterator_context)); - std::unique_ptr iterator; - TF_ASSERT_OK(tensor_slice_dataset->MakeIterator(iterator_context.get(), - "Iterator", &iterator)); - const std::vector produced_output_shapes = - iterator->output_shapes(); - EXPECT_EQ(produced_output_shapes.size(), num_tensors_per_slice); - for (int i = 0; i < num_tensors_per_slice; ++i) { - EXPECT_TRUE( - produced_output_shapes[i].IsIdenticalTo(expected_outputs[i].shape())); - } + std::unique_ptr iterator_context; + TF_ASSERT_OK(CreateIteratorContext(tensor_slice_dataset_context.get(), + &iterator_context)); + std::unique_ptr iterator; + TF_ASSERT_OK(tensor_slice_dataset->MakeIterator(iterator_context.get(), + "Iterator", &iterator)); + const std::vector produced_output_shapes = + iterator->output_shapes(); + EXPECT_EQ(produced_output_shapes.size(), num_tensors_per_slice); + for (int i = 0; i < num_tensors_per_slice; ++i) { + EXPECT_TRUE( + produced_output_shapes[i].IsIdenticalTo(expected_outputs[i].shape())); } } @@ -488,12 +532,21 @@ TEST_F(TensorSliceDatasetOpTest, IteratorOutputPrefix) { TF_ASSERT_OK(InitThreadPool(thread_num)); TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); - Tensor t1 = CreateTensor(TensorShape({2, 2}), {1, 2, 3, 4}); - Tensor t2 = CreateTensor(TensorShape({2, 2}), {5, 6, 7, 8}); - gtl::InlinedVector inputs = {&t1, &t2}; - DataTypeVector dtypes({DT_INT64, DT_INT64}); - std::vector shapes = {PartialTensorShape({2}), - PartialTensorShape({2})}; + const TestCase &test_case = PlainTensorTestCase(); + const std::vector &expected_outputs = test_case.expected_outputs; + std::vector components = test_case.components; + DataTypeVector dtypes; + gtl::InlinedVector inputs; + for (auto &component : components) { + inputs.emplace_back(&component); + dtypes.emplace_back(component.dtype()); + } + size_t num_tensors_per_slice = components.size(); + std::vector shapes; + shapes.reserve(num_tensors_per_slice); + for (int i = 0; i < num_tensors_per_slice; ++i) { + shapes.emplace_back(expected_outputs[i].shape()); + } std::unique_ptr tensor_slice_dataset_kernel; TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, &tensor_slice_dataset_kernel)); @@ -516,95 +569,96 @@ TEST_F(TensorSliceDatasetOpTest, IteratorOutputPrefix) { EXPECT_EQ(iterator->prefix(), "Iterator::TensorSlice"); } -TEST_F(TensorSliceDatasetOpTest, Roundtrip) { +TEST_P(ParameterizedTensorSliceDatasetOpTest, Roundtrip) { int thread_num = 2, cpu_num = 2; - for (auto &test_case : TestCases()) { - std::vector components = test_case.components; - std::vector expected_outputs = test_case.expected_outputs; - std::vector breakpoints = test_case.breakpoints; - size_t num_tensors_per_slice = components.size(); + TF_ASSERT_OK(InitThreadPool(thread_num)); + TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); - TF_ASSERT_OK(InitThreadPool(thread_num)); - TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num)); + const TestCase &test_case = GetParam(); + const std::vector &expected_outputs = test_case.expected_outputs; + std::vector components = test_case.components; + DataTypeVector dtypes; + gtl::InlinedVector inputs; + for (auto &component : components) { + inputs.emplace_back(&component); + dtypes.emplace_back(component.dtype()); + } + size_t num_tensors_per_slice = components.size(); + std::vector shapes; + shapes.reserve(num_tensors_per_slice); + for (int i = 0; i < num_tensors_per_slice; ++i) { + shapes.emplace_back(expected_outputs[i].shape()); + } + std::unique_ptr tensor_slice_dataset_kernel; + TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, + &tensor_slice_dataset_kernel)); + std::unique_ptr tensor_slice_dataset_context; + TF_ASSERT_OK( + CreateTensorSliceDatasetContext(tensor_slice_dataset_kernel.get(), + &inputs, &tensor_slice_dataset_context)); + DatasetBase *tensor_slice_dataset; + TF_ASSERT_OK(CreateDataset(tensor_slice_dataset_kernel.get(), + tensor_slice_dataset_context.get(), + &tensor_slice_dataset)); + core::ScopedUnref scored_unref(tensor_slice_dataset); - DataTypeVector dtypes; - std::vector shapes; - gtl::InlinedVector inputs; - for (auto &component : components) { - inputs.emplace_back(&component); - dtypes.emplace_back(component.dtype()); - } - for (int i = 0; i < num_tensors_per_slice; ++i) { - shapes.emplace_back(expected_outputs[i].shape()); + std::unique_ptr iterator_context; + TF_ASSERT_OK(CreateIteratorContext(tensor_slice_dataset_context.get(), + &iterator_context)); + std::unique_ptr iterator; + TF_ASSERT_OK(tensor_slice_dataset->MakeIterator(iterator_context.get(), + "Iterator", &iterator)); + std::unique_ptr serialization_context; + TF_ASSERT_OK(CreateSerializationContext(&serialization_context)); + + int cur_iteration = 0; + bool end_of_sequence = false; + int64 num_slices = inputs[0].tensor->dim_size(0); + std::vector out_tensors; + const std::vector &breakpoints = test_case.breakpoints; + for (int breakpoint : breakpoints) { + while (cur_iteration < breakpoint) { + TF_EXPECT_OK(iterator->GetNext(iterator_context.get(), &out_tensors, + &end_of_sequence)); + cur_iteration++; } - std::unique_ptr tensor_slice_dataset_kernel; - TF_ASSERT_OK(CreateTensorSliceDatasetKernel(dtypes, shapes, - &tensor_slice_dataset_kernel)); - std::unique_ptr tensor_slice_dataset_context; - TF_ASSERT_OK(CreateTensorSliceDatasetContext( - tensor_slice_dataset_kernel.get(), &inputs, - &tensor_slice_dataset_context)); - DatasetBase *tensor_slice_dataset; - TF_ASSERT_OK(CreateDataset(tensor_slice_dataset_kernel.get(), - tensor_slice_dataset_context.get(), - &tensor_slice_dataset)); - core::ScopedUnref scored_unref(tensor_slice_dataset); - - std::unique_ptr iterator_context; - TF_ASSERT_OK(CreateIteratorContext(tensor_slice_dataset_context.get(), - &iterator_context)); - std::unique_ptr iterator; - TF_ASSERT_OK(tensor_slice_dataset->MakeIterator(iterator_context.get(), - "Iterator", &iterator)); - std::unique_ptr serialization_context; - TF_ASSERT_OK(CreateSerializationContext(&serialization_context)); - - int cur_iteration = 0; - bool end_of_sequence = false; - int64 num_slices = inputs[0].tensor->dim_size(0); - std::vector out_tensors; - - for (int breakpoint : breakpoints) { - while (cur_iteration < breakpoint) { - TF_EXPECT_OK(iterator->GetNext(iterator_context.get(), &out_tensors, - &end_of_sequence)); - cur_iteration++; - } - - if (breakpoint == 0) { - EXPECT_FALSE(end_of_sequence); - } else if (breakpoint <= num_slices) { - for (int i = 0; i < out_tensors.size(); ++i) { - if (out_tensors[i].dtype() == DT_VARIANT) { - const Tensor *output = - out_tensors[i].scalar()().get(); - const Tensor *expected_output = - expected_outputs[i + - num_tensors_per_slice * (cur_iteration - 1)] - .scalar()() - .get(); - TF_EXPECT_OK(ExpectEqual(*output, *expected_output)); - } else { - TF_EXPECT_OK(ExpectEqual( - out_tensors[i], expected_outputs[i + num_tensors_per_slice * - (cur_iteration - 1)])); - } + if (breakpoint == 0) { + EXPECT_FALSE(end_of_sequence); + } else if (breakpoint <= num_slices) { + for (int i = 0; i < out_tensors.size(); ++i) { + if (out_tensors[i].dtype() == DT_VARIANT) { + const Tensor *output = + out_tensors[i].scalar()().get(); + const Tensor *expected_output = + expected_outputs[i + num_tensors_per_slice * (cur_iteration - 1)] + .scalar()() + .get(); + TF_EXPECT_OK(ExpectEqual(*output, *expected_output)); + } else { + TF_EXPECT_OK(ExpectEqual( + out_tensors[i], expected_outputs[i + num_tensors_per_slice * + (cur_iteration - 1)])); } - } else { - EXPECT_TRUE(end_of_sequence); } - - VariantTensorData data; - VariantTensorDataWriter writer(&data); - TF_ASSERT_OK(iterator->Save(serialization_context.get(), &writer)); - TF_ASSERT_OK(writer.Flush()); - VariantTensorDataReader reader(&data); - TF_ASSERT_OK(iterator->Restore(iterator_context.get(), &reader)); + } else { + EXPECT_TRUE(end_of_sequence); } + + VariantTensorData data; + VariantTensorDataWriter writer(&data); + TF_ASSERT_OK(iterator->Save(serialization_context.get(), &writer)); + TF_ASSERT_OK(writer.Flush()); + VariantTensorDataReader reader(&data); + TF_ASSERT_OK(iterator->Restore(iterator_context.get(), &reader)); } } +INSTANTIATE_TEST_SUITE_P(TensorSliceDatasetOpTest, + ParameterizedTensorSliceDatasetOpTest, + ::testing::ValuesIn(std::vector( + {PlainTensorTestCase(), NestedTensorTestCase()}))); + } // namespace } // namespace data } // namespace tensorflow