1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 #include <utility>
16 #include <vector>
17
18 #include "absl/strings/match.h"
19 #include "tensorflow/core/framework/fake_input.h"
20 #include "tensorflow/core/framework/node_def_builder.h"
21 #include "tensorflow/core/framework/shape_inference.h"
22 #include "tensorflow/core/framework/shape_inference_testutil.h"
23 #include "tensorflow/core/framework/tensor.h"
24 #include "tensorflow/core/framework/tensor_shape.h"
25 #include "tensorflow/core/framework/tensor_testutil.h"
26 #include "tensorflow/core/framework/variant.h"
27 #include "tensorflow/core/framework/variant_encode_decode.h"
28 #include "tensorflow/core/kernels/ops_testutil.h"
29 #include "tensorflow/core/kernels/ragged_tensor_variant.h"
30 #include "tensorflow/core/lib/core/status_test_util.h"
31 #include "tensorflow/core/platform/test.h"
32
33 namespace tensorflow {
34 namespace {
35
36 class RaggedTensorFromVariantKernelTest : public ::tensorflow::OpsTestBase {
37 protected:
38 // Builds the tensorflow test graph for the RaggedTensorFromVariant op, and
39 // populates the variant input with the given values.
40 template <typename VALUE_TYPE, typename SPLIT_TYPE>
BuildDecodeRaggedTensorGraph(const int input_ragged_rank,const int output_ragged_rank,const TensorShape & variant_shape,const std::vector<Variant> & variant_values)41 void BuildDecodeRaggedTensorGraph(
42 const int input_ragged_rank, const int output_ragged_rank,
43 const TensorShape& variant_shape,
44 const std::vector<Variant>& variant_values) {
45 const auto value_dtype = DataTypeToEnum<VALUE_TYPE>::v();
46 const auto split_dtype = DataTypeToEnum<SPLIT_TYPE>::v();
47 TF_ASSERT_OK(NodeDefBuilder("tested_op", "RaggedTensorFromVariant")
48 .Input(FakeInput(DT_VARIANT))
49 .Attr("input_ragged_rank", input_ragged_rank)
50 .Attr("output_ragged_rank", output_ragged_rank)
51 .Attr("Tvalues", value_dtype)
52 .Attr("Tsplits", split_dtype)
53 .Finalize(node_def()));
54 TF_ASSERT_OK(InitOp());
55 AddInputFromArray<Variant>(variant_shape, variant_values);
56 }
57
58 template <typename VALUE_TYPE, typename SPLIT_TYPE>
CreateVariantFromRagged(const std::vector<std::vector<SPLIT_TYPE>> & ragged_splits,const TensorShape & ragged_values_shape,const std::vector<VALUE_TYPE> & ragged_values)59 RaggedTensorVariant CreateVariantFromRagged(
60 const std::vector<std::vector<SPLIT_TYPE>>& ragged_splits,
61 const TensorShape& ragged_values_shape,
62 const std::vector<VALUE_TYPE>& ragged_values) {
63 RaggedTensorVariant encoded;
64 for (auto ragged_split : ragged_splits) {
65 int splits_size = ragged_split.size();
66 Tensor splits(DataTypeToEnum<SPLIT_TYPE>::v(),
67 TensorShape({splits_size}));
68 test::FillValues<SPLIT_TYPE>(&splits, ragged_split);
69 encoded.append_splits(splits);
70 }
71 Tensor values(DataTypeToEnum<VALUE_TYPE>::v(), ragged_values_shape);
72 test::FillValues<VALUE_TYPE>(&values, ragged_values);
73 encoded.set_values(values);
74 return encoded;
75 }
76 };
77
TEST_F(RaggedTensorFromVariantKernelTest,ScalarInput)78 TEST_F(RaggedTensorFromVariantKernelTest, ScalarInput) {
79 const std::vector<int64> split_1 = {0, 1, 2, 3, 4, 5};
80 const std::vector<int64> split_2 = {0, 1, 2, 5, 6, 7};
81 const std::vector<int> values = {0, 1, 1, 2, 2, 3, 4};
82
83 auto encoded_variant = CreateVariantFromRagged<int, int64>(
84 {split_1, split_2}, TensorShape({7}), values);
85 Tensor expected_splits_1(DT_INT64, TensorShape({6}));
86 Tensor expected_splits_2(DT_INT64, TensorShape({6}));
87 Tensor expected_values(DT_INT32, TensorShape({7}));
88
89 test::FillValues<int64>(&expected_splits_1, split_1);
90 test::FillValues<int64>(&expected_splits_2, split_2);
91 test::FillValues<int>(&expected_values, values);
92
93 int input_ragged_rank = 2;
94 int output_ragged_rank = 2;
95 BuildDecodeRaggedTensorGraph<int, int64>(input_ragged_rank,
96 output_ragged_rank, TensorShape({}),
97 {encoded_variant});
98 TF_ASSERT_OK(RunOpKernel());
99
100 test::ExpectTensorEqual<int64>(*GetOutput(0), expected_splits_1);
101 test::ExpectTensorEqual<int64>(*GetOutput(1), expected_splits_2);
102 test::ExpectTensorEqual<int>(*GetOutput(2), expected_values);
103 }
104
TEST_F(RaggedTensorFromVariantKernelTest,OneInputElement)105 TEST_F(RaggedTensorFromVariantKernelTest, OneInputElement) {
106 const std::vector<int64> split_1 = {0, 1, 2, 3, 4, 5};
107 const std::vector<int64> split_2 = {0, 1, 2, 5, 6, 7};
108 const std::vector<int> values = {0, 1, 1, 2, 2, 3, 4};
109 const std::vector<int64> batched_splits_1 = {0, 5};
110
111 auto encoded_variant = CreateVariantFromRagged<int, int64>(
112 {split_1, split_2}, TensorShape({7}), values);
113 Tensor expected_splits_1(DT_INT64, TensorShape({2}));
114 Tensor expected_splits_2(DT_INT64, TensorShape({6}));
115 Tensor expected_splits_3(DT_INT64, TensorShape({6}));
116 Tensor expected_values(DT_INT32, TensorShape({7}));
117
118 test::FillValues<int64>(&expected_splits_1, batched_splits_1);
119 test::FillValues<int64>(&expected_splits_2, split_1);
120 test::FillValues<int64>(&expected_splits_3, split_2);
121 test::FillValues<int>(&expected_values, values);
122
123 int input_ragged_rank = 2;
124 int output_ragged_rank = 3;
125 BuildDecodeRaggedTensorGraph<int, int64>(input_ragged_rank,
126 output_ragged_rank, TensorShape({1}),
127 {encoded_variant});
128 TF_ASSERT_OK(RunOpKernel());
129
130 test::ExpectTensorEqual<int64>(*GetOutput(0), expected_splits_1);
131 test::ExpectTensorEqual<int64>(*GetOutput(1), expected_splits_2);
132 test::ExpectTensorEqual<int64>(*GetOutput(2), expected_splits_3);
133 test::ExpectTensorEqual<int>(*GetOutput(3), expected_values);
134 }
135
TEST_F(RaggedTensorFromVariantKernelTest,TensorIn2DOut)136 TEST_F(RaggedTensorFromVariantKernelTest, TensorIn2DOut) {
137 // component_1 = [x, x, x]
138 // component_2 = []
139 // component_3 = [x, x]
140 // component_4 = []
141 // batched_ragged =
142 // [[component_1, component_2], [component_3, component_4]]
143 // [
144 // [ [x, x, x], [] ],
145 // [ [x, x], [x] ]
146 // ]
147 const std::vector<int> values_1 = {1, 2, 3};
148 const std::vector<int> values_2 = {};
149 const std::vector<int> values_3 = {4, 5};
150 const std::vector<int> values_4 = {6};
151 const std::vector<int64> batched_splits_1 = {0, 2, 4};
152 const std::vector<int64> batched_splits_2 = {0, 3, 3, 5, 6};
153 const std::vector<int> batched_values = {1, 2, 3, 4, 5, 6};
154
155 auto component_variant_1 =
156 CreateVariantFromRagged<int, int64>({}, TensorShape({3}), values_1);
157 auto component_variant_2 =
158 CreateVariantFromRagged<int, int64>({}, TensorShape({0}), values_2);
159 auto component_variant_3 =
160 CreateVariantFromRagged<int, int64>({}, TensorShape({2}), values_3);
161 auto component_variant_4 =
162 CreateVariantFromRagged<int, int64>({}, TensorShape({1}), values_4);
163
164 Tensor expected_splits_1(DT_INT64, TensorShape({3}));
165 Tensor expected_splits_2(DT_INT64, TensorShape({5}));
166 Tensor expected_values(DT_INT32, TensorShape({6}));
167
168 test::FillValues<int64>(&expected_splits_1, batched_splits_1);
169 test::FillValues<int64>(&expected_splits_2, batched_splits_2);
170 test::FillValues<int>(&expected_values, batched_values);
171
172 int input_ragged_rank = 0;
173 int output_ragged_rank = 2;
174 BuildDecodeRaggedTensorGraph<int, int64>(
175 input_ragged_rank, output_ragged_rank, TensorShape({2, 2}),
176 {component_variant_1, component_variant_2, component_variant_3,
177 component_variant_4});
178 TF_ASSERT_OK(RunOpKernel());
179
180 test::ExpectTensorEqual<int64>(*GetOutput(0), expected_splits_1);
181 test::ExpectTensorEqual<int64>(*GetOutput(1), expected_splits_2);
182 test::ExpectTensorEqual<int>(*GetOutput(2), expected_values);
183 }
184
TEST_F(RaggedTensorFromVariantKernelTest,NonEmpty1DIn3DOut)185 TEST_F(RaggedTensorFromVariantKernelTest, NonEmpty1DIn3DOut) {
186 // ragged_component_1 = [[x]]
187 // ragged_component_2 = [[x], [x]]
188 // ragged_component_3 = [[x, x]]
189 // ragged_component_4 = [[x, x], [x]]
190 // ragged_component_5 = [[x], [x, x]]
191 // batched_ragged = [[rc1, rc2, rc3, rc4, rc5], [rc4, rc5, rc1, rc3, rc2]]
192 const std::vector<int64> component_split_1_1 = {0, 1};
193 const std::vector<int64> component_split_2_1 = {0, 1, 2};
194 const std::vector<int64> component_split_3_1 = {0, 2};
195 const std::vector<int64> component_split_4_1 = {0, 2, 3};
196 const std::vector<int64> component_split_5_1 = {0, 1, 3};
197 const std::vector<int> component_values_1 = {0};
198 const std::vector<int> component_values_2 = {0, 1};
199 const std::vector<int> component_values_3 = {0, 1};
200 const std::vector<int> component_values_4 = {0, 1, 2};
201 const std::vector<int> component_values_5 = {0, 1, 2};
202
203 const std::vector<int64> batched_splits_1 = {0, 5, 10};
204 const std::vector<int64> batched_splits_2 = {0, 1, 3, 4, 6, 8,
205 10, 12, 13, 14, 16};
206 const std::vector<int64> batched_splits_3 = {
207 0, 1, 2, 3, 5, 7, 8, 9, 11, 13, 14, 15, 17, 18, 20, 21, 22};
208 const std::vector<int> batched_values = {0, 0, 1, 0, 1, 0, 1, 2, 0, 1, 2,
209 0, 1, 2, 0, 1, 2, 0, 0, 1, 0, 1};
210
211 Tensor expected_splits_1(DT_INT64, TensorShape({3}));
212 Tensor expected_splits_2(DT_INT64, TensorShape({11}));
213 Tensor expected_splits_3(DT_INT64, TensorShape({17}));
214 Tensor expected_values(DT_INT32, TensorShape({22}));
215
216 test::FillValues<int64>(&expected_splits_1, batched_splits_1);
217 test::FillValues<int64>(&expected_splits_2, batched_splits_2);
218 test::FillValues<int64>(&expected_splits_3, batched_splits_3);
219 test::FillValues<int>(&expected_values, batched_values);
220
221 auto variant_component_1 = CreateVariantFromRagged<int, int64>(
222 {component_split_1_1}, TensorShape({1}), component_values_1);
223 auto variant_component_2 = CreateVariantFromRagged<int, int64>(
224 {component_split_2_1}, TensorShape({2}), component_values_2);
225 auto variant_component_3 = CreateVariantFromRagged<int, int64>(
226 {component_split_3_1}, TensorShape({2}), component_values_3);
227 auto variant_component_4 = CreateVariantFromRagged<int, int64>(
228 {component_split_4_1}, TensorShape({3}), component_values_4);
229 auto variant_component_5 = CreateVariantFromRagged<int, int64>(
230 {component_split_5_1}, TensorShape({3}), component_values_5);
231 int input_ragged_rank = 1;
232 int output_ragged_rank = 3;
233 BuildDecodeRaggedTensorGraph<int, int64>(
234 input_ragged_rank, output_ragged_rank, TensorShape({2, 5}),
235 {variant_component_1, variant_component_2, variant_component_3,
236 variant_component_4, variant_component_5, variant_component_4,
237 variant_component_5, variant_component_1, variant_component_3,
238 variant_component_2});
239 TF_ASSERT_OK(RunOpKernel());
240
241 test::ExpectTensorEqual<int64>(*GetOutput(0), expected_splits_1);
242 test::ExpectTensorEqual<int64>(*GetOutput(1), expected_splits_2);
243 test::ExpectTensorEqual<int64>(*GetOutput(2), expected_splits_3);
244 test::ExpectTensorEqual<int>(*GetOutput(3), expected_values);
245 }
246
TEST_F(RaggedTensorFromVariantKernelTest,NonEmpty2DIn4DOutInferredInputRaggedRank)247 TEST_F(RaggedTensorFromVariantKernelTest,
248 NonEmpty2DIn4DOutInferredInputRaggedRank) {
249 // ragged_component_1 =
250 // [
251 // [ [x] ],
252 // [ [x], [x] ],
253 // [ [x, x] ],
254 // [ [x, x], [x] ],
255 // [ [x], [x, x] ]
256 // ]
257 // ragged_component_2 =
258 // [
259 // [ [x, x], [x] ],
260 // [ [x], [x, x] ],
261 // [ [x] ],
262 // [ [x, x] ],
263 // [ [x], [x] ]
264 // ]
265 // batched_ragged = [[rc1, rc2], [rc2, rc1]]
266 const std::vector<int64> component_split_1_1 = {0, 1, 3, 4, 6, 8};
267 const std::vector<int64> component_split_1_2 = {0, 1, 2, 3, 5, 7, 8, 9, 11};
268 const std::vector<int64> component_split_2_1 = {0, 2, 4, 5, 6, 8};
269 const std::vector<int64> component_split_2_2 = {0, 2, 3, 4, 6, 7, 9, 10, 11};
270 const std::vector<int> component_values_1 = {0, 0, 1, 0, 1, 0, 1, 2, 0, 1, 2};
271 const std::vector<int> component_values_2 = {0, 1, 2, 0, 1, 2, 0, 0, 1, 0, 1};
272 const std::vector<int64> batched_splits_1 = {0, 2, 4};
273 const std::vector<int64> batched_splits_2 = {0, 5, 10, 15, 20};
274 const std::vector<int64> batched_splits_3 = {0, 1, 3, 4, 6, 8, 10,
275 12, 13, 14, 16, 18, 20, 21,
276 22, 24, 25, 27, 28, 30, 32};
277 const std::vector<int64> batched_splits_4 = {
278 0, 1, 2, 3, 5, 7, 8, 9, 11, 13, 14, 15, 17, 18, 20, 21, 22,
279 24, 25, 26, 28, 29, 31, 32, 33, 34, 35, 36, 38, 40, 41, 42, 44};
280 const std::vector<int> batched_values = {
281 0, 0, 1, 0, 1, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 0, 1, 0, 1,
282 0, 1, 2, 0, 1, 2, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 1, 2, 0, 1, 2};
283
284 Tensor expected_splits_1(DT_INT64, TensorShape({3}));
285 Tensor expected_splits_2(DT_INT64, TensorShape({5}));
286 Tensor expected_splits_3(DT_INT64, TensorShape({21}));
287 Tensor expected_splits_4(DT_INT64, TensorShape({33}));
288 Tensor expected_values(DT_INT32, TensorShape({44}));
289 test::FillValues<int64>(&expected_splits_1, batched_splits_1);
290 test::FillValues<int64>(&expected_splits_2, batched_splits_2);
291 test::FillValues<int64>(&expected_splits_3, batched_splits_3);
292 test::FillValues<int64>(&expected_splits_4, batched_splits_4);
293 test::FillValues<int>(&expected_values, batched_values);
294
295 auto variant_component_1 = CreateVariantFromRagged<int, int64>(
296 {component_split_1_1, component_split_1_2}, TensorShape({11}),
297 component_values_1);
298 auto variant_component_2 = CreateVariantFromRagged<int, int64>(
299 {component_split_2_1, component_split_2_2}, TensorShape({11}),
300 component_values_2);
301 int input_ragged_rank = -1;
302 int output_ragged_rank = 4;
303 BuildDecodeRaggedTensorGraph<int, int64>(
304 input_ragged_rank, output_ragged_rank, TensorShape({2, 2}),
305 {variant_component_1, variant_component_2, variant_component_2,
306 variant_component_1});
307 TF_ASSERT_OK(RunOpKernel());
308
309 test::ExpectTensorEqual<int64>(*GetOutput(0), expected_splits_1);
310 test::ExpectTensorEqual<int64>(*GetOutput(1), expected_splits_2);
311 test::ExpectTensorEqual<int64>(*GetOutput(2), expected_splits_3);
312 test::ExpectTensorEqual<int64>(*GetOutput(3), expected_splits_4);
313 test::ExpectTensorEqual<int>(*GetOutput(4), expected_values);
314 }
315
TEST_F(RaggedTensorFromVariantKernelTest,EmptyRow1DIn2DOut)316 TEST_F(RaggedTensorFromVariantKernelTest, EmptyRow1DIn2DOut) {
317 // ragged_component_1 = [[x, x, x], []]
318 // ragged_component_2 = []
319 // batched_ragged = [rc1, rc2] = [[[x, x, x], []], []]
320 const std::vector<int64> component_split_1_1 = {0, 3, 3};
321 const std::vector<int> component_values_1 = {1, 2, 3};
322 const std::vector<int64> component_split_2_1 = {0};
323 const std::vector<int64> batched_splits_1 = {0, 2, 2};
324 const std::vector<int64> batched_splits_2 = {0, 3, 3};
325 const std::vector<int> batched_values = {1, 2, 3};
326
327 Tensor expected_splits_1(DT_INT64, TensorShape({3}));
328 Tensor expected_splits_2(DT_INT64, TensorShape({3}));
329 Tensor expected_values(DT_INT32, TensorShape({3}));
330 test::FillValues<int64>(&expected_splits_1, batched_splits_1);
331 test::FillValues<int64>(&expected_splits_2, batched_splits_2);
332 test::FillValues<int>(&expected_values, batched_values);
333
334 auto variant_component_1 = CreateVariantFromRagged<int, int64>(
335 {component_split_1_1}, TensorShape({3}), component_values_1);
336 auto variant_component_2 = CreateVariantFromRagged<int, int64>(
337 {component_split_2_1}, TensorShape({0}), {}); // Empty row.
338 int input_ragged_rank = 1;
339 int output_ragged_rank = 2;
340 BuildDecodeRaggedTensorGraph<int, int64>(
341 input_ragged_rank, output_ragged_rank, TensorShape({2}),
342 {variant_component_1, variant_component_2});
343 TF_ASSERT_OK(RunOpKernel());
344
345 test::ExpectTensorEqual<int64>(*GetOutput(0), expected_splits_1);
346 test::ExpectTensorEqual<int64>(*GetOutput(1), expected_splits_2);
347 test::ExpectTensorEqual<int>(*GetOutput(2), expected_values);
348 }
349
TEST_F(RaggedTensorFromVariantKernelTest,NDValues1DIn2DOut)350 TEST_F(RaggedTensorFromVariantKernelTest, NDValues1DIn2DOut) {
351 // ragged_component_1 = [[x]]
352 // ragged_component_1 = [[x], [x]]
353 // batched_ragged = [rc1, rc2] = [[[x]], [[x], [x]]]
354 const std::vector<int64> component_split_1_1 = {0, 1};
355 const std::vector<int> component_values_1 = {1, 2};
356 const std::vector<int64> component_split_2_1 = {0, 1, 2};
357 const std::vector<int> component_values_2 = {1, 2, 3, 4};
358 const std::vector<int64> batched_splits_1 = {0, 1, 3};
359 const std::vector<int64> batched_splits_2 = {0, 1, 2, 3};
360 const std::vector<int> batched_values = {1, 2, 1, 2, 3, 4};
361
362 Tensor expected_splits_1(DT_INT64, TensorShape({3}));
363 Tensor expected_splits_2(DT_INT64, TensorShape({4}));
364 Tensor expected_values(DT_INT32, TensorShape({3, 2}));
365 test::FillValues<int64>(&expected_splits_1, batched_splits_1);
366 test::FillValues<int64>(&expected_splits_2, batched_splits_2);
367 test::FillValues<int>(&expected_values, batched_values);
368
369 auto variant_component_1 = CreateVariantFromRagged<int, int64>(
370 {component_split_1_1}, TensorShape({1, 2}), component_values_1);
371 auto variant_component_2 = CreateVariantFromRagged<int, int64>(
372 {component_split_2_1}, TensorShape({2, 2}), component_values_2);
373 int input_ragged_rank = 1;
374 int output_ragged_rank = 2;
375 BuildDecodeRaggedTensorGraph<int, int64>(
376 input_ragged_rank, output_ragged_rank, TensorShape({2}),
377 {variant_component_1, variant_component_2});
378
379 TF_ASSERT_OK(RunOpKernel());
380 test::ExpectTensorEqual<int64>(*GetOutput(0), expected_splits_1);
381 test::ExpectTensorEqual<int64>(*GetOutput(1), expected_splits_2);
382 test::ExpectTensorEqual<int>(*GetOutput(2), expected_values);
383 }
384
TEST_F(RaggedTensorFromVariantKernelTest,NonEmpty1DIn3DOutInt32Splits)385 TEST_F(RaggedTensorFromVariantKernelTest, NonEmpty1DIn3DOutInt32Splits) {
386 // ragged_component_1 = [[x]]
387 // ragged_component_2 = [[x], [x]]
388 // ragged_component_3 = [[x, x]]
389 // ragged_component_4 = [[x, x], [x]]
390 // ragged_component_5 = [[x], [x, x]]
391 // batched_ragged = [[rc1, rc2, rc3, rc4, rc5], [rc4, rc5, rc1, rc3, rc2]]
392 const std::vector<int> component_split_1_1 = {0, 1};
393 const std::vector<int> component_split_2_1 = {0, 1, 2};
394 const std::vector<int> component_split_3_1 = {0, 2};
395 const std::vector<int> component_split_4_1 = {0, 2, 3};
396 const std::vector<int> component_split_5_1 = {0, 1, 3};
397 const std::vector<int> component_values_1 = {0};
398 const std::vector<int> component_values_2 = {0, 1};
399 const std::vector<int> component_values_3 = {0, 1};
400 const std::vector<int> component_values_4 = {0, 1, 2};
401 const std::vector<int> component_values_5 = {0, 1, 2};
402
403 const std::vector<int> batched_splits_1 = {0, 5, 10};
404 const std::vector<int> batched_splits_2 = {0, 1, 3, 4, 6, 8,
405 10, 12, 13, 14, 16};
406 const std::vector<int> batched_splits_3 = {0, 1, 2, 3, 5, 7, 8, 9, 11,
407 13, 14, 15, 17, 18, 20, 21, 22};
408 const std::vector<int> batched_values = {0, 0, 1, 0, 1, 0, 1, 2, 0, 1, 2,
409 0, 1, 2, 0, 1, 2, 0, 0, 1, 0, 1};
410
411 Tensor expected_splits_1(DT_INT32, TensorShape({3}));
412 Tensor expected_splits_2(DT_INT32, TensorShape({11}));
413 Tensor expected_splits_3(DT_INT32, TensorShape({17}));
414 Tensor expected_values(DT_INT32, TensorShape({22}));
415
416 test::FillValues<int>(&expected_splits_1, batched_splits_1);
417 test::FillValues<int>(&expected_splits_2, batched_splits_2);
418 test::FillValues<int>(&expected_splits_3, batched_splits_3);
419 test::FillValues<int>(&expected_values, batched_values);
420
421 auto variant_component_1 = CreateVariantFromRagged<int, int>(
422 {component_split_1_1}, TensorShape({1}), component_values_1);
423 auto variant_component_2 = CreateVariantFromRagged<int, int>(
424 {component_split_2_1}, TensorShape({2}), component_values_2);
425 auto variant_component_3 = CreateVariantFromRagged<int, int>(
426 {component_split_3_1}, TensorShape({2}), component_values_3);
427 auto variant_component_4 = CreateVariantFromRagged<int, int>(
428 {component_split_4_1}, TensorShape({3}), component_values_4);
429 auto variant_component_5 = CreateVariantFromRagged<int, int>(
430 {component_split_5_1}, TensorShape({3}), component_values_5);
431 int input_ragged_rank = 1;
432 int output_ragged_rank = 3;
433 BuildDecodeRaggedTensorGraph<int, int>(
434 input_ragged_rank, output_ragged_rank, TensorShape({2, 5}),
435 {variant_component_1, variant_component_2, variant_component_3,
436 variant_component_4, variant_component_5, variant_component_4,
437 variant_component_5, variant_component_1, variant_component_3,
438 variant_component_2});
439 TF_ASSERT_OK(RunOpKernel());
440
441 test::ExpectTensorEqual<int>(*GetOutput(0), expected_splits_1);
442 test::ExpectTensorEqual<int>(*GetOutput(1), expected_splits_2);
443 test::ExpectTensorEqual<int>(*GetOutput(2), expected_splits_3);
444 test::ExpectTensorEqual<int>(*GetOutput(3), expected_values);
445 }
446
447 // Tests for invalid inputs.
TEST_F(RaggedTensorFromVariantKernelTest,InvalidInferredInputRaggedRank)448 TEST_F(RaggedTensorFromVariantKernelTest, InvalidInferredInputRaggedRank) {
449 auto component_variant_1 =
450 CreateVariantFromRagged<int, int64>({}, TensorShape({3}), {1, 2, 3});
451 auto component_variant_2 =
452 CreateVariantFromRagged<int, int64>({}, TensorShape({0}), {});
453 auto component_variant_3 =
454 CreateVariantFromRagged<int, int64>({}, TensorShape({2}), {1, 2});
455 auto component_variant_4 =
456 CreateVariantFromRagged<int, int64>({}, TensorShape({1}), {1});
457
458 int input_ragged_rank = -1;
459 int output_ragged_rank = 2;
460 BuildDecodeRaggedTensorGraph<int, int64>(
461 input_ragged_rank, output_ragged_rank, TensorShape({1, 1, 1, 4}),
462 {component_variant_1, component_variant_2, component_variant_3,
463 component_variant_4});
464 EXPECT_TRUE(
465 absl::StartsWith(RunOpKernel().error_message(),
466 "Inferred input_ragged_rank (output_ragged_rank - "
467 "encoded_variant.dims()) must be >= 0"));
468 }
469
TEST_F(RaggedTensorFromVariantKernelTest,InputDimsAndRaggedRankAttrsMismatch)470 TEST_F(RaggedTensorFromVariantKernelTest, InputDimsAndRaggedRankAttrsMismatch) {
471 const std::vector<int64> component_split_1_1 = {0, 1};
472 const std::vector<int64> component_split_2_1 = {0, 1, 2};
473 const std::vector<int> component_values_1 = {0};
474 const std::vector<int> component_values_2 = {0, 1};
475
476 auto variant_component_1 = CreateVariantFromRagged<int, int64>(
477 {component_split_1_1}, TensorShape({1}), component_values_1);
478 auto variant_component_2 = CreateVariantFromRagged<int, int64>(
479 {component_split_2_1}, TensorShape({2}), component_values_2);
480
481 int input_ragged_rank = 1;
482 int output_ragged_rank = 4;
483 BuildDecodeRaggedTensorGraph<int, int64>(
484 input_ragged_rank, output_ragged_rank, TensorShape({2}),
485 {variant_component_1, variant_component_2});
486 EXPECT_TRUE(absl::StartsWith(RunOpKernel().error_message(),
487 "output_ragged_rank must be equal to "
488 "input_ragged_rank + encoded_ragged.dims()"));
489 }
490
TEST_F(RaggedTensorFromVariantKernelTest,InputDoesNotHoldRaggedTensorVariant)491 TEST_F(RaggedTensorFromVariantKernelTest, InputDoesNotHoldRaggedTensorVariant) {
492 int input_ragged_rank = 1;
493 int output_ragged_rank = 2;
494 BuildDecodeRaggedTensorGraph<int, int64>(
495 input_ragged_rank, output_ragged_rank, TensorShape({2}), {1, 2});
496 EXPECT_TRUE(absl::StartsWith(
497 RunOpKernel().error_message(),
498 "Input Variant element at index 0 doesn't hold a RaggedTensorVariant"));
499 }
500
TEST_F(RaggedTensorFromVariantKernelTest,InputScalarElementDoesNotMatchInputRaggedRank)501 TEST_F(RaggedTensorFromVariantKernelTest,
502 InputScalarElementDoesNotMatchInputRaggedRank) {
503 const std::vector<int64> component_split_1_1 = {0, 1};
504 const std::vector<int> component_values_1 = {1, 2};
505 auto variant_component_1 = CreateVariantFromRagged<int, int64>(
506 {component_split_1_1}, TensorShape({1, 2}), component_values_1);
507
508 int input_ragged_rank = 2;
509 int output_ragged_rank = 3;
510 BuildDecodeRaggedTensorGraph<int, int64>(input_ragged_rank,
511 output_ragged_rank, TensorShape({1}),
512 {variant_component_1});
513 EXPECT_TRUE(
514 absl::StartsWith(RunOpKernel().error_message(),
515 "Encoded input RaggedTensorVariant has ragged_rank=1. "
516 "Expected ragged_rank=2."));
517 }
518
TEST_F(RaggedTensorFromVariantKernelTest,RaggedSplitTypeMismatch)519 TEST_F(RaggedTensorFromVariantKernelTest, RaggedSplitTypeMismatch) {
520 const std::vector<int64> component_split_1_1 = {0, 1};
521 const std::vector<int> component_values_1 = {0};
522
523 auto variant_component_1 = CreateVariantFromRagged<int, int64>(
524 {component_split_1_1}, TensorShape({1}), component_values_1);
525
526 int input_ragged_rank = 1;
527 int output_ragged_rank = 2;
528 BuildDecodeRaggedTensorGraph<int, int>(input_ragged_rank, output_ragged_rank,
529 TensorShape({1}),
530 {variant_component_1});
531 EXPECT_TRUE(absl::StartsWith(
532 RunOpKernel().error_message(),
533 "Expected row_splits Tensor dtype: int32, found: int64"));
534 }
535
TEST_F(RaggedTensorFromVariantKernelTest,RaggedSplitRankNotOne)536 TEST_F(RaggedTensorFromVariantKernelTest, RaggedSplitRankNotOne) {
537 RaggedTensorVariant encoded(Tensor(DT_INT32, {2}),
538 {Tensor(DT_INT64, {2, 1})});
539 test::FillValues<int64>(encoded.mutable_splits(0), {1, 2});
540 test::FillValues<int>(encoded.mutable_values(), {1, 2});
541
542 int input_ragged_rank = 1;
543 int output_ragged_rank = 2;
544 BuildDecodeRaggedTensorGraph<int, int64>(
545 input_ragged_rank, output_ragged_rank, TensorShape({1}), {encoded});
546 EXPECT_TRUE(absl::StartsWith(RunOpKernel().error_message(),
547 "Ragged splits must have rank 1"));
548 }
549
TEST_F(RaggedTensorFromVariantKernelTest,RaggedValuesTypeMismatch)550 TEST_F(RaggedTensorFromVariantKernelTest, RaggedValuesTypeMismatch) {
551 const std::vector<int64> component_split_1_1 = {0, 1};
552 const std::vector<int> component_values_1 = {0};
553 auto variant_component_1 = CreateVariantFromRagged<int, int64>(
554 {component_split_1_1}, TensorShape({1}), component_values_1);
555 int input_ragged_rank = 1;
556 int output_ragged_rank = 2;
557 BuildDecodeRaggedTensorGraph<tstring, int64>(
558 input_ragged_rank, output_ragged_rank, TensorShape({1}),
559 {variant_component_1});
560 EXPECT_TRUE(
561 absl::StartsWith(RunOpKernel().error_message(),
562 "Expected values Tensor dtype: string, found: int32"));
563 }
564
TEST_F(RaggedTensorFromVariantKernelTest,RaggedValuesRankNotGreaterThanOne)565 TEST_F(RaggedTensorFromVariantKernelTest, RaggedValuesRankNotGreaterThanOne) {
566 auto variant_component_1 =
567 CreateVariantFromRagged<int, int64>({{0, 1}}, TensorShape({}), {1});
568 int input_ragged_rank = 1;
569 int output_ragged_rank = 2;
570 BuildDecodeRaggedTensorGraph<int, int64>(input_ragged_rank,
571 output_ragged_rank, TensorShape({1}),
572 {variant_component_1});
573 EXPECT_TRUE(absl::StartsWith(RunOpKernel().error_message(),
574 "Ragged values must have rank >= 1"));
575 }
576
TEST_F(RaggedTensorFromVariantKernelTest,RaggedValuesRankMismatch)577 TEST_F(RaggedTensorFromVariantKernelTest, RaggedValuesRankMismatch) {
578 const std::vector<int64> component_split_1_1 = {0, 1};
579 const std::vector<int64> component_split_2_1 = {0, 1, 2};
580 const std::vector<int> component_values_1 = {0};
581 const std::vector<int> component_values_2 = {0, 1, 2, 3};
582
583 auto variant_component_1 = CreateVariantFromRagged<int, int64>(
584 {component_split_1_1}, TensorShape({1}), component_values_1);
585 auto variant_component_2 = CreateVariantFromRagged<int, int64>(
586 {component_split_2_1}, TensorShape({2, 2}), component_values_2);
587 int input_ragged_rank = 1;
588 int output_ragged_rank = 2;
589 BuildDecodeRaggedTensorGraph<int, int64>(
590 input_ragged_rank, output_ragged_rank, TensorShape({2}),
591 {variant_component_1, variant_component_2});
592 EXPECT_TRUE(absl::StartsWith(RunOpKernel().error_message(),
593 "Rank of values must match for all components"));
594 }
595
TEST_F(RaggedTensorFromVariantKernelTest,ShapeFnTest)596 TEST_F(RaggedTensorFromVariantKernelTest, ShapeFnTest) {
597 ShapeInferenceTestOp op("RaggedTensorFromVariant");
598
599 // Tests with input_ragged_rank == 0.
600 (*op.node_def.mutable_attr())["input_ragged_rank"].set_i(0);
601 (*op.node_def.mutable_attr())["output_ragged_rank"].set_i(1);
602 INFER_OK(op, "?", "[?];?");
603 INFER_OK(op, "[?]", "[?];?");
604 INFER_ERROR("Shape must be rank 1 but is rank 2", op, "[?,?]");
605
606 // Tests with input_ragged_rank == 1.
607 (*op.node_def.mutable_attr())["input_ragged_rank"].set_i(1);
608
609 (*op.node_def.mutable_attr())["output_ragged_rank"].set_i(1);
610 INFER_OK(op, "?", "[?];?");
611 INFER_ERROR("Shape must be rank 0 but is rank 1", op, "[?]");
612 INFER_ERROR("Shape must be rank 0 but is rank 2", op, "[?,?]");
613
614 (*op.node_def.mutable_attr())["output_ragged_rank"].set_i(2);
615 INFER_OK(op, "?", "[?];[?];?");
616 INFER_OK(op, "[?]", "[?];[?];?");
617 INFER_ERROR("Shape must be rank 1 but is rank 2", op, "[?,?]");
618
619 (*op.node_def.mutable_attr())["output_ragged_rank"].set_i(3);
620 INFER_OK(op, "?", "[?];[?];[?];?");
621 INFER_ERROR("Shape must be rank 2 but is rank 1", op, "[?]");
622 INFER_OK(op, "[?,?]", "[?];[?];[?];?");
623 INFER_ERROR("Shape must be rank 2 but is rank 3", op, "[?,?,?]");
624
625 // Tests with input_ragged_rank == 3.
626 (*op.node_def.mutable_attr())["input_ragged_rank"].set_i(3);
627
628 (*op.node_def.mutable_attr())["output_ragged_rank"].set_i(3);
629 INFER_OK(op, "?", "[?];[?];[?];?");
630 INFER_ERROR("Shape must be rank 0 but is rank 1", op, "[?]");
631
632 (*op.node_def.mutable_attr())["output_ragged_rank"].set_i(4);
633 INFER_OK(op, "?", "[?];[?];[?];[?];?");
634 INFER_OK(op, "[?]", "[?];[?];[?];[?];?");
635 INFER_ERROR("Shape must be rank 1 but is rank 2", op, "[?,?]");
636
637 (*op.node_def.mutable_attr())["output_ragged_rank"].set_i(5);
638 INFER_OK(op, "?", "[?];[?];[?];[?];[?];?");
639 INFER_ERROR("Shape must be rank 2 but is rank 1", op, "[?]");
640 INFER_OK(op, "[?,?]", "[?];[?];[?];[?];[?];?");
641
642 (*op.node_def.mutable_attr())["output_ragged_rank"].set_i(6);
643 INFER_OK(op, "?", "[?];[?];[?];[?];[?];[?];?");
644 INFER_ERROR("Shape must be rank 3 but is rank 1", op, "[?]");
645 INFER_ERROR("Shape must be rank 3 but is rank 2", op, "[?,?]");
646 INFER_OK(op, "[?,?,?]", "[?];[?];[?];[?];[?];[?];?");
647 }
648
649 TEST_F(RaggedTensorFromVariantKernelTest, 2DValuesTensorIn1DOut) {
650 // [
651 // [
652 // [[x, x], [x, x]],
653 // [[x, x], [x, x]]
654 // ],
655 // [[[x, x], [x, x]]],
656 // [],
657 // [
658 // [[x, x], [x, x]],
659 // [[x, x], [x, x]]
660 // ]
661 // ]
662 const std::vector<int64> batched_splits_1 = {0, 2, 3, 3, 5};
663 const std::vector<int> batched_values = {1, 1, 1, 1, 2, 2, 2, 2, 3, 3,
664 3, 3, 4, 4, 4, 4, 5, 5, 5, 5};
665
666 auto variant_component_1 = CreateVariantFromRagged<int, int64>(
667 {}, TensorShape({2, 2, 2}), {1, 1, 1, 1, 2, 2, 2, 2});
668 auto variant_component_2 = CreateVariantFromRagged<int, int64>(
669 {}, TensorShape({1, 2, 2}), {3, 3, 3, 3});
670 auto variant_component_3 =
671 CreateVariantFromRagged<int, int64>({}, TensorShape({0, 2, 2}), {});
672 auto variant_component_4 = CreateVariantFromRagged<int, int64>(
673 {}, TensorShape({2, 2, 2}), {4, 4, 4, 4, 5, 5, 5, 5});
674
675 Tensor expected_splits_1(DT_INT64, TensorShape({5}));
676 Tensor expected_values(DT_INT32, TensorShape({5, 2, 2}));
677 test::FillValues<int64>(&expected_splits_1, batched_splits_1);
678 test::FillValues<int>(&expected_values, batched_values);
679
680 int input_ragged_rank = 0;
681 int output_ragged_rank = 1;
682 BuildDecodeRaggedTensorGraph<int, int64>(
683 input_ragged_rank, output_ragged_rank, TensorShape({4}),
684 {variant_component_1, variant_component_2, variant_component_3,
685 variant_component_4});
686 TF_ASSERT_OK(RunOpKernel());
687
688 test::ExpectTensorEqual<int64>(*GetOutput(0), expected_splits_1);
689 test::ExpectTensorEqual<int>(*GetOutput(1), expected_values);
690 }
691 } // namespace
692 } // namespace tensorflow
693