• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /* Copyright 2018 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 
16 #include "tensorflow/core/framework/attr_value_util.h"
17 #include "tensorflow/core/framework/fake_input.h"
18 #include "tensorflow/core/framework/node_def_builder.h"
19 #include "tensorflow/core/framework/shape_inference.h"
20 #include "tensorflow/core/framework/shape_inference_testutil.h"
21 #include "tensorflow/core/framework/tensor.h"
22 #include "tensorflow/core/framework/tensor_shape.h"
23 #include "tensorflow/core/framework/tensor_testutil.h"
24 #include "tensorflow/core/kernels/ops_testutil.h"
25 #include "tensorflow/core/lib/core/status_test_util.h"
26 #include "tensorflow/core/platform/test.h"
27 
28 namespace tensorflow {
29 namespace {
30 
31 template <typename VALUE_TYPE>
32 struct ShapeAndValues {
33   TensorShape shape;
34   std::vector<VALUE_TYPE> values;
35 };
36 
37 template <typename VALUE_TYPE>
createVector(const std::vector<VALUE_TYPE> & values)38 ShapeAndValues<VALUE_TYPE> createVector(const std::vector<VALUE_TYPE>& values) {
39   TensorShape shape({static_cast<int64>(values.size())});
40   return {shape, values};
41 }
42 
43 template <typename VALUE_TYPE>
createScalar(const VALUE_TYPE & values)44 ShapeAndValues<VALUE_TYPE> createScalar(const VALUE_TYPE& values) {
45   TensorShape shape({});
46   return {shape, {values}};
47 }
48 
49 class RaggedTensorToTensorOpTest : public ::tensorflow::OpsTestBase {
50  protected:
51   // Builds the tensorflow test graph for RaggedTensorToTensor.
52   template <typename VALUE_TYPE, typename INDEX_TYPE>
BuildRaggedTensorToTensorGraph(const TensorShape & shape,const std::vector<string> & row_partition_types,const ShapeAndValues<VALUE_TYPE> & values,const ShapeAndValues<VALUE_TYPE> & default_value,const std::vector<ShapeAndValues<INDEX_TYPE>> & row_partition_tensors)53   void BuildRaggedTensorToTensorGraph(
54       const TensorShape& shape, const std::vector<string>& row_partition_types,
55       const ShapeAndValues<VALUE_TYPE>& values,
56       const ShapeAndValues<VALUE_TYPE>& default_value,
57       const std::vector<ShapeAndValues<INDEX_TYPE>>& row_partition_tensors) {
58     const auto& value_dtype = DataTypeToEnum<VALUE_TYPE>::v();
59     const auto& index_dtype = DataTypeToEnum<INDEX_TYPE>::v();
60     int num_row_partition_tensors = row_partition_tensors.size();
61     TF_ASSERT_OK(
62         NodeDefBuilder("tested_op", "RaggedTensorToTensor")
63             .Attr("T", value_dtype)
64             .Attr("Tindex", index_dtype)
65             .Attr("num_row_partition_tensors", num_row_partition_tensors)
66             .Attr("row_partition_types", row_partition_types)
67             .Input(FakeInput(index_dtype))
68             .Input(FakeInput(value_dtype))  // values
69             .Input(FakeInput(value_dtype))  // default_value
70             .Input(FakeInput(num_row_partition_tensors,
71                              index_dtype))  // row_partition_tensors
72             .Finalize(node_def()));
73     TF_ASSERT_OK(InitOp());
74     {
75       std::vector<INDEX_TYPE> shape_as_vector;
76       for (const auto& dim : shape.dim_sizes()) {
77         shape_as_vector.push_back(dim);
78       }
79       ShapeAndValues<INDEX_TYPE> shape_as_tensor =
80           createVector(shape_as_vector);
81       AddInputFromArray<INDEX_TYPE>(shape_as_tensor.shape,
82                                     shape_as_tensor.values);
83     }
84     AddInputFromArray<VALUE_TYPE>(values.shape, values.values);
85     AddInputFromArray<VALUE_TYPE>(default_value.shape, default_value.values);
86 
87     for (const auto& row_partition_tensor : row_partition_tensors) {
88       AddInputFromArray<INDEX_TYPE>(row_partition_tensor.shape,
89                                     row_partition_tensor.values);
90     }
91   }
92 };
93 
TEST_F(RaggedTensorToTensorOpTest,RaggedTensorToTensor)94 TEST_F(RaggedTensorToTensorOpTest, RaggedTensorToTensor) {
95   // indices = [2, 1, 0, 3]
96   // params = [[.1, .2, .3], [], [.4, .5, .6, .7], [.8, .9]]
97   // params.shape = [4, None]
98   BuildRaggedTensorToTensorGraph<float, int32>(
99       TensorShape({4, 4}),                 // shape
100       {"FIRST_DIM_SIZE", "VALUE_ROWIDS"},  // row_partition_types
101       createVector<float>({.1, .2, .3, .4, .5, .6, .7, .8, .9}),  // values
102       createScalar<float>(1.5),  // default_value
103       {createScalar<int32>(4), createVector<int32>({0, 0, 0, 2, 2, 2, 2, 3, 3})}
104       // row_partition_tensors
105   );
106 
107   TF_ASSERT_OK(RunOpKernel());
108 
109   test::ExpectTensorNear<float>(
110       *GetOutput(0),
111       test::AsTensor<float>({.1, .2, .3, 1.5, 1.5, 1.5, 1.5, 1.5, .4, .5, .6,
112                              .7, .8, .9, 1.5, 1.5},
113                             TensorShape({4, 4})),
114       0.01);
115 }
116 
TEST_F(RaggedTensorToTensorOpTest,RaggedTensorToTensorRowSplits)117 TEST_F(RaggedTensorToTensorOpTest, RaggedTensorToTensorRowSplits) {
118   // indices = [2, 1, 0, 3]
119   // params = [[.1, .2, .3], [], [.4, .5, .6, .7], [.8, .9]]
120   BuildRaggedTensorToTensorGraph<float, int32>(
121       TensorShape({4, 4}),  // shape
122       {"ROW_SPLITS"},       // row_partition_types
123       createVector<float>({.1, .2, .3, .4, .5, .6, .7, .8, .9}),  // values
124       createScalar<float>(1.5),               // default_value
125       {createVector<int32>({0, 3, 3, 7, 9})}  // row_partition_tensors
126   );
127 
128   TF_ASSERT_OK(RunOpKernel());
129 
130   test::ExpectTensorNear<float>(
131       *GetOutput(0),
132       test::AsTensor<float>({.1, .2, .3, 1.5, 1.5, 1.5, 1.5, 1.5, .4, .5, .6,
133                              .7, .8, .9, 1.5, 1.5},
134                             TensorShape({4, 4})),
135       0.01);
136 }
137 
TEST_F(RaggedTensorToTensorOpTest,RaggedTensorToTensor_3DParams)138 TEST_F(RaggedTensorToTensorOpTest, RaggedTensorToTensor_3DParams) {
139   // params = [
140   //           [[]],
141   //           [[.1, .2], [.3]],
142   //           [],
143   //           [[.4, .5], [.6, .7, .8]],
144   //           [[.9]]
145   //          ]
146   BuildRaggedTensorToTensorGraph<float, int32>(
147       TensorShape({5, 2, 3}),  // shape
148       {"FIRST_DIM_SIZE", "VALUE_ROWIDS",
149        "VALUE_ROWIDS"},  // row_partition_types
150       createVector<float>({.1, .2, .3, .4, .5, .6, .7, .8, .9}),  // values
151       createScalar<float>(1.5),  // default_value
152       {
153           createScalar<int32>(5),
154           createVector<int32>({0, 1, 1, 3, 3, 4}),
155           createVector<int32>({1, 1, 2, 3, 3, 4, 4, 4, 5}),
156       }  // row_partition_tensors
157   );
158   TF_ASSERT_OK(RunOpKernel());
159 
160   // Expected = [
161   //              [[1.5, 1.5, 1.5], [1.5, 1.5, 1.5]],
162   //              [[.1, .2, 1.5], [.3, 1.5, 1.5]],
163   //              [[1.5, 1.5, 1.5], [1.5, 1.5, 1.5]],
164   //              [[.4, .5, 1.5], [.6, .7, .8]],
165   //              [[.9, 1.5, 1.5], [1.5, 1.5, 1.5]]
166   //            ]
167   test::ExpectTensorNear<float>(
168       *GetOutput(0),
169       test::AsTensor<float>({1.5, 1.5, 1.5, 1.5, 1.5, 1.5, .1,  .2,  1.5, .3,
170                              1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, .4,  .5,
171                              1.5, .6,  .7,  .8,  .9,  1.5, 1.5, 1.5, 1.5, 1.5},
172                             TensorShape({5, 2, 3})),
173       0.1);
174 }
175 
TEST_F(RaggedTensorToTensorOpTest,RaggedTensorToTensor_3DParamsRowSplits)176 TEST_F(RaggedTensorToTensorOpTest, RaggedTensorToTensor_3DParamsRowSplits) {
177   // params = [
178   //           [[]],
179   //           [[.1, .2], [.3]],
180   //           [],
181   //           [[.4, .5], [.6, .7, .8]],
182   //           [[.9]]
183   //          ]
184   BuildRaggedTensorToTensorGraph<float, int32>(
185       TensorShape({5, 2, 3}),        // shape
186       {"ROW_SPLITS", "ROW_SPLITS"},  // row_partition_types
187       createVector<float>({.1, .2, .3, .4, .5, .6, .7, .8, .9}),  // values
188       createScalar<float>(1.5),  // default_value
189       {
190           createVector<int32>({0, 1, 3, 3, 5, 6}),
191           createVector<int32>({0, 0, 2, 3, 5, 8, 9}),
192       }  // row_partition_tensors
193   );
194   TF_ASSERT_OK(RunOpKernel());
195 
196   // Expected = [
197   //              [[1.5, 1.5, 1.5], [1.5, 1.5, 1.5]],
198   //              [[.1, .2, 1.5], [.3, 1.5, 1.5]],
199   //              [[1.5, 1.5, 1.5], [1.5, 1.5, 1.5]],
200   //              [[.4, .5, 1.5], [.6, .7, .8]],
201   //              [[.9, 1.5, 1.5], [1.5, 1.5, 1.5]]
202   //            ]
203   test::ExpectTensorNear<float>(
204       *GetOutput(0),
205       test::AsTensor<float>({1.5, 1.5, 1.5, 1.5, 1.5, 1.5, .1,  .2,  1.5, .3,
206                              1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, .4,  .5,
207                              1.5, .6,  .7,  .8,  .9,  1.5, 1.5, 1.5, 1.5, 1.5},
208                             TensorShape({5, 2, 3})),
209       0.1);
210 }
211 
212 // test_three_dimensional_ragged fails, want to try it at a lower level.
TEST_F(RaggedTensorToTensorOpTest,RaggedTensorToTensor_3DParamsRowSplits2)213 TEST_F(RaggedTensorToTensorOpTest, RaggedTensorToTensor_3DParamsRowSplits2) {
214   // params = [
215   //           [[0, 1, 2], []],
216   //           [],
217   //           [[3]]
218   //          ]
219   BuildRaggedTensorToTensorGraph<int64, int64>(
220       TensorShape({3, 2, 3}),             // shape
221       {"ROW_SPLITS", "ROW_SPLITS"},       // row_partition_types
222       createVector<int64>({0, 1, 2, 3}),  // values
223       createScalar<int64>(5),             // default_value
224       {
225           createVector<int64>({0, 2, 2, 3}),
226           createVector<int64>({0, 3, 3, 4}),
227       }  // row_partition_tensors
228   );
229   TF_ASSERT_OK(RunOpKernel());
230 
231   // Expected = [
232   //              [[0, 1, 2], [5, 5, 5]],
233   //              [[5, 5, 5], [5, 5, 5]],
234   //              [[3, 5, 5], [5, 5, 5]]
235   //            ]
236   test::ExpectTensorEqual<int64>(
237       *GetOutput(0), test::AsTensor<int64>(
238                          {0, 1, 2, 5, 5, 5, 5, 5, 5, 5, 5, 5, 3, 5, 5, 5, 5, 5},
239                          TensorShape({3, 2, 3})));
240 }
241 
TEST_F(RaggedTensorToTensorOpTest,RaggedTensorToTensor_4DParams)242 TEST_F(RaggedTensorToTensorOpTest, RaggedTensorToTensor_4DParams) {
243   // Input:    [[],
244   //            [
245   //             [[1, 2], [3, 4], [5, 6]],
246   //             [[7, 8]]
247   //            ],
248   //            [[]],
249   //            []
250   // ]
251   // params.shape = [3, 2, 3, 2]
252   BuildRaggedTensorToTensorGraph<int32, int32>(
253       TensorShape({4, 2, 3, 2}),  // shape
254       {"FIRST_DIM_SIZE", "VALUE_ROWIDS", "VALUE_ROWIDS",
255        "VALUE_ROWIDS"},                               // row_partition_types
256       createVector<int32>({1, 2, 3, 4, 5, 6, 7, 8}),  // values
257       createScalar<int32>(15),                        // default_value
258       {createScalar<int32>(5), createVector<int32>({0, 1, 1}),
259        createVector<int32>({1, 1, 1, 2}),
260        createVector<int32>({0, 0, 1, 1, 2, 2, 3, 3})}  // row_partition_tensors
261   );
262 
263   TF_ASSERT_OK(RunOpKernel());
264   // params = [
265   //           [
266   //             [[15,15],[15,15],[15,15]],
267   //             [[15,15],[15,15],[15,15]],
268   //           ],
269   //           [
270   //             [[1, 2], [3, 4], [5, 6]],
271   //             [[7, 8], [15, 15], [15,15]],
272   //           ],
273   //             [[15,15],[15,15],[15,15]],
274   //             [[15,15],[15,15],[15,15]],
275   //           ],
276   //             [[15,15],[15,15],[15,15]],
277   //             [[15,15],[15,15],[15,15]],
278   //           ]
279   // params.shape = [3, 2, 3, 2]
280   test::ExpectTensorEqual<int32>(
281       *GetOutput(0),
282       test::AsTensor<int32>(
283           {15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 1,  2,  3,  4,
284            5,  6,  7,  8,  15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15,
285            15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15},
286           TensorShape({4, 2, 3, 2})));
287 }
288 
TEST_F(RaggedTensorToTensorOpTest,RaggedTensorToTensor_4DParamsRowSplit)289 TEST_F(RaggedTensorToTensorOpTest, RaggedTensorToTensor_4DParamsRowSplit) {
290   // Input:    [[],
291   //            [
292   //             [[1, 2], [3, 4], [5, 6]],
293   //             [[7, 8]]
294   //            ],
295   //            [[]],
296   //            []
297   // ]
298   // params.shape = [3, 2, 3, 2]
299   BuildRaggedTensorToTensorGraph<int32, int32>(
300       TensorShape({4, 2, 3, 2}),  // shape
301       {"ROW_SPLITS", "ROW_SPLITS", "ROW_SPLITS"},
302       // row_partition_types
303       createVector<int32>({1, 2, 3, 4, 5, 6, 7, 8}),  // values
304       createScalar<int32>(15),                        // default_value
305       {createVector<int32>({0, 1, 3}), createVector<int32>({0, 0, 3, 4}),
306        createVector<int32>({0, 2, 4, 6, 8})}  // row_partition_tensors
307   );
308 
309   TF_ASSERT_OK(RunOpKernel());
310   // params = [
311   //           [
312   //             [[15,15],[15,15],[15,15]],
313   //             [[15,15],[15,15],[15,15]],
314   //           ],
315   //           [
316   //             [[1, 2], [3, 4], [5, 6]],
317   //             [[7, 8], [15, 15], [15,15]],
318   //           ],
319   //             [[15,15],[15,15],[15,15]],
320   //             [[15,15],[15,15],[15,15]],
321   //           ],
322   //             [[15,15],[15,15],[15,15]],
323   //             [[15,15],[15,15],[15,15]],
324   //           ]
325   // params.shape = [3, 2, 3, 2]
326   test::ExpectTensorEqual<int32>(
327       *GetOutput(0),
328       test::AsTensor<int32>(
329           {15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 1,  2,  3,  4,
330            5,  6,  7,  8,  15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15,
331            15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15},
332           TensorShape({4, 2, 3, 2})));
333 }
334 
TEST_F(RaggedTensorToTensorOpTest,RaggedTensorToTensorContractExpanded)335 TEST_F(RaggedTensorToTensorOpTest, RaggedTensorToTensorContractExpanded) {
336   // params = [[.1, .2, .3], [], [.4, .5, .6, .7], [.8, .9]]
337   BuildRaggedTensorToTensorGraph<float, int32>(
338       TensorShape({3, 5}),                 // shape
339       {"FIRST_DIM_SIZE", "VALUE_ROWIDS"},  // row_partition_types
340       createVector<float>({.1, .2, .3, .4, .5, .6, .7, .8, .9}),  // values
341       createScalar<float>(1.5),  // default_value
342       {createScalar<int32>(4), createVector<int32>({0, 0, 0, 2, 2, 2, 2, 3, 3})}
343       // row_partition_tensors
344   );
345 
346   TF_ASSERT_OK(RunOpKernel());
347 
348   test::ExpectTensorNear<float>(
349       *GetOutput(0),
350       test::AsTensor<float>({.1, .2, .3, 1.5, 1.5,     //
351                              1.5, 1.5, 1.5, 1.5, 1.5,  //
352                              .4, .5, .6, .7, 1.5},     //
353                             TensorShape({3, 5})),
354       0.01);
355 }
356 
357 // Adds a dense dimension.
TEST_F(RaggedTensorToTensorOpTest,RaggedTensorToTensorContractExpandedDense)358 TEST_F(RaggedTensorToTensorOpTest, RaggedTensorToTensorContractExpandedDense) {
359   // params = [[.1, .2, .3], [], [.4, .5, .6, .7], [.8, .9]]
360   BuildRaggedTensorToTensorGraph<float, int32>(
361       TensorShape({3, 5, 2}),              // shape
362       {"FIRST_DIM_SIZE", "VALUE_ROWIDS"},  // row_partition_types
363       ShapeAndValues<float>{TensorShape({9, 2}),
364                             {.1, 1.1, .2, 1.2, .3, 1.3, .4, 1.4, .5, 1.5, .6,
365                              1.6, .7, 1.7, .8, 1.8, .9, 1.9}},  // values
366       createScalar<float>(1.5),                                 // default_value
367       {createScalar<int32>(4), createVector<int32>({0, 0, 0, 2, 2, 2, 2, 3, 3})}
368       // row_partition_tensors
369   );
370 
371   TF_ASSERT_OK(RunOpKernel());
372 
373   test::ExpectTensorNear<float>(
374       *GetOutput(0),
375       test::AsTensor<float>(
376           {.1,  1.1, .2,  1.2, .3,  1.3, 1.5, 1.5, 1.5, 1.5,   //
377            1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5,   //
378            .4,  1.4, .5,  1.5, .6,  1.6, .7,  1.7, 1.5, 1.5},  //
379           TensorShape({3, 5, 2})),
380       0.01);
381 }
382 
TEST_F(RaggedTensorToTensorOpTest,RaggedTensorToTensorConstrained)383 TEST_F(RaggedTensorToTensorOpTest, RaggedTensorToTensorConstrained) {
384   // params = [[.1, .2, .3],
385   //           [],
386   //           [.4, .5, .6, .7],
387   //           [.8, .9]]
388   // constrained to (3, 3)
389   BuildRaggedTensorToTensorGraph<float, int32>(
390       TensorShape({3, 3}),                 // shape
391       {"FIRST_DIM_SIZE", "VALUE_ROWIDS"},  // row_partition_types
392       createVector<float>({.1, .2, .3, .4, .5, .6, .7, .8, .9}),  // values
393       createScalar<float>(1.5),  // default_value
394       {createScalar<int32>(4), createVector<int32>({0, 0, 0, 2, 2, 2, 2, 3, 3})}
395       // row_partition_tensors
396   );
397 
398   TF_ASSERT_OK(RunOpKernel());
399 
400   test::ExpectTensorNear<float>(*GetOutput(0),
401                                 test::AsTensor<float>(
402                                     {
403                                         //
404                                         .1, .2, .3,     //
405                                         1.5, 1.5, 1.5,  //
406                                         .4, .5, .6      //
407                                     },
408                                     TensorShape({3, 3})),
409                                 0.01);
410 }
411 
TEST_F(RaggedTensorToTensorOpTest,RaggedTensorToTensor_3DParamsConstrained)412 TEST_F(RaggedTensorToTensorOpTest, RaggedTensorToTensor_3DParamsConstrained) {
413   // params = [
414   //           [[]],
415   //           [[.1, .2], [.3]],
416   //           [],
417   //           [[.4, .5], [.6, .7, .8]],
418   //           [[.9]]
419   //          ]
420   // params.shape = [5, None, None]
421   BuildRaggedTensorToTensorGraph<float, int32>(
422       TensorShape({4, 1, 2}),  // shape
423       {"FIRST_DIM_SIZE", "VALUE_ROWIDS",
424        "VALUE_ROWIDS"},  // row_partition_types
425       createVector<float>({.1, .2, .3, .4, .5, .6, .7, .8, .9}),  // values
426       createScalar<float>(1.5),  // default_value
427       {
428           createScalar<int32>(5),
429           createVector<int32>({0, 1, 1, 3, 3, 4}),
430           createVector<int32>({1, 1, 2, 3, 3, 4, 4, 4, 5}),
431       }  // row_partition_tensors
432   );
433   TF_ASSERT_OK(RunOpKernel());
434 
435   // Expected = [
436   //              [[1.5, 1.5]],
437   //              [[.1, .2]],
438   //              [[1.5, 1.5]],
439   //              [[.4, .5]],
440   //            ]
441   test::ExpectTensorNear<float>(
442       *GetOutput(0),
443       test::AsTensor<float>({1.5, 1.5, .1, .2, 1.5, 1.5, .4, .5},
444                             TensorShape({4, 1, 2})),
445       0.01);
446 }
447 
448 // Seg fault but removing this does not make the problem go away.
449 // This tests is labeled as flaky. Removing it to find out.
TEST_F(RaggedTensorToTensorOpTest,RaggedTensorToTensor_4DParamsConstrained)450 TEST_F(RaggedTensorToTensorOpTest, RaggedTensorToTensor_4DParamsConstrained) {
451   // Input:    [[],
452   //            [
453   //             [[1, 2], [3, 4], [5, 6]],
454   //             [[7, 8]]
455   //            ],
456   //            [[]],
457   //            []
458   // ]
459   // params.shape = [3, 2, 3, 2]
460   BuildRaggedTensorToTensorGraph<int32, int32>(
461       TensorShape({2, 2, 2, 2}),  // shape
462       {"FIRST_DIM_SIZE", "VALUE_ROWIDS", "VALUE_ROWIDS",
463        "VALUE_ROWIDS"},                               // row_partition_types
464       createVector<int32>({1, 2, 3, 4, 5, 6, 7, 8}),  // values
465       createScalar<int32>(15),                        // default_value
466       {createScalar<int32>(5), createVector<int32>({0, 1, 1}),
467        createVector<int32>({1, 1, 1, 2}),
468        createVector<int32>({0, 0, 1, 1, 2, 2, 3, 3})}  // row_partition_tensors
469   );
470 
471   TF_ASSERT_OK(RunOpKernel());
472   // params = [
473   //           [
474   //             [[15,15],[15,15]],
475   //             [[15,15],[15,15]],
476   //           ],
477   //           [
478   //             [[1, 2], [3, 4]],
479   //             [[7, 8], [15, 15]],
480   //           ],
481   //          ]
482   // params.shape = [3, 2, 3, 2]
483   test::ExpectTensorEqual<int32>(*GetOutput(0), test::AsTensor<int32>(
484                                                     {
485                                                         15, 15, 15, 15,  //
486                                                         15, 15, 15, 15,  //
487                                                         1, 2, 3, 4,      //
488                                                         7, 8, 15, 15,    //
489                                                     },
490                                                     TensorShape({2, 2, 2, 2})));
491 }
492 
TEST_F(RaggedTensorToTensorOpTest,ShapeWrongDimensions)493 TEST_F(RaggedTensorToTensorOpTest, ShapeWrongDimensions) {
494   BuildRaggedTensorToTensorGraph<int32, int32>(
495       TensorShape({10, 7, 10, 20}),  // shape
496       {"FIRST_DIM_SIZE", "VALUE_ROWIDS",
497        "VALUE_ROWIDS"},                   // row_partition_types
498       createVector<int32>({1, 2, 3, 4}),  // values
499       createScalar<int32>(15),            // default_value
500       {createScalar<int32>(5), createVector<int32>({0, 1, 1}),
501        createVector<int32>({1, 1, 1, 2})}  // row_partition_tensors
502   );
503   // Fails with an invalid argument.
504   EXPECT_EQ(RunOpKernel().code(), errors::Code::INVALID_ARGUMENT);
505 }
506 
507 class RaggedTensorToTensorOpUnknownShapeTest
508     : public ::tensorflow::OpsTestBase {
509  protected:
510   std::unique_ptr<ShapeInferenceTestOp> op_;
SetAttributes(const gtl::ArraySlice<string> row_partition_types,int num_row_partition_tensors)511   void SetAttributes(const gtl::ArraySlice<string> row_partition_types,
512                      int num_row_partition_tensors) {
513     op_ = absl::make_unique<ShapeInferenceTestOp>("RaggedTensorToTensor");
514     SetAttrValue(row_partition_types,
515                  &((*op_->node_def.mutable_attr())["row_partition_types"]));
516     (*op_->node_def.mutable_attr())["num_row_partition_tensors"].set_i(
517         num_row_partition_tensors);
518   }
519 };
520 
TEST_F(RaggedTensorToTensorOpUnknownShapeTest,ValueRowIDs)521 TEST_F(RaggedTensorToTensorOpUnknownShapeTest, ValueRowIDs) {
522   SetAttributes(gtl::ArraySlice<string>{"FIRST_DIM_SIZE", "VALUE_ROWIDS"}, 2);
523 
524   INFER_OK(*op_, "?;?;?;?;?", "?");
525   INFER_OK(*op_, "?;[6];[];[];[6]", "[?,?]");
526   INFER_OK(*op_, "?;[6];?;[];[6]", "[?,?]");
527   INFER_OK(*op_, "?;?;[];[];[6]", "?");
528   INFER_OK(*op_, "?;[6];?;[];[6]", "[?,?]");
529   INFER_OK(*op_, "?;[6,2];?;[];[6]", "[?,?,2]");
530   INFER_OK(*op_, "?;[6,2];[2];[];[6]", "[?,?,2]");
531   INFER_OK(*op_, "?;[6,2,7];[2,7];[];[6]", "[?,?,2,7]");
532   INFER_ERROR(
533       "default_value.shape=[3] and rt_input.flat_values.shape=[6,2] "
534       "are incompatible",
535       *op_, "?;[6,2];[3];[];[6]");
536   INFER_ERROR(
537       "default_value.shape=[2,2] and rt_input.flat_values.shape="
538       "[6,2,1,2] are incompatible",
539       *op_, "?;[6,2,1,2];[2,2];[];[6]");
540   INFER_ERROR("must be a vector", *op_, "?;[6];[];[];[3,6]");
541   INFER_ERROR("must be a scalar", *op_, "?;[6];[];[7];[3]");
542 }
543 
TEST_F(RaggedTensorToTensorOpUnknownShapeTest,RowSplits)544 TEST_F(RaggedTensorToTensorOpUnknownShapeTest, RowSplits) {
545   // RaggedTensorToTensor(param_splits+, param_values, indices) -> [splits+,
546   // values]
547   SetAttributes(gtl::ArraySlice<string>{"ROW_SPLITS"}, 1);
548 
549   // value, default_value, ROW_SPLITS
550   INFER_OK(*op_, "?;?;?;?", "?");
551   INFER_OK(*op_, "?;[3];[];[6]", "[?,?]");
552   INFER_OK(*op_, "?;?;?;?", "?");
553   INFER_OK(*op_, "?;[3,2];[2];[6]", "[?,?,2]");
554   INFER_OK(*op_, "?;[3,2,7];[2,7];[6]", "[?,?,2,7]");
555   INFER_OK(*op_, "?;[3,2,7];[2,7];[6]", "[?,?,2,7]");
556 }
557 
558 }  // namespace
559 }  // namespace tensorflow
560