1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 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/node_def_builder.h" 17 #include "tensorflow/core/framework/node_def_util.h" 18 #include "tensorflow/core/framework/op.h" 19 #include "tensorflow/core/framework/shape_inference.h" 20 #include "tensorflow/core/framework/shape_inference_testutil.h" 21 #include "tensorflow/core/framework/tensor.pb.h" 22 #include "tensorflow/core/framework/tensor_shape.pb.h" 23 #include "tensorflow/core/framework/tensor_testutil.h" 24 #include "tensorflow/core/lib/core/status_test_util.h" 25 #include "tensorflow/core/platform/test.h" 26 #include "tensorflow/core/public/version.h" 27 28 namespace tensorflow { 29 30 TEST(ArrayOpsTest, UnravelIndex_ShapeFn) { 31 ShapeInferenceTestOp op("UnravelIndex"); 32 33 INFER_OK(op, "?;?", "?"); 34 35 INFER_OK(op, "[];[?]", "[d1_0]"); 36 37 INFER_OK(op, "[4,5];[?]", "[d1_0,20]"); 38 INFER_OK(op, "[2,3,4];[?]", "[d1_0,24]"); 39 INFER_OK(op, "?;[?]", "?"); 40 INFER_OK(op, "[?];[?]", "[d1_0,?]"); 41 42 INFER_ERROR("Shape must be rank 1 but is rank 2", op, "?;[1,1]"); 43 } 44 45 TEST(ArrayOpsTest, Pack_ShapeFn) { 46 ShapeInferenceTestOp op("Pack"); 47 auto set_axis = [&op](int axis) { 48 int n = 3; 49 std::vector<NodeDefBuilder::NodeOut> src_list; 50 src_list.reserve(n); 51 for (int i = 0; i < n; ++i) src_list.emplace_back("a", 0, DT_FLOAT); 52 TF_ASSERT_OK(NodeDefBuilder("test", "Pack") 53 .Input(src_list) 54 .Attr("N", n) 55 .Attr("axis", axis) 56 .Finalize(&op.node_def)); 57 }; 58 59 set_axis(0); 60 INFER_OK(op, "?;?;?", "?"); 61 62 for (int axis : {0, -3}) { 63 set_axis(axis); 64 INFER_OK(op, "?;?;?", "?"); 65 INFER_OK(op, "[1,3];[1,3];?", "[3,d0_0|d1_0,d0_1|d1_1]"); 66 INFER_OK(op, "[?,3];[1,3];?", "[3,d1_0,d0_1|d1_1]"); 67 INFER_OK(op, "[?,?];[1,3];?", "[3,d1_0,d1_1]"); 68 } 69 for (int axis : {1, -2}) { 70 set_axis(axis); 71 INFER_OK(op, "?;?;?", "?"); 72 INFER_OK(op, "[1,3];[1,3];?", "[d0_0|d1_0,3,d0_1|d1_1]"); 73 INFER_OK(op, "[?,3];[1,3];?", "[d1_0,3,d0_1|d1_1]"); 74 INFER_OK(op, "[?,?];[1,3];?", "[d1_0,3,d1_1]"); 75 } 76 for (int axis : {2, -1}) { 77 set_axis(axis); 78 INFER_OK(op, "?;?;?", "?"); 79 INFER_OK(op, "[1,3];[1,3];?", "[d0_0|d1_0,d0_1|d1_1,3]"); 80 INFER_OK(op, "[?,3];[1,3];?", "[d1_0,d0_1|d1_1,3]"); 81 INFER_OK(op, "[?,?];[1,3];?", "[d1_0,d1_1,3]"); 82 } 83 84 set_axis(-4); 85 INFER_ERROR("Invalid axis: -4; must be in [-3,3)", op, "[1,3];[1,3];?"); 86 set_axis(3); 87 INFER_ERROR("Invalid axis: 3; must be in [-3,3)", op, "[1,3];[1,3];?"); 88 89 set_axis(0); 90 91 // Check that both components of error message are there. 92 INFER_ERROR("Shapes must be equal rank, but are 3 and 2", op, 93 "[1,2,3];?;[1,4]"); 94 INFER_ERROR("From merging shape 0 with other shapes.", op, "[1,2,3];?;[1,4]"); 95 } 96 97 TEST(ArrayOpsTest, UnPack_ShapeFn) { 98 ShapeInferenceTestOp op("Unpack"); 99 auto set_axis_and_num = [&op](int axis, int num) { 100 TF_ASSERT_OK(NodeDefBuilder("test", "Unpack") 101 .Input("a", 0, DT_FLOAT) 102 .Attr("axis", axis) 103 .Attr("num", num) 104 .Finalize(&op.node_def)); 105 }; 106 107 set_axis_and_num(0, 1); 108 INFER_OK(op, "?", "?"); 109 110 for (int axis : {0, -3}) { 111 set_axis_and_num(axis, 1); 112 INFER_OK(op, "?", "?"); 113 INFER_OK(op, "[1,2,3]", "[d0_1,d0_2]"); 114 INFER_OK(op, "[?,?,?]", "[d0_1,d0_2]"); 115 } 116 for (int axis : {1, -2}) { 117 set_axis_and_num(axis, 2); 118 INFER_OK(op, "[1,2,3]", "[d0_0,d0_2];[d0_0,d0_2]"); 119 INFER_OK(op, "[?,?,?]", "[d0_0,d0_2];[d0_0,d0_2]"); 120 } 121 for (int axis : {2, -1}) { 122 set_axis_and_num(axis, 3); 123 INFER_OK(op, "[1,2,3]", "[d0_0,d0_1];[d0_0,d0_1];[d0_0,d0_1]"); 124 INFER_OK(op, "[?,?,?]", "[d0_0,d0_1];[d0_0,d0_1];[d0_0,d0_1]"); 125 } 126 127 set_axis_and_num(2, 2); 128 INFER_ERROR("Dimension must be 2 but is 3", op, "[1,2,3]"); 129 130 set_axis_and_num(-4, 3); 131 INFER_ERROR("Invalid axis: -4; must be in [-3,3)", op, "[1,2,3]"); 132 set_axis_and_num(3, 3); 133 INFER_ERROR("Invalid axis: 3; must be in [-3,3)", op, "[1,2,3]"); 134 } 135 136 TEST(ArrayOpsTest, Const_ShapeFn) { 137 ShapeInferenceTestOp op("Const"); 138 TensorProto tensor_proto; 139 auto* shape_proto = tensor_proto.mutable_tensor_shape(); 140 auto rebuild_node_def = [&op, &tensor_proto]() { 141 TF_ASSERT_OK(NodeDefBuilder("test", "Const") 142 .Attr("value", tensor_proto) 143 .Finalize(&op.node_def)); 144 }; 145 146 TensorShape{}.AsProto(shape_proto); 147 rebuild_node_def(); 148 INFER_OK(op, "", "[]"); 149 TensorShape{1, 2, 3, 4}.AsProto(shape_proto); 150 rebuild_node_def(); 151 INFER_OK(op, "", "[1,2,3,4]"); 152 153 shape_proto->add_dim()->set_size(-1); 154 rebuild_node_def(); 155 INFER_ERROR("Shape [1,2,3,4,?] is not fully defined", op, ""); 156 } 157 158 TEST(ArrayOpsTest, UnchangedShapes_ShapeFn) { 159 for (const char* op_name : { 160 "CheckNumerics", 161 "Identity", 162 "RefIdentity", 163 "QuantizeAndDequantize", 164 "StopGradient", 165 "ZerosLike", 166 "OnesLike", 167 }) { 168 ShapeInferenceTestOp op(op_name); 169 INFER_OK(op, "?", "in0"); 170 INFER_OK(op, "[]", "in0"); 171 INFER_OK(op, "[1,2,?,4,5]", "in0"); 172 } 173 174 // inputs 1 and 2 are ignored; input 0 is transferred to output 0. 175 ShapeInferenceTestOp op("MatrixBandPart"); 176 INFER_OK(op, "?;?;?", "in0"); 177 INFER_OK(op, "[];?;?", "in0"); 178 INFER_OK(op, "[1,2,?,4,5];?;?", "in0"); 179 } 180 181 TEST(ArrayOpsTest, GuaranteeConst_ShapeFn) { 182 ShapeInferenceTestOp op("GuaranteeConst"); 183 INFER_OK(op, "?", "in0"); 184 INFER_OK(op, "[]", "in0"); 185 INFER_OK(op, "[1,2,?,4,5]", "in0"); 186 } 187 188 TEST(ArrayOpsTest, Identity_ShapeFnHandles) { 189 const char* op_name = "Identity"; 190 ShapeInferenceTestOp op(op_name); 191 // Check that handle dtypes are preserved. 192 const OpRegistrationData* op_reg_data; 193 TF_ASSERT_OK(OpRegistry::Global()->LookUp(op.name, &op_reg_data)); 194 std::vector< 195 std::unique_ptr<std::vector<std::pair<TensorShapeProto, DataType>>>> 196 handle_data; 197 handle_data.emplace_back( 198 new std::vector<std::pair<TensorShapeProto, DataType>>{ 199 {TensorShapeProto(), DT_BOOL}}); 200 shape_inference::InferenceContext c(TF_GRAPH_DEF_VERSION, &op.node_def, 201 op_reg_data->op_def, {TensorShapeProto()}, 202 {}, {}, handle_data); 203 TF_ASSERT_OK(c.construction_status()); 204 ASSERT_TRUE(op_reg_data->shape_inference_fn != nullptr); 205 TF_ASSERT_OK(c.Run(op_reg_data->shape_inference_fn)); 206 207 const auto* shapes_and_types = c.output_handle_shapes_and_types(0); 208 ASSERT_TRUE(shapes_and_types != nullptr); 209 ASSERT_EQ(1, shapes_and_types->size()); 210 EXPECT_EQ((*shapes_and_types)[0].dtype, DT_BOOL); 211 } 212 213 TEST(ArrayOpsTest, Diag_ShapeFn) { 214 ShapeInferenceTestOp op("Diag"); 215 INFER_OK(op, "?", "?"); 216 INFER_OK(op, "[1,?,3]", "[d0_0,d0_1,d0_2,d0_0,d0_1,d0_2]"); 217 INFER_OK(op, "[?,1,2,3]", "[d0_0,d0_1,d0_2,d0_3,d0_0,d0_1,d0_2,d0_3]"); 218 INFER_ERROR("Shape must be at least rank 1 but is rank 0", op, "[]"); 219 } 220 221 TEST(ArrayOpsTest, DiagPart_ShapeFn) { 222 ShapeInferenceTestOp op("DiagPart"); 223 INFER_OK(op, "?", "?"); 224 INFER_OK(op, "[1,?,?,4]", "[d0_0,d0_3]"); 225 INFER_OK(op, "[1,?,3,?,4,3]", "[d0_0,d0_4,d0_2|d0_5]"); 226 INFER_OK(op, "[1,2,3,?,?,?,?,4]", "[d0_0,d0_1,d0_2,d0_7]"); 227 INFER_ERROR("Input must have even and non-zero rank", op, "[]"); 228 INFER_ERROR("Input must have even and non-zero rank", op, "[?]"); 229 INFER_ERROR("Input must have even and non-zero rank", op, "[1,2,3]"); 230 INFER_ERROR("Dimensions must be equal, but are 2 and 10", op, "[1,2,?,10]"); 231 } 232 233 TEST(ArrayOpsTest, MatrixDiag_ShapeFn) { 234 ShapeInferenceTestOp op("MatrixDiag"); 235 INFER_OK(op, "?", "?"); 236 INFER_ERROR("Shape must be at least rank 1 but is rank 0", op, "[]"); 237 INFER_OK(op, "[?]", "[d0_0,d0_0]"); 238 INFER_OK(op, "[1,?,?,4]", "[d0_0,d0_1,d0_2,d0_3,d0_3]"); 239 } 240 241 TEST(ArrayOpsTest, MatrixDiagPart_ShapeFn) { 242 ShapeInferenceTestOp op("MatrixDiagPart"); 243 INFER_OK(op, "?", "?"); 244 INFER_ERROR("Shape must be at least rank 2 but is rank 1", op, "[?]"); 245 INFER_OK(op, "[?,1,2,2]", "[d0_0,d0_1,d0_2|d0_3]"); 246 INFER_OK(op, "[?,1,2,3]", "[d0_0,d0_1,d0_2]"); 247 INFER_OK(op, "[?,1,3,2]", "[d0_0,d0_1,d0_3]"); 248 } 249 250 TEST(ArrayOpsTest, Reverse_ShapeFn) { 251 ShapeInferenceTestOp op("Reverse"); 252 INFER_OK(op, "?;?", "in0"); 253 INFER_ERROR("Shape must be rank 1 but is rank 0", op, "?;[]"); 254 INFER_ERROR("Shape must be rank 1 but is rank 2", op, "?;[?,2]"); 255 INFER_ERROR("Shape must be rank 4 but is rank 3", op, "[1,2,3];[4]"); 256 INFER_ERROR("reverse does not work on tensors with more than 8 dimensions", 257 op, "[1,2,3,4,5,6,7,8,9];[9]"); 258 INFER_OK(op, "[1,2,3,?];[4]", "in0"); 259 INFER_OK(op, "[1,2,3,?,5,6,7,8];[8]", "in0"); 260 } 261 262 TEST(ArrayOpsTest, ReverseV2_ShapeFn) { 263 ShapeInferenceTestOp op("ReverseV2"); 264 INFER_OK(op, "?;?", "in0"); 265 INFER_ERROR("Shape must be rank 1 but is rank 0", op, "?;[]"); 266 INFER_ERROR("Shape must be rank 1 but is rank 2", op, "?;[?,2]"); 267 INFER_OK(op, "[1,2,3];[2]", "in0"); 268 INFER_ERROR("reverse does not work on tensors with more than 8 dimensions", 269 op, "[1,2,3,4,5,6,7,8,9];[9]"); 270 INFER_OK(op, "[1,2,3,?];[4]", "in0"); 271 INFER_OK(op, "[1,2,3,?,5,6,7,8];[8]", "in0"); 272 } 273 274 TEST(ArrayOpsTest, Fill_ShapeFn) { 275 ShapeInferenceTestOp op("Fill"); 276 AddNodeAttr("index_type", DT_INT32, &op.node_def); 277 op.input_tensors.resize(2); 278 INFER_OK(op, "?;?", "?"); 279 INFER_OK(op, "[?];?", "?"); 280 INFER_OK(op, "[4];?", "[?,?,?,?]"); 281 282 Tensor in_t = test::AsTensor<int32>({1, 2, 3, 4}); 283 op.input_tensors[0] = &in_t; 284 INFER_OK(op, "[4];?", "[1,2,3,4]"); 285 } 286 287 TEST(ArrayOpsTest, Gather_ShapeFn) { 288 ShapeInferenceTestOp op("Gather"); 289 INFER_OK(op, "?;?", "?"); 290 INFER_OK(op, "[1,?,2];[3]", "[d1_0,d0_1,d0_2]"); 291 INFER_ERROR("Shape must be at least rank 1 but is rank 0", op, "[];[1,2,3]"); 292 } 293 294 TEST(ArrayOpsTest, GatherV2_ShapeFn) { 295 ShapeInferenceTestOp op("GatherV2"); 296 297 // Tests when axis is unknown. 298 INFER_OK(op, "?;?;?", "?"); 299 INFER_OK(op, "[1,2,3];[3];[]", "[?,?,?]"); 300 INFER_ERROR("Shape must be at least rank 1 but is rank 0", op, 301 "[];[1,2,3];[]"); 302 303 // Non-scalar axis. 304 INFER_ERROR("Shape must be rank 0 but is rank 1", op, "[1];[1,2,3];[1]"); 305 306 // Test when axis dim is known. 307 Tensor axis_dim_t; 308 op.input_tensors.resize(3); 309 op.input_tensors[2] = &axis_dim_t; 310 311 // Out of range axis. 312 axis_dim_t = test::AsScalar(1); 313 INFER_ERROR("Shape must be at least rank 2 but is rank 1", op, 314 "[1];[1,2];[]"); 315 316 // Rank 0 indices. 317 axis_dim_t = test::AsScalar(0); 318 INFER_OK(op, "[1,2,3];[];[]", "[d0_1,d0_2]"); 319 axis_dim_t = test::AsScalar(1); 320 INFER_OK(op, "[1,2,3];[];[]", "[d0_0,d0_2]"); 321 axis_dim_t = test::AsScalar(2); 322 INFER_OK(op, "[1,2,3];[];[]", "[d0_0,d0_1]"); 323 324 // Rank 1 indices. 325 axis_dim_t = test::AsScalar(0); 326 INFER_OK(op, "[1,2,3];[5];[]", "[d1_0,d0_1,d0_2]"); 327 axis_dim_t = test::AsScalar(1); 328 INFER_OK(op, "[1,2,3];[5];[]", "[d0_0,d1_0,d0_2]"); 329 axis_dim_t = test::AsScalar(2); 330 INFER_OK(op, "[1,2,3];[5];[]", "[d0_0,d0_1,d1_0]"); 331 332 // Rank 2 indices. 333 axis_dim_t = test::AsScalar(0); 334 INFER_OK(op, "[1,2,3];[5,6];[]", "[d1_0,d1_1,d0_1,d0_2]"); 335 axis_dim_t = test::AsScalar(1); 336 INFER_OK(op, "[1,2,3];[5,6];[]", "[d0_0,d1_0,d1_1,d0_2]"); 337 axis_dim_t = test::AsScalar(2); 338 INFER_OK(op, "[1,2,3];[5,6];[]", "[d0_0,d0_1,d1_0,d1_1]"); 339 340 // Negative axis. 341 axis_dim_t = test::AsScalar(-3); 342 INFER_OK(op, "[1,2,3];[5,6];[]", "[d1_0,d1_1,d0_1,d0_2]"); 343 axis_dim_t = test::AsScalar(-2); 344 INFER_OK(op, "[1,2,3];[5,6];[]", "[d0_0,d1_0,d1_1,d0_2]"); 345 axis_dim_t = test::AsScalar(-1); 346 INFER_OK(op, "[1,2,3];[5,6];[]", "[d0_0,d0_1,d1_0,d1_1]"); 347 } 348 349 TEST(ArrayOpsTest, GatherNd_ShapeFn) { 350 ShapeInferenceTestOp op("GatherNd"); 351 352 // Inputs are (params, indices). 353 INFER_OK(op, "?;?", "?"); 354 INFER_OK(op, "[1,?,3,?];[?,0]", "[d1_0,d0_0,d0_1,d0_2,d0_3]"); 355 INFER_OK(op, "[1,?,3,?];[?,4]", "[d1_0]"); 356 357 // params.rank >= indices.dim(-1). 358 INFER_ERROR("indices.shape[-1] must be <= params.rank", op, "[1,2,3];[4]"); 359 } 360 361 TEST(ArrayOpsTest, Shape_ShapeFn) { 362 ShapeInferenceTestOp op("Shape"); 363 INFER_OK(op, "?", "[?]"); 364 INFER_OK(op, "[?]", "[1]"); 365 INFER_OK(op, "[?,2,3,4,5]", "[5]"); 366 } 367 368 TEST(ArrayOpsTest, ShapeN_ShapeFn) { 369 ShapeInferenceTestOp op("ShapeN"); 370 int n = 3; 371 std::vector<NodeDefBuilder::NodeOut> src_list; 372 src_list.reserve(n); 373 for (int i = 0; i < n; ++i) src_list.emplace_back("a", 0, DT_FLOAT); 374 TF_ASSERT_OK(NodeDefBuilder("test", "ShapeN") 375 .Input(src_list) 376 .Attr("N", n) 377 .Finalize(&op.node_def)); 378 INFER_OK(op, "?;?;?", "[?];[?];[?]"); 379 INFER_OK(op, "[?];[?];[?]", "[1];[1];[1]"); 380 INFER_OK(op, "[?,2,3,4,5];?;[1,?,3]", "[5];[?];[3]"); 381 } 382 383 TEST(ArrayOpsTest, Unique_ShapeFn) { 384 ShapeInferenceTestOp op("Unique"); 385 INFER_OK(op, "?", "[?];in0"); 386 INFER_OK(op, "[5]", "[?];in0"); 387 INFER_ERROR( 388 "Shape must be rank 1 but is rank 5 for '' (op: '') with input shapes: " 389 "[1,2,3,?,5].", 390 op, "[1,2,3,?,5]"); 391 } 392 393 TEST(ArrayOpsTest, UniqueWithCounts_ShapeFn) { 394 ShapeInferenceTestOp op("UniqueWithCounts"); 395 INFER_OK(op, "?", "[?];in0;[?]"); 396 INFER_OK(op, "[1,2,3,?,5]", "[?];in0;[?]"); 397 } 398 399 TEST(ArrayOpsTest, InvertPermutation_ShapeFn) { 400 ShapeInferenceTestOp op("InvertPermutation"); 401 INFER_OK(op, "?", "[?]"); 402 INFER_OK(op, "[1]", "in0"); 403 INFER_ERROR("Shape must be rank 1 but is rank 0", op, "[]"); 404 } 405 406 TEST(ArrayOpsTest, PadD_ShapeFn) { 407 for (const char* op_name : {"Pad", "MirrorPad"}) { 408 ShapeInferenceTestOp op(op_name); 409 op.input_tensors.resize(2); 410 411 // Inputs are input and paddings. 412 413 INFER_OK(op, "?;?", "?"); 414 415 // Check shape of paddings. 416 INFER_ERROR("Shape must be rank 2 but is rank 3", op, "?;[1,2,3]"); 417 INFER_ERROR("Dimension must be 2 but is 4", op, "?;[1,4]"); 418 419 // input.rank and paddings.dim(0) are equal. This is the number of dims in 420 // output. 421 INFER_ERROR("Shape must be rank 4 but is rank 3", op, "[1,2,3];[4,2]"); 422 INFER_OK(op, "[1,2,3];?", "[?,?,?]"); 423 INFER_OK(op, "?;[3,2]", "[?,?,?]"); 424 425 // Make the paddings tensor known and verify padding values get added. 426 // E.g., if padding is ((1,10),(2,20),(3,30)) then values 11,22,23 are added 427 // to input dims to get output. 428 Tensor paddings_t(DT_INT64, TensorShape{3, 2}); 429 test::FillValues<int64>(&paddings_t, {1, 10, 2, 20, 3, 30}); 430 op.input_tensors[1] = &paddings_t; 431 INFER_OK(op, "[100,200,300];[3,2]", "[111,222,333]"); 432 INFER_OK(op, "[100,?,300];[3,2]", "[111,?,333]"); 433 INFER_OK(op, "?;[3,2]", "[?,?,?]"); 434 INFER_OK(op, "?;?", "[?,?,?]"); 435 } 436 } 437 438 TEST(ArrayOpsTest, PadV2_ShapeFn) { 439 ShapeInferenceTestOp op("PadV2"); 440 op.input_tensors.resize(3); 441 442 // Inputs are input, paddings and constant_values. 443 444 INFER_OK(op, "?;?;?", "?"); 445 446 // Check shape of paddings. 447 INFER_ERROR("Shape must be rank 2 but is rank 3", op, "?;[1,2,3];?"); 448 INFER_ERROR("Dimension must be 2 but is 4", op, "?;[1,4];?"); 449 450 // input.rank and paddings.dim(0) are equal. This is the number of dims in 451 // output. 452 INFER_ERROR("Shape must be rank 4 but is rank 3", op, "[1,2,3];[4,2];[]"); 453 INFER_OK(op, "[1,2,3];?;[]", "[?,?,?]"); 454 INFER_OK(op, "?;[3,2];[]", "[?,?,?]"); 455 456 // Make the paddings tensor known and verify padding values get added. 457 // E.g., if padding is ((1,10),(2,20),(3,30)) then values 11,22,23 are added 458 // to input dims to get output. 459 Tensor paddings_t(DT_INT64, TensorShape{3, 2}); 460 test::FillValues<int64>(&paddings_t, {1, 10, 2, 20, 3, 30}); 461 op.input_tensors[1] = &paddings_t; 462 INFER_OK(op, "[100,200,300];[3,2];[]", "[111,222,333]"); 463 INFER_OK(op, "[100,?,300];[3,2];[]", "[111,?,333]"); 464 INFER_OK(op, "?;[3,2];[]", "[?,?,?]"); 465 INFER_OK(op, "?;?;[]", "[?,?,?]"); 466 } 467 468 TEST(ArrayOpsTest, MirrorPadGrad_ShapeFn) { 469 ShapeInferenceTestOp op("MirrorPadGrad"); 470 op.input_tensors.resize(2); 471 472 // Inputs are input and paddings. 473 INFER_OK(op, "?;?", "?"); 474 475 // First padding dimension is unknown, so rank is unknown. 476 INFER_OK(op, "?;[?,4]", "?"); 477 478 // Input tensor rank doesn't match paddings dimension. 479 INFER_ERROR("must be rank 3 but is rank 2", op, "[?,?];[3,2]"); 480 481 // Paddings tensor is not a [rank x 2] matrix. 482 INFER_ERROR("Dimension 1 in both shapes must be equal, but are 3 and 2", op, 483 "[?,?,?];[3,3]"); 484 485 // Paddings tensor is unknown, but rank is known, so the output 486 // shape is a rank 3 unknown shape. 487 INFER_OK(op, "[?,?,?];[3,2]", "[?,?,?]"); 488 489 // Make the paddings tensor known and verify padding values get 490 // subtracted. E.g., if padding is ((1,10),(2,20),(3,30)) then 491 // values 11,22,23 are subtracted to input dims to get output. 492 Tensor paddings_t(DT_INT64, TensorShape{3, 2}); 493 test::FillValues<int64>(&paddings_t, {1, 10, 2, 20, 3, 30}); 494 op.input_tensors[1] = &paddings_t; 495 496 INFER_OK(op, "[111,222,333];[3,2]", "[100,200,300]"); 497 INFER_OK(op, "[111,?,333];[3,2]", "[100,?,300]"); 498 } 499 500 TEST(ArrayOpsTest, BroadcastArgs_ShapeFn) { 501 ShapeInferenceTestOp op("BroadcastArgs"); 502 INFER_OK(op, "?;?", "[?]"); 503 INFER_OK(op, "[123];[1]", "[123]"); 504 INFER_OK(op, "[1];[123]", "[123]"); 505 INFER_OK(op, "[123];[121]", "[123]"); 506 507 // Rank checks 508 INFER_ERROR("Shape must be rank 1 but is rank 0", op, "[];?"); 509 INFER_ERROR("Shape must be rank 1 but is rank 0", op, "?;[]"); 510 } 511 512 TEST(ArrayOpsTest, BroadcastTo_ShapeFn) { 513 ShapeInferenceTestOp op("BroadcastTo"); 514 op.input_tensors.resize(2); 515 516 INFER_OK(op, "?;[?]", "?"); 517 INFER_OK(op, "[];[1]", "[?]"); 518 INFER_OK(op, "[1];[1]", "[?]"); 519 INFER_OK(op, "[1];[2]", "[?,?]"); 520 INFER_OK(op, "[2,2];[3]", "[?,d0_0,d0_1]"); 521 522 // Rank checks 523 INFER_ERROR("Shape must be rank 1 but is rank 2", op, "?;[?,?]"); 524 INFER_ERROR("Shape must be rank 1 but is rank 0", op, "[2];[]"); 525 INFER_ERROR("Shape must be at most rank 1 but is rank 2", op, "[2,2];[1]"); 526 527 Tensor shape_t(DT_INT64, TensorShape{3}); 528 test::FillValues<int64>(&shape_t, {2, 10, 3}); 529 op.input_tensors[1] = &shape_t; 530 INFER_OK(op, "[1,?,1];[3]", "[2,10,3]"); 531 INFER_OK(op, "[1,1,1];[3]", "[2,10,3]"); 532 INFER_OK(op, "[10,1];[3]", "[2,d0_0,3]"); 533 INFER_ERROR("Dimensions must be equal, but are 3 and 2 for", op, 534 "[3,1,1];[3]"); 535 INFER_ERROR("Dimensions must be equal, but are 2 and 10 for", op, 536 "[2,2,1];[3]"); 537 } 538 539 TEST(ArrayOpsTest, BroadcastGradientArgs_ShapeFn) { 540 ShapeInferenceTestOp op("BroadcastGradientArgs"); 541 // Output is always two unknown vectors. 542 INFER_OK(op, "?;?", "[?];[?]"); 543 INFER_OK(op, "[123];[456]", "[?];[?]"); 544 545 // Rank checks 546 INFER_ERROR("Shape must be rank 1 but is rank 0", op, "[];?"); 547 INFER_ERROR("Shape must be rank 1 but is rank 0", op, "?;[]"); 548 } 549 550 TEST(ArrayOpsTest, ListDiff_ShapeFn) { 551 ShapeInferenceTestOp op("BroadcastGradientArgs"); 552 // Output is always two matching unknown vectors. 553 INFER_OK(op, "?;?", "[?];[?]"); 554 INFER_OK(op, "[123];[456]", "[?];[?]"); 555 556 // Rank checks 557 INFER_ERROR("Shape must be rank 1 but is rank 0", op, "[];?"); 558 INFER_ERROR("Shape must be rank 1 but is rank 0", op, "?;[]"); 559 } 560 561 TEST(ArrayOpsTest, MatrixSetDiag_ShapeFn) { 562 ShapeInferenceTestOp op("MatrixSetDiag"); 563 564 // Inputs are input and diagonal. 565 566 // Rank checks. 567 INFER_ERROR("Shape must be at least rank 2 but is rank 1", op, "[1];?"); 568 INFER_ERROR("Shape must be at least rank 1 but is rank 0", op, "?;[]"); 569 INFER_ERROR("Shape must be at least rank 1 but is rank 0", op, "[2,2];[]"); 570 INFER_ERROR("Shape must be rank 1 but is rank 2", op, "[2,2];[2,2]"); 571 572 // diagonal[-1] must match smallest matrix dimension. 573 INFER_ERROR("Dimensions must be equal, but are 2 and 3", op, "[2,3];[3]"); 574 575 // Output matches input. 576 INFER_OK(op, "?;?", "in0"); 577 INFER_OK(op, "[1,2,2];[1,2]", "in0"); 578 INFER_OK(op, "[1,2,3];?", "in0"); 579 INFER_OK(op, "[1,3,2];?", "in0"); 580 INFER_OK(op, "[1,?,2];[?,?]", "in0"); 581 INFER_OK(op, "[1,?,?];[?,2]", "in0"); 582 583 // Infer batch shape from diag when input is not fully specified. 584 INFER_OK(op, "?;[1,2]", "[d1_0,?,?]"); 585 INFER_OK(op, "[?,?,3];[1,2]", "[d1_0,d0_1,d0_2]"); 586 INFER_OK(op, "[?,3,?];[1,2]", "[d1_0,d0_1,d0_2]"); 587 INFER_OK(op, "[?,3,2];[1,2]", "[d1_0,d0_1,d0_2]"); 588 } 589 590 TEST(ArrayOpsTest, ExpandDims_ShapeFn) { 591 ShapeInferenceTestOp op("ExpandDims"); 592 op.input_tensors.resize(2); 593 594 // With unknown dim tensor value, output is unknown. 595 INFER_OK(op, "?;?", "?"); 596 Tensor dim_t; 597 op.input_tensors[1] = &dim_t; 598 599 // Expand at front of tensor. 600 for (int32 idx : {0, -4}) { 601 dim_t = test::AsScalar<int32>(idx); 602 INFER_OK(op, "?;?", "?"); 603 INFER_OK(op, "[5,?,7];?", "[1,d0_0,d0_1,d0_2]"); 604 } 605 606 // Expand at middle of tensor. 607 for (int32 idx : {1, -3}) { 608 dim_t = test::AsScalar<int32>(idx); 609 INFER_OK(op, "?;?", "?"); 610 INFER_OK(op, "[5,?,7];?", "[d0_0,1,d0_1,d0_2]"); 611 612 // Repeat with int64. 613 dim_t = test::AsScalar<int64>(idx); 614 INFER_OK(op, "?;?", "?"); 615 INFER_OK(op, "[5,?,7];?", "[d0_0,1,d0_1,d0_2]"); 616 } 617 for (int32 idx : {2, -2}) { 618 dim_t = test::AsScalar<int32>(idx); 619 INFER_OK(op, "?;?", "?"); 620 INFER_OK(op, "[5,?,7];?", "[d0_0,d0_1,1,d0_2]"); 621 622 // Repeat with int64. 623 dim_t = test::AsScalar<int64>(idx); 624 INFER_OK(op, "?;?", "?"); 625 INFER_OK(op, "[5,?,7];?", "[d0_0,d0_1,1,d0_2]"); 626 } 627 628 for (int32 idx : {3, -1}) { 629 // Expand at the end. 630 dim_t = test::AsScalar<int32>(idx); 631 INFER_OK(op, "?;?", "?"); 632 INFER_OK(op, "[5,?,7];?", "[d0_0,d0_1,d0_2,1]"); 633 634 // Repeat with int64. 635 dim_t = test::AsScalar<int64>(idx); 636 INFER_OK(op, "?;?", "?"); 637 INFER_OK(op, "[5,?,7];?", "[d0_0,d0_1,d0_2,1]"); 638 } 639 for (int32 idx : {4, -5}) { 640 // Invalid idx. 641 dim_t = test::AsScalar<int32>(idx); 642 INFER_ERROR("not in the interval [-4, 3]", op, "[5,?,7];?"); 643 dim_t = test::AsScalar<int64>(idx); 644 INFER_ERROR("not in the interval [-4, 3]", op, "[5,?,7];?"); 645 } 646 647 // Expand using an input vector tensor. 648 std::vector<int32> dims; 649 dims.push_back(0); 650 dim_t = test::AsTensor<int32>(dims); 651 INFER_OK(op, "?;?", "?"); 652 INFER_OK(op, "[5,?,7];?", "[1,d0_0,d0_1,d0_2]"); 653 654 // Expand using too many input elements. 655 dims.push_back(1); 656 dim_t = test::AsTensor<int32>(dims); 657 INFER_ERROR("'dim' input must be a tensor with a single", op, "?;?"); 658 INFER_ERROR("'dim' input must be a tensor with a single", op, "[5,6,7];?"); 659 660 // Examples from ExpandDims doc. 661 dim_t = test::AsScalar<int32>(0); 662 INFER_OK(op, "[2];[]", "[1,d0_0]"); 663 dim_t = test::AsScalar<int32>(1); 664 INFER_OK(op, "[2];[]", "[d0_0,1]"); 665 dim_t = test::AsScalar<int32>(-1); 666 INFER_OK(op, "[2];[]", "[d0_0,1]"); 667 } 668 669 TEST(ArrayOpsTest, ImmutableConst_ShapeFn) { 670 ShapeInferenceTestOp op("ImmutableConst"); 671 672 TF_ASSERT_OK(NodeDefBuilder("test", "ImmutableConst") 673 .Attr("dtype", DT_FLOAT) 674 .Attr("shape", TensorShape({1, 2, 3})) 675 .Attr("memory_region_name", "test_region") 676 .Finalize(&op.node_def)); 677 INFER_OK(op, "", "[1,2,3]"); 678 679 TF_ASSERT_OK(NodeDefBuilder("test", "ImmutableConst") 680 .Attr("dtype", DT_FLOAT) 681 .Attr("shape", TensorShape({})) 682 .Attr("memory_region_name", "test_region") 683 .Finalize(&op.node_def)); 684 INFER_OK(op, "", "[]"); 685 686 TF_ASSERT_OK(NodeDefBuilder("test", "ImmutableConst") 687 .Attr("dtype", DT_FLOAT) 688 .Attr("shape", "invalid") 689 .Attr("memory_region_name", "test_region") 690 .Finalize(&op.node_def)); 691 INFER_ERROR("AttrValue had value with type 'string' when 'shape' expected", 692 op, ""); 693 } 694 695 TEST(ArrayOpsTest, Concat_ShapeFn) { 696 ShapeInferenceTestOp op("Concat"); 697 auto set_n = [&op](int n) { 698 std::vector<NodeDefBuilder::NodeOut> src_list; 699 src_list.reserve(n); 700 for (int i = 0; i < n; ++i) src_list.emplace_back("a", 0, DT_FLOAT); 701 TF_ASSERT_OK(NodeDefBuilder("test", "Concat") 702 .Input({"concat_dim", 0, DT_INT32}) 703 .Input(src_list) 704 .Attr("n", n) 705 .Finalize(&op.node_def)); 706 }; 707 708 // Confirm dimension[0] of the input (the concat_dim) is a scalar. 709 set_n(2); 710 INFER_ERROR("Shape must be rank 0 but is rank 1", op, "[1];?;?"); 711 712 // Test with the input concat_dim tensor not known. This takes the known rank 713 // of the inputs and makes a tensor of that many unknown dims. 714 set_n(7); 715 INFER_OK(op, "?;?;?;?;[1,2,3];?;[3,2,1];?", "[?,?,?]"); 716 set_n(4); 717 INFER_OK(op, "?;?;?;[1,2,3,4];[4,3,2,1]", "[?,?,?,?]"); 718 INFER_OK(op, "?;?;?;?;?", "?"); // output rank unknown 719 INFER_ERROR("Can't concatenate scalars (use tf.stack instead)", op, 720 "?;?;?;[];[]"); 721 INFER_ERROR("Shape must be rank 2 but is rank 3", op, "?;?;?;[1,2];[1,2,3]"); 722 723 // Test when the concat_dim tensor is known. The concatenated dimension is 724 // summed across all input tensors, and other dimensions are merged. 725 Tensor concat_dim_t; 726 op.input_tensors.push_back(&concat_dim_t); 727 set_n(2); 728 729 // Sum dim 0, merge the other two dims. 730 for (int concat_dim : {0, -3}) { 731 concat_dim_t = test::AsScalar(concat_dim); 732 INFER_OK(op, "[];[100,2,?];[10,?,3]", "[110,d1_1,d2_2]"); 733 INFER_ERROR("Dimension 1 in both shapes must be equal, but are 5 and 3", op, 734 "[];[100,2,5];[10,?,3]"); 735 // concat_dim can't be summed, as one value is unknown. 736 INFER_OK(op, "[];[100,2,?];[?,?,3]", "[?,d1_1,d2_2]"); 737 INFER_OK(op, "[];[?,2,?];[10,?,3]", "[?,d1_1,d2_2]"); 738 } 739 740 // Test with a higher concat_dim. 741 for (bool use_negative : {false, true}) { 742 concat_dim_t = test::AsScalar(use_negative ? -2 : 1); 743 INFER_OK(op, "[];[1,100,?];[?,10,3]", "[d1_0,110,d2_2]"); 744 concat_dim_t = test::AsScalar(use_negative ? -1 : 1); 745 INFER_OK(op, "[];[1,100];[?,10]", "[d1_0,110]"); 746 INFER_OK(op, "[];[?,100];[1,10]", "[d2_0,110]"); 747 748 // concat_dim is out of bounds. 749 concat_dim_t = test::AsScalar(use_negative ? -2 : 1); 750 INFER_ERROR("Shape must be at least rank 2 but is rank 1", op, 751 "[];[100];[10,?]"); 752 INFER_ERROR("Shape must be at least rank 2 but is rank 1", op, 753 "[];[100,5];[10]"); 754 } 755 756 // concat_dim is too low. 757 concat_dim_t = test::AsScalar(-2); 758 INFER_ERROR("Shape must be at least rank 2 but is rank 1", op, 759 "[];[100];[10,?]"); 760 INFER_ERROR("Shape must be at least rank 2 but is rank 1", op, 761 "[];[100,5];[10]"); 762 763 // Repeat successful case with several unknown inputs. 764 set_n(5); 765 concat_dim_t = test::AsScalar(1); 766 INFER_OK(op, "[];?;[1,100,?];[?,?,?];[?,10,3];?", "[d2_0,?,d4_2]"); 767 } 768 769 TEST(ArrayOpsTest, ConcatV2_ShapeFn) { 770 ShapeInferenceTestOp op("ConcatV2"); 771 auto set_n = [&op](int n) { 772 std::vector<NodeDefBuilder::NodeOut> src_list; 773 src_list.reserve(n); 774 for (int i = 0; i < n; ++i) src_list.emplace_back("a", 0, DT_FLOAT); 775 TF_ASSERT_OK(NodeDefBuilder("test", "ConcatV2") 776 .Input(src_list) 777 .Input({"axis", 0, DT_INT32}) 778 .Attr("n", n) 779 .Finalize(&op.node_def)); 780 }; 781 782 // Confirm dimension[0] of the input (the concat_dim) is a scalar. 783 set_n(2); 784 INFER_ERROR("Shape must be rank 0 but is rank 1", op, "?;?;[1]"); 785 786 // Test with the input concat_dim tensor not known. This takes the known rank 787 // of the inputs and makes a tensor of that many unknown dims. 788 set_n(7); 789 INFER_OK(op, "?;?;?;?;[1,2,3];?;[3,2,1];?", "[?,?,?]"); 790 set_n(4); 791 INFER_OK(op, "?;?;[1,2,3,4];[4,3,2,1];?", "[?,?,?,?]"); 792 INFER_OK(op, "?;?;?;?;?", "?"); // output rank unknown 793 INFER_ERROR("Can't concatenate scalars (use tf.stack instead)", op, 794 "?;?;[];[];?"); 795 INFER_ERROR("Shape must be rank 2 but is rank 3", op, "?;?;[1,2];[1,2,3];?"); 796 797 // Test when the concat_dim tensor is known. The concatenated dimension is 798 // summed across all input tensors, and other dimensions are merged. 799 Tensor concat_dim_t; 800 op.input_tensors.resize(3); 801 op.input_tensors[2] = &concat_dim_t; 802 803 set_n(2); 804 805 // Invalid concat dim value. 806 // concat_dim_t = test::AsScalar(-1); 807 // INFER_ERROR("Expected concat_dim >= 0, but got -1", op, "?;?;?"); 808 809 // Sum dim 0, merge the other two dims. 810 concat_dim_t = test::AsScalar(0); 811 INFER_OK(op, "[100,2,?];[10,?,3];[]", "[110,d0_1,d1_2]"); 812 INFER_ERROR("Dimension 1 in both shapes must be equal, but are 5 and 3", op, 813 "[100,2,5];[10,?,3];[]"); 814 // concat_dim can't be summed, as one value is unknown. 815 INFER_OK(op, "[100,2,?];[?,?,3];[]", "[?,d0_1,d1_2]"); 816 INFER_OK(op, "[?,2,?];[10,?,3];[]", "[?,d0_1,d1_2]"); 817 818 // Test with a higher concat_dim. 819 concat_dim_t = test::AsScalar(1); 820 INFER_OK(op, "[1,100,?];[?,10,3];[]", "[d0_0,110,d1_2]"); 821 INFER_OK(op, "[1,100];[?,10];[]", "[d0_0,110]"); 822 INFER_OK(op, "[?,100];[1,10];[]", "[d1_0,110]"); 823 // concat_dim is too high. 824 INFER_ERROR("Shape must be at least rank 2 but is rank 1", op, 825 "[100];[10,?];[]"); 826 INFER_ERROR("Shape must be at least rank 2 but is rank 1", op, 827 "[100,5];[10];[]"); 828 // concat_dim is too low. 829 concat_dim_t = test::AsScalar(-2); 830 INFER_ERROR("Shape must be at least rank 2 but is rank 1", op, 831 "[100];[10,?];[]"); 832 INFER_ERROR("Shape must be at least rank 2 but is rank 1", op, 833 "[100,5];[10];[]"); 834 835 // Repeat successful case with several unknown inputs. 836 op.input_tensors.resize(6); 837 op.input_tensors[3] = nullptr; 838 op.input_tensors[5] = &concat_dim_t; 839 concat_dim_t = test::AsScalar(1); 840 841 set_n(5); 842 INFER_OK(op, "?;[1,100,?];[?,?,?];[?,10,3];?;[]", "[d1_0,?,d3_2]"); 843 } 844 845 TEST(ArrayOpsTest, ConcatOffset_ShapeFn) { 846 ShapeInferenceTestOp op("ConcatOffset"); 847 848 const int n = 4; 849 std::vector<NodeDefBuilder::NodeOut> src_list; 850 src_list.reserve(n); 851 for (int i = 0; i < n; ++i) src_list.emplace_back("a", 0, DT_INT32); 852 TF_ASSERT_OK(NodeDefBuilder("test", "ConcatOffset") 853 .Input({"concat_dim", 0, DT_INT32}) 854 .Input(src_list) 855 .Attr("n", n) 856 .Finalize(&op.node_def)); 857 INFER_OK(op, "?;?;?;?;?", "in1;in2;in3;in4"); 858 } 859 860 TEST(ArrayOpsTest, Reshape_ShapeFn) { 861 ShapeInferenceTestOp op("Reshape"); 862 op.input_tensors.resize(2); 863 864 // No valid shape provided. 865 INFER_OK(op, "?;?", "?"); 866 INFER_OK(op, "[?];?", "?"); 867 INFER_OK(op, "[?];[?]", "?"); 868 INFER_OK(op, "[4];[?]", "?"); 869 870 // All dimensions provided. 871 Tensor new_shape = test::AsTensor<int32>({1, 2, 3}); 872 op.input_tensors[1] = &new_shape; 873 INFER_OK(op, "[?];[3]", "[1,2,3]"); 874 INFER_OK(op, "[6];[3]", "[1,2,3]"); 875 // The number of elements should match for the reshape to succeed. 876 INFER_ERROR( 877 "Cannot reshape a tensor with 12 elements to shape [1,2,3] (6 elements)", 878 op, "[3,4];[3]"); 879 880 // Unknown dimensions. 881 // Flatten: 882 new_shape = test::AsTensor<int32>({-1}); 883 INFER_OK(op, "[?];[1]", "[d0_0]"); 884 INFER_OK(op, "[2,2];[1]", "[4]"); 885 // The first dimension is inferred: 886 new_shape = test::AsTensor<int32>({2, -1}); 887 INFER_OK(op, "[3,4];[2]", "[2,6]"); 888 // The total number of elements must be evenly divisible by the known 889 // dimensions. 890 INFER_ERROR("Dimension size must be evenly divisible by 2 but is 7", op, 891 "[7];[2]"); 892 // Multiple missing dimensions cannot be inferred. 893 new_shape = test::AsTensor<int32>({-1, -1, 2}); 894 INFER_OK(op, "[8];[3]", "[?,?,2]"); 895 896 // Symbolic shape propagation 897 new_shape = test::AsTensor<int32>({-1, 2, 3}); 898 INFER_OK(op, "[?,2,3];[3]", "[d0_0,2,3]"); 899 900 // Reshaping to a scalar. 901 new_shape = test::AsTensor<int32>({}); 902 INFER_OK(op, "[1];[0]", "[]"); 903 INFER_ERROR( 904 "Cannot reshape a tensor with 2 elements to shape [] (1 elements)", op, 905 "[1,2];[0]"); 906 907 // Reshaping a tensor with no elements. 908 new_shape = test::AsTensor<int32>({-1}); 909 INFER_OK(op, "[0];[1]", "[0]"); 910 new_shape = test::AsTensor<int32>({-1, 6}); 911 INFER_OK(op, "[0,2];[1]", "[0,6]"); 912 new_shape = test::AsTensor<int32>({0, -1}); 913 INFER_OK(op, "[0,2];[1]", "[0,?]"); 914 } 915 916 TEST(ArrayOpsTest, QuantizedReshape_ShapeFn) { 917 ShapeInferenceTestOp op("QuantizedReshape"); 918 op.input_tensors.resize(2); 919 920 // First test a subset of the Reshape_ShapeFn tests. Not all are tested, as 921 // QuantizedReshape uses the same code for the reshape part of the operation. 922 INFER_OK(op, "?;?;?;?", "?;[];[]"); 923 INFER_OK(op, "[?];?;?;?", "?;[];[]"); 924 INFER_OK(op, "[?];[?];?;?", "?;[];[]"); 925 INFER_OK(op, "[4];[?];?;?", "?;[];[]"); 926 Tensor new_shape = test::AsTensor<int32>({1, 2, 3}); 927 op.input_tensors[1] = &new_shape; 928 INFER_OK(op, "[?];[3];?;?", "[1,2,3];[];[]"); 929 INFER_OK(op, "[6];[3];?;?", "[1,2,3];[];[]"); 930 INFER_ERROR( 931 "Cannot reshape a tensor with 12 elements to shape [1,2,3] (6 elements)", 932 op, "[3,4];[3];?;?"); 933 934 // Test the scalar rank checks on input_min and input_max. 935 INFER_ERROR("must be rank 0", op, "?;?;[1];?"); 936 INFER_ERROR("must be rank 0", op, "?;?;?;[1]"); 937 } 938 939 TEST(ArrayOpsTest, Placeholder_ShapeFn) { 940 { 941 // 2D shape 942 ShapeInferenceTestOp op("Placeholder"); 943 TensorShape shape({1, 2}); 944 TF_ASSERT_OK(NodeDefBuilder("test", "Placeholder") 945 .Attr("shape", shape) 946 .Attr("dtype", DT_FLOAT) 947 .Finalize(&op.node_def)); 948 INFER_OK(op, "", "[1,2]"); 949 } 950 951 { 952 // Scalar shapes are supported 953 ShapeInferenceTestOp op("Placeholder"); 954 TensorShape shape({}); 955 TF_ASSERT_OK(NodeDefBuilder("test", "Placeholder") 956 .Attr("shape", shape) 957 .Attr("dtype", DT_FLOAT) 958 .Finalize(&op.node_def)); 959 INFER_OK(op, "", "[]"); 960 } 961 962 { 963 // Partial shape 964 ShapeInferenceTestOp op("Placeholder"); 965 const int64 dims[2] = {1, -1}; 966 PartialTensorShape shape; 967 TF_ASSERT_OK(PartialTensorShape::MakePartialShape(dims, 2, &shape)); 968 TF_ASSERT_OK(NodeDefBuilder("test", "Placeholder") 969 .Attr("shape", shape) 970 .Attr("dtype", DT_FLOAT) 971 .Finalize(&op.node_def)); 972 INFER_OK(op, "", "[1,?]"); 973 } 974 975 { 976 // Unknown shape 977 ShapeInferenceTestOp op("Placeholder"); 978 PartialTensorShape shape; 979 TF_ASSERT_OK(NodeDefBuilder("test", "Placeholder") 980 .Attr("shape", shape) 981 .Attr("dtype", DT_FLOAT) 982 .Finalize(&op.node_def)); 983 INFER_OK(op, "", "?"); 984 } 985 } 986 987 TEST(ArrayOpsTest, Transpose_ShapeFn) { 988 ShapeInferenceTestOp op("Transpose"); 989 op.input_tensors.resize(2); 990 991 // Missing shape information. 992 INFER_OK(op, "?;?", "?"); 993 INFER_OK(op, "?;[?]", "?"); 994 INFER_OK(op, "?;[2]", "[?,?]"); 995 INFER_OK(op, "[?];?", "[?]"); 996 INFER_OK(op, "[?,?];[2]", "[?,?]"); 997 INFER_ERROR("Dimension must be 3 but is 2", op, "[1,2,3];[2]"); 998 Tensor perm = test::AsTensor<int32>({0}); 999 op.input_tensors[1] = &perm; 1000 INFER_OK(op, "[?];[?]", "[d0_0]"); 1001 perm = test::AsTensor<int32>({1, 0}); 1002 INFER_OK(op, "?;[2]", "[?,?]"); 1003 INFER_OK(op, "[?,?];[2]", "[d0_1,d0_0]"); 1004 INFER_OK(op, "[1,?];[2]", "[d0_1,d0_0]"); 1005 INFER_OK(op, "?;[0]", "in0"); 1006 1007 // Invalid arguments. 1008 perm = test::AsTensor<int32>({1, 2}); 1009 INFER_ERROR("perm dim 2 is out of range of input rank 2", op, "[1,2];[2]"); 1010 perm = test::AsTensor<int32>({0}); 1011 INFER_ERROR("Dimension must be 2 but is 1", op, "[1,2];[1]"); 1012 1013 // Larger valid cases. 1014 perm = test::AsTensor<int32>({1, 0, 3, 4, 2}); 1015 INFER_OK(op, "[0,1,2,3,4];[5]", "[d0_1,d0_0,d0_3,d0_4,d0_2]"); 1016 INFER_OK(op, "[0,?,2,3,4];[5]", "[d0_1,d0_0,d0_3,d0_4,d0_2]"); 1017 } 1018 1019 TEST(ArrayOpsTest, Bitcast_ShapeFn) { 1020 ShapeInferenceTestOp op("Bitcast"); 1021 auto rebuild_node_def = [&op](DataType input_type, DataType output_type) { 1022 TF_ASSERT_OK(NodeDefBuilder("test", "Bitcast") 1023 .Input("input", 0, input_type) 1024 .Attr("type", output_type) 1025 .Finalize(&op.node_def)); 1026 }; 1027 1028 rebuild_node_def(DT_FLOAT, DT_INT32); 1029 // No valid shape provided, so output is unknown. 1030 INFER_OK(op, "?", "?"); 1031 1032 // Bitcasting from two equal sizes propagates shape. 1033 INFER_OK(op, "[1,2]", "in0"); 1034 1035 // Bitcasting from smaller to larger reduces the size of the last dimension. 1036 rebuild_node_def(DT_INT32, DT_INT64); 1037 INFER_OK(op, "[1,2]", "[d0_0]"); // last dimension matches divisor. 1038 // TODO(vrv): Seems like a bug, or at least, too lenient. 1039 INFER_OK(op, "[1,?]", "[d0_0]"); 1040 // 4 is divisible by 2, but the shape function signature requires 1041 // that the last dimension matches the last value exactly. 1042 INFER_ERROR("does not match", op, "[1,4]"); 1043 INFER_ERROR("does not match", op, "[1,3]"); 1044 1045 // Bitcasting from a larger type to a smaller type extends the dimension 1046 rebuild_node_def(DT_INT64, DT_INT32); 1047 INFER_OK(op, "[4,5]", "[d0_0,d0_1,2]"); 1048 rebuild_node_def(DT_COMPLEX128, DT_INT32); 1049 INFER_OK(op, "[4,5]", "[d0_0,d0_1,4]"); 1050 rebuild_node_def(DT_COMPLEX128, DT_HALF); 1051 INFER_OK(op, "[4,5]", "[d0_0,d0_1,8]"); 1052 rebuild_node_def(DT_COMPLEX128, DT_INT8); 1053 INFER_OK(op, "[4,5]", "[d0_0,d0_1,16]"); 1054 1055 // Bitcasting from a POD or quantized datatype is not allowed. 1056 rebuild_node_def(DT_STRING, DT_INT32); 1057 INFER_ERROR("one of the type sizes is zero", op, "[1,2,3]"); 1058 rebuild_node_def(DT_INT32, DT_STRING); 1059 INFER_ERROR("one of the type sizes is zero", op, "[1,2,3]"); 1060 } 1061 1062 TEST(ArrayOpsTest, Squeeze_ShapeFn) { 1063 ShapeInferenceTestOp op("Squeeze"); 1064 1065 auto rebuild_node_def = [&op](const std::vector<int32>& squeeze_dims) { 1066 TF_ASSERT_OK(NodeDefBuilder("test", "Squeeze") 1067 .Input("input", 0, DT_FLOAT) 1068 .Attr("squeeze_dims", squeeze_dims) 1069 .Finalize(&op.node_def)); 1070 }; 1071 1072 // Default squeeze_dims = [] 1073 rebuild_node_def({}); 1074 1075 // No valid shape provided, so output is unknown. 1076 INFER_OK(op, "?", "?"); 1077 1078 INFER_OK(op, "[1,4,1,5,1]", "[d0_1,d0_3]"); 1079 1080 // Squeezing all dimensions, but see some unknown values. 1081 INFER_OK(op, "[1,?,1,?,1]", "?"); 1082 1083 // Test simple squeeze of an explicit dimension 1084 rebuild_node_def({1}); 1085 INFER_OK(op, "[4,1,5]", "[d0_0,d0_2]"); 1086 // Squeezing unknown dim explicitly, assumes it's 1 at runtime. 1087 INFER_OK(op, "[4,?,5]", "[d0_0,d0_2]"); 1088 1089 // Attempt to squeeze non-one dimension 1090 INFER_ERROR("Can not squeeze dim[1]", op, "[4,6,5]"); 1091 1092 // Squeeze multiple dimensions 1093 rebuild_node_def({1, 2}); 1094 INFER_OK(op, "[4,1,1,5]", "[d0_0,d0_3]"); 1095 rebuild_node_def({1, -2}); 1096 INFER_OK(op, "[4,1,1,5]", "[d0_0,d0_3]"); 1097 1098 // Negative squeeze dim 1099 rebuild_node_def({-2}); 1100 INFER_OK(op, "[4,1,5]", "[d0_0,d0_2]"); 1101 1102 // Test validation of squeeze dimensions 1103 rebuild_node_def({-4}); 1104 INFER_ERROR("not in [-3,3)", op, "[1,2,3]"); 1105 rebuild_node_def({3}); 1106 INFER_ERROR("not in [-3,3)", op, "[1,2,3]"); 1107 } 1108 1109 TEST(ArrayOpsTest, ReverseSequence_ShapeFn) { 1110 ShapeInferenceTestOp op("ReverseSequence"); 1111 auto rebuild_node_def = [&op](const int32 seq_dim, const int32 batch_dim) { 1112 TF_ASSERT_OK(NodeDefBuilder("test", "ReverseSequence") 1113 .Input("input", 0, DT_FLOAT) 1114 .Input("seq_lengths", 1, DT_INT64) 1115 .Attr("seq_dim", seq_dim) 1116 .Attr("batch_dim", batch_dim) 1117 .Finalize(&op.node_def)); 1118 }; 1119 1120 rebuild_node_def(1, 2); 1121 // No valid shape provided, so output is unknown. 1122 INFER_OK(op, "?;[10]", "?"); 1123 1124 // Bad rank for seq_lengths 1125 INFER_ERROR("Shape must be rank 1 but is rank 2", op, "?;[10,10]"); 1126 1127 // Validate seq_dim and batch_dim 1128 rebuild_node_def(1, 4); 1129 INFER_ERROR("batch_dim must be < input rank", op, "[1,2,3];[3]"); 1130 rebuild_node_def(4, 1); 1131 INFER_ERROR("seq_dim must be < input rank", op, "[1,2,3];[3]"); 1132 1133 rebuild_node_def(1, 2); 1134 INFER_OK(op, "[1,2,3];[3]", "[d0_0,d0_1,d0_2]"); 1135 // Resolves uncertainty on batch dimension by merging. 1136 INFER_OK(op, "[1,2,?];[3]", "[d0_0,d0_1,d1_0]"); 1137 INFER_OK(op, "[1,2,3];[?]", "[d0_0,d0_1,d0_2]"); 1138 } 1139 1140 TEST(ArrayOpsTest, Split_ShapeFn) { 1141 ShapeInferenceTestOp op("Split"); 1142 op.input_tensors.resize(2); 1143 1144 // No value for split_dim and no input. 1145 TF_ASSERT_OK(NodeDefBuilder("test", "Split") 1146 .Input("split_dim", 0, DT_INT32) 1147 .Input("value", 1, DT_FLOAT) 1148 .Attr("num_split", 2) 1149 .Finalize(&op.node_def)); 1150 INFER_OK(op, "?;?", "?;?"); 1151 // If the rank is known, we know the rank of each output. 1152 INFER_OK(op, "?;[?,?]", "[?,?];[?,?]"); 1153 1154 // split_dim is unknown but other inputs are known. 1155 INFER_OK(op, "?;[1,4]", "[?,?];[?,?]"); 1156 1157 // split_dim is known. 1158 Tensor split_dim = test::AsTensor<int32>({1, 2}); 1159 op.input_tensors[0] = &split_dim; 1160 INFER_ERROR("Input must be scalar but has rank 1", op, "[?];[?,?]"); 1161 split_dim = test::AsScalar<int32>(1); 1162 INFER_OK(op, "?;?", "?;?"); 1163 INFER_OK(op, "?;[?,?]", "[d1_0,?];[d1_0,?]"); 1164 INFER_OK(op, "?;[1,4]", "[d1_0,2];[d1_0,2]"); 1165 INFER_OK(op, "?;[1,?]", "[d1_0,?];[d1_0,?]"); 1166 INFER_ERROR("Dimension size must be evenly divisible by 2 but is 5", op, 1167 "?;[1,5]"); 1168 1169 // split_dim too large. 1170 split_dim = test::AsScalar<int32>(3); 1171 INFER_ERROR( 1172 "Dimension size, given by scalar input 3 must be in range [-3, 3)", op, 1173 "?;[1,4,8]"); 1174 1175 // Negative split_dim. 1176 split_dim = test::AsScalar<int32>(-1); 1177 INFER_OK(op, "?;?", "?;?"); 1178 INFER_OK(op, "?;[?,?]", "[d1_0,?];[d1_0,?]"); 1179 INFER_OK(op, "?;[1,?]", "[d1_0,?];[d1_0,?]"); 1180 INFER_OK(op, "?;[1,4]", "[d1_0,2];[d1_0,2]"); 1181 INFER_OK(op, "?;[1,4,8]", "[d1_0,d1_1,4];[d1_0,d1_1,4]"); 1182 split_dim = test::AsScalar<int32>(-2); 1183 INFER_OK(op, "?;[1,4,8]", "[d1_0,2,d1_2];[d1_0,2,d1_2]"); 1184 split_dim = test::AsScalar<int32>(-4); 1185 INFER_ERROR( 1186 "Dimension size, given by scalar input -4 must be in range [-3, 3)", op, 1187 "?;[1,4,8]"); 1188 } 1189 1190 TEST(ArrayOpsTest, Tile_ShapeFn) { 1191 ShapeInferenceTestOp op("Tile"); 1192 op.input_tensors.resize(2); 1193 1194 // No value for split_dim and no input. 1195 TF_ASSERT_OK(NodeDefBuilder("test", "Tile") 1196 .Input("input", 0, DT_FLOAT) 1197 .Input("multiples", 1, DT_INT32) 1198 .Finalize(&op.node_def)); 1199 1200 // If both are unknown, output is unknown. 1201 INFER_OK(op, "?;?", "?"); 1202 1203 // If multiples rank is unknown but input is, output rank is known. 1204 INFER_OK(op, "[2,3,1,4];?", "[?,?,?,?]"); 1205 1206 // Bad rank for 'multiples' 1207 INFER_ERROR("Shape must be rank 1 but is rank 2", op, "[2,3,1,4];[4,1]"); 1208 1209 // No multiples tensor available, but output rank is known from multiples. 1210 INFER_OK(op, "?;[4]", "[?,?,?,?]"); 1211 1212 // Test a tile of a 4D input. 1213 Tensor multiples = test::AsTensor<int32>({2, 3, 4, 5}); 1214 op.input_tensors[1] = &multiples; 1215 INFER_OK(op, "[2,3,1,4];[4]", "[4,9,4,20]"); 1216 // Test 64-bit tensor type 1217 multiples = test::AsTensor<int64>({2, 3, 4, 5}); 1218 INFER_OK(op, "[2,3,1,4];[4]", "[4,9,4,20]"); 1219 } 1220 1221 TEST(ArrayOpsTest, EditDistance_ShapeFn) { 1222 ShapeInferenceTestOp op("EditDistance"); 1223 op.input_tensors.resize(6); 1224 1225 // If the shape tensors are not available, the output shape is unknown. 1226 INFER_OK(op, "[?,?];[?];[4];[?,?];[?];[4]", "?"); 1227 1228 Tensor hypothesis_shape = test::AsTensor<int64>({2, 30, 4, 50}); 1229 op.input_tensors[2] = &hypothesis_shape; 1230 Tensor truth_shape = test::AsTensor<int64>({20, 3, 40, 5}); 1231 op.input_tensors[5] = &truth_shape; 1232 INFER_OK(op, "[?,?];[?];[4];[?,?];[?];[4]", "[20,30,40]"); 1233 1234 // Shape elements don't match 1235 hypothesis_shape = test::AsTensor<int64>({2}); 1236 op.input_tensors[2] = &hypothesis_shape; 1237 INFER_ERROR("Num elements of hypothesis_shape does not match truth_shape", op, 1238 "[?,?];[?];[1];[?,?];[?];[4]"); 1239 } 1240 1241 TEST(ArrayOpsTest, OneHot_ShapeFn) { 1242 ShapeInferenceTestOp op("OneHot"); 1243 op.input_tensors.resize(4); 1244 auto set_axis = [&op](int axis) { 1245 TF_ASSERT_OK(NodeDefBuilder("test", "OneHot") 1246 .Input("indices", 0, DT_FLOAT) 1247 .Input("depth", 1, DT_INT32) 1248 .Input("on_value", 2, DT_FLOAT) 1249 .Input("off_value", 3, DT_FLOAT) 1250 .Attr("axis", axis) 1251 .Finalize(&op.node_def)); 1252 }; 1253 1254 // Invalid axis value. 1255 set_axis(-2); 1256 INFER_ERROR("axis must be >= -1", op, "?;?;?;?"); 1257 set_axis(1); 1258 1259 // If indices shape is unknown, we return an unknown shape. 1260 INFER_OK(op, "?;[];?;?", "?"); 1261 1262 // Depth must be scalar. 1263 Tensor depth = test::AsTensor<int32>({1, 2}); 1264 op.input_tensors[1] = &depth; 1265 INFER_ERROR("Input must be scalar but has rank 1", op, "?;[2];?;?"); 1266 1267 // Full information is available. 1268 depth = test::AsScalar<int32>(2); 1269 INFER_OK(op, "[1,3,4];[];?;?", "[d0_0,2,d0_1,d0_2]"); 1270 set_axis(-1); 1271 INFER_OK(op, "[1,3,4];[];?;?", "[d0_0,d0_1,d0_2,2]"); 1272 } 1273 1274 TEST(ArrayOpsTest, ExtractImagePatchesShapeTest) { 1275 ShapeInferenceTestOp op("ExtractImagePatches"); 1276 auto set_op = [&op](const std::vector<int32>& ksizes, 1277 const std::vector<int32>& strides, 1278 const std::vector<int32>& rates, const string& padding) { 1279 TF_ASSERT_OK(NodeDefBuilder("test", "ExtractImagePatches") 1280 .Input("input", 0, DT_FLOAT) 1281 .Attr("ksizes", ksizes) 1282 .Attr("strides", strides) 1283 .Attr("rates", rates) 1284 .Attr("padding", padding) 1285 .Finalize(&op.node_def)); 1286 }; 1287 1288 // Just tests that the ksize calculation with rates works. Most of 1289 // the other code is boilerplate that is tested by a variety of 1290 // other ops. 1291 // 1292 // ksizes is 2x2. rate rows and cols is 2, so ksize_rows and 1293 // cols are changed to be 2 + (2 - 1) = 3. 7x7 input with 3x3 1294 // filter and 1x1 stride gives a 5x5 output. 1295 set_op({1, 2, 2, 1}, {1, 1, 1, 1}, {1, 2, 2, 1}, "VALID"); 1296 INFER_OK(op, "[1,7,7,2]", "[d0_0,5,5,8]"); 1297 // With ksizes as 1x1, the output depth is now exactly the last value of the 1298 // input and output spatial is reduced as well. 1299 set_op({1, 1, 1, 1}, {1, 1, 1, 1}, {1, 2, 2, 1}, "VALID"); 1300 INFER_OK(op, "[1,7,7,2]", "[d0_0,7,7,d0_3]"); 1301 1302 // Bad ksize rank 1303 set_op({1, 2, 2, 1, 1}, {1, 1, 1, 1}, {1, 2, 2, 1}, "VALID"); 1304 INFER_ERROR( 1305 "ExtractImagePatches requires the ksizes attribute to contain 4 values, " 1306 "but got: 5", 1307 op, "[1,7,7,2]"); 1308 } 1309 1310 TEST(ArrayOpsTest, QuantizeAndDequantizeV2_ShapeFn) { 1311 ShapeInferenceTestOp op("QuantizeAndDequantizeV2"); 1312 INFER_OK(op, "?;?;?", "in0"); 1313 INFER_OK(op, "[];?;?", "in0"); 1314 INFER_OK(op, "[1,2,?,4,5];?;?", "in0"); 1315 1316 INFER_ERROR("Shape must be rank 0 but is rank 1", op, "[1,2,?,4,5];[1];[]"); 1317 INFER_ERROR("Shape must be rank 0 but is rank 1", op, "[1,2,?,4,5];[];[1]"); 1318 INFER_ERROR("Shape must be rank 0 but is rank 1", op, "[1,2,?,4,5];[1];[1]"); 1319 } 1320 1321 TEST(ArrayOpsTest, SpaceToBatch_ShapeFn) { 1322 ShapeInferenceTestOp op("SpaceToBatch"); 1323 op.input_tensors.resize(2); 1324 TF_ASSERT_OK(NodeDefBuilder("test", "SpaceToBatch") 1325 .Input("input", 0, DT_FLOAT) 1326 .Input("paddings", 1, DT_INT32) 1327 .Attr("block_size", 2) 1328 .Finalize(&op.node_def)); 1329 1330 // Paddings not known, but batch size can be computed. 1331 INFER_OK(op, "[1,10,10,3];[2,2]", "[4,?,?,d0_3]"); 1332 1333 // Unknown paddings means width and height. 1334 INFER_OK(op, "[1,10,10,3];?", "[4,?,?,d0_3]"); 1335 1336 // Paddings not correct shape 1337 INFER_ERROR("rank", op, "[1,10,10,3];[4]"); 1338 INFER_ERROR("3 and 2", op, "[1,10,10,3];[2,3]"); 1339 1340 Tensor paddings = test::AsTensor<int32>({4, 2, 2, 4}, {{2, 2}}); 1341 op.input_tensors[1] = &paddings; 1342 INFER_OK(op, "[1,10,10,3];[2,2]", "[4,8,8,d0_3]"); 1343 paddings = test::AsTensor<int64>({4, 2, 2, 4}, {{2, 2}}); 1344 INFER_OK(op, "[1,10,10,3];[2,2]", "[4,8,8,d0_3]"); 1345 1346 // Bad paddings values 1347 paddings = test::AsTensor<int32>({1, 2, 3, 4}, {{2, 2}}); 1348 op.input_tensors[1] = &paddings; 1349 INFER_ERROR("Dimension size must be evenly divisible by 2 but is 13", op, 1350 "[1,10,10,3];[2,2]"); 1351 1352 // Negative paddsings 1353 paddings = test::AsTensor<int32>({1, -2, 3, 4}, {{2, 2}}); 1354 op.input_tensors[1] = &paddings; 1355 INFER_ERROR("cannot be negative", op, "[1,10,10,3];[2,2]"); 1356 } 1357 1358 TEST(ArrayOpsTest, SpaceToBatchND_ShapeFn) { 1359 ShapeInferenceTestOp op("SpaceToBatchND"); 1360 op.input_tensors.resize(3); 1361 TF_ASSERT_OK(NodeDefBuilder("test", "SpaceToBatchND") 1362 .Input("input", 0, DT_FLOAT) 1363 .Input("block_shape", 1, DT_INT32) 1364 .Input("paddings", 2, DT_INT32) 1365 .Finalize(&op.node_def)); 1366 1367 // Verify that input shape and paddings shape can be unknown. 1368 INFER_OK(op, "?;[2];?", "?"); 1369 1370 // Only number of input dimensions is known. 1371 INFER_OK(op, "[?,?,?,?];[2];?", "[?,?,?,d0_3]"); 1372 1373 // Dimensions are partially known. 1374 INFER_OK(op, "[?,?,?,2];[2];?", "[?,?,?,d0_3]"); 1375 1376 { 1377 // Dimensions are partially known, block_shape known. 1378 Tensor block_shape = test::AsTensor<int32>({2, 3}); 1379 op.input_tensors[1] = &block_shape; 1380 INFER_OK(op, "[3,?,?,2];[2];?", "[18,?,?,d0_3]"); 1381 1382 // Dimensions are partially known, block_shape and paddings known. 1383 { 1384 Tensor paddings = test::AsTensor<int32>({1, 1, 0, 1}, {{2, 2}}); 1385 op.input_tensors[2] = &paddings; 1386 INFER_OK(op, "[3,?,2,2];[2];[2,2]", "[18,?,1,d0_3]"); 1387 op.input_tensors[2] = nullptr; 1388 } 1389 1390 // Dimensions are fully known, block_shape and paddings are known. 1391 { 1392 Tensor paddings = test::AsTensor<int32>({1, 1, 0, 0}, {{2, 2}}); 1393 op.input_tensors[2] = &paddings; 1394 INFER_OK(op, "[3,2,3,2];[2];[2,2]", "[18,2,1,d0_3]"); 1395 op.input_tensors[2] = nullptr; 1396 } 1397 1398 op.input_tensors[1] = nullptr; 1399 } 1400 1401 INFER_ERROR("block_shape must have rank 1", op, "?;[1,1];?"); 1402 INFER_ERROR("block_shape must have known size", op, "?;[?];?"); 1403 1404 { 1405 Tensor block_shape = test::AsTensor<int32>({0, 2}); 1406 op.input_tensors[1] = &block_shape; 1407 INFER_ERROR("block_shape must be positive", op, "[1,2,2];[2];[2,2]"); 1408 op.input_tensors[1] = nullptr; 1409 } 1410 1411 { 1412 Tensor block_shape = test::AsTensor<int32>({1, 1}); 1413 op.input_tensors[1] = &block_shape; 1414 Tensor paddings = test::AsTensor<int32>({0, -1, 0, 0}, {{2, 2}}); 1415 op.input_tensors[2] = &paddings; 1416 INFER_ERROR("paddings cannot be negative", op, "[1,2,2];[2];[2,2]"); 1417 op.input_tensors[1] = nullptr; 1418 op.input_tensors[2] = nullptr; 1419 } 1420 1421 { 1422 Tensor block_shape = test::AsTensor<int32>({3, 3}); 1423 op.input_tensors[1] = &block_shape; 1424 Tensor paddings = test::AsTensor<int32>({0, 0, 0, 0}, {{2, 2}}); 1425 op.input_tensors[2] = &paddings; 1426 INFER_ERROR("divisible", op, "[1,2,3,1];[2];[2,2]"); 1427 op.input_tensors[1] = nullptr; 1428 op.input_tensors[2] = nullptr; 1429 } 1430 1431 INFER_ERROR("rank", op, "[1,3,3,1];[2];[1]"); 1432 INFER_ERROR("shape", op, "[1,3,3,1];[2];[1,2]"); 1433 } 1434 1435 TEST(ArrayOpsTest, BatchToSpace_ShapeFn) { 1436 ShapeInferenceTestOp op("BatchToSpace"); 1437 op.input_tensors.resize(2); 1438 TF_ASSERT_OK(NodeDefBuilder("test", "BatchToSpace") 1439 .Input("input", 0, DT_FLOAT) 1440 .Input("crops", 1, DT_INT32) 1441 .Attr("block_size", 2) 1442 .Finalize(&op.node_def)); 1443 1444 // croppings not known, but batch size can be computed. 1445 INFER_OK(op, "[4,8,8,3];[2,2]", "[1,?,?,d0_3]"); 1446 1447 // block_size not compatible with batch size 1448 INFER_ERROR("Dimension size must be evenly divisible by", op, 1449 "[5,8,8,3];[2,2]"); 1450 1451 // Unknown croppings means unknown width and height. 1452 INFER_OK(op, "[4,8,8,3];?", "[1,?,?,d0_3]"); 1453 1454 // croppings not correct shape 1455 INFER_ERROR("rank", op, "[4,8,8,3];[4]"); 1456 INFER_ERROR("3 and 2", op, "[4,8,8,3];[2,3]"); 1457 1458 Tensor croppings = test::AsTensor<int64>({4, 2, 2, 4}, {{2, 2}}); 1459 op.input_tensors[1] = &croppings; 1460 INFER_OK(op, "[4,8,8,3];[2,2]", "[1,10,10,d0_3]"); 1461 1462 // Bad croppings values 1463 croppings = test::AsTensor<int32>({100, 2, 3, 4}, {{2, 2}}); 1464 op.input_tensors[1] = &croppings; 1465 INFER_ERROR("Negative dimension size caused by subtracting", op, 1466 "[4,8,8,3];[2,2]"); 1467 croppings = test::AsTensor<int32>({1, 2, 3, 400}, {{2, 2}}); 1468 op.input_tensors[1] = &croppings; 1469 INFER_ERROR("Negative dimension size caused by subtracting", op, 1470 "[4,8,8,3];[2,2]"); 1471 1472 // Negative paddsings 1473 croppings = test::AsTensor<int32>({1, -2, 3, 4}, {{2, 2}}); 1474 op.input_tensors[1] = &croppings; 1475 INFER_ERROR("cannot be negative", op, "[4,8,8,3];[2,2]"); 1476 } 1477 1478 TEST(ArrayOpsTest, BatchToSpaceND_ShapeFn) { 1479 ShapeInferenceTestOp op("BatchToSpaceND"); 1480 op.input_tensors.resize(3); 1481 TF_ASSERT_OK(NodeDefBuilder("test", "BatchToSpaceND") 1482 .Input("input", 0, DT_FLOAT) 1483 .Input("block_shape", 1, DT_INT32) 1484 .Input("crops", 2, DT_INT32) 1485 .Finalize(&op.node_def)); 1486 1487 // Verify that input shape and crops shape can be unknown. 1488 INFER_OK(op, "?;[2];?", "?"); 1489 1490 // Only number of input dimensions is known. 1491 INFER_OK(op, "[?,?,?,?];[2];?", "[?,?,?,d0_3]"); 1492 1493 { 1494 // Dimensions are partially known, block_shape known. 1495 Tensor block_shape = test::AsTensor<int32>({2, 3}); 1496 op.input_tensors[1] = &block_shape; 1497 INFER_OK(op, "[?,?,?,2];[2];?", "[?,?,?,d0_3]"); 1498 1499 INFER_OK(op, "[18,?,?,2];[2];?", "[3,?,?,d0_3]"); 1500 1501 // Dimensions are partially known, block_shape and crops known. 1502 { 1503 Tensor crops = test::AsTensor<int32>({1, 1, 0, 1}, {{2, 2}}); 1504 op.input_tensors[2] = &crops; 1505 INFER_OK(op, "[18,?,2,2];[2];[2,2]", "[3,?,5,d0_3]"); 1506 op.input_tensors[2] = nullptr; 1507 } 1508 1509 // Dimensions are fully known, block_shape and crops are known. 1510 { 1511 Tensor crops = test::AsTensor<int32>({1, 1, 0, 0}, {{2, 2}}); 1512 op.input_tensors[2] = &crops; 1513 INFER_OK(op, "[18,2,1,2];[2];[2,2]", "[3,2,3,d0_3]"); 1514 op.input_tensors[2] = nullptr; 1515 } 1516 1517 op.input_tensors[1] = nullptr; 1518 } 1519 1520 INFER_ERROR("block_shape must have rank 1", op, "?;[1,1];?"); 1521 INFER_ERROR("block_shape must have known size", op, "?;[?];?"); 1522 INFER_ERROR("rank", op, "[2,2];[2];[2,2]"); 1523 INFER_ERROR("rank", op, "[2,2,3];[3];[3,2]"); 1524 1525 { 1526 Tensor block_shape = test::AsTensor<int32>({0, 2}); 1527 op.input_tensors[1] = &block_shape; 1528 INFER_ERROR("block_shape must be positive", op, "[1,2,2];[2];[2,2]"); 1529 op.input_tensors[1] = nullptr; 1530 } 1531 1532 { 1533 Tensor block_shape = test::AsTensor<int32>({1, 1}); 1534 op.input_tensors[1] = &block_shape; 1535 Tensor paddings = test::AsTensor<int32>({0, -1, 0, 0}, {{2, 2}}); 1536 op.input_tensors[2] = &paddings; 1537 INFER_ERROR("crops cannot be negative", op, "[1,2,2];[2];[2,2]"); 1538 op.input_tensors[1] = nullptr; 1539 op.input_tensors[2] = nullptr; 1540 } 1541 1542 // The amount to crop exceeds the padded size. 1543 { 1544 Tensor block_shape = test::AsTensor<int32>({2, 2}); 1545 op.input_tensors[1] = &block_shape; 1546 Tensor crops = test::AsTensor<int32>({3, 2, 0, 0}, {{2, 2}}); 1547 op.input_tensors[2] = &crops; 1548 INFER_ERROR("Negative", op, "[4,2,3,1];[2];[2,2]"); 1549 op.input_tensors[1] = nullptr; 1550 op.input_tensors[2] = nullptr; 1551 } 1552 1553 // The batch size is not divisible by the product of the block_shape. 1554 { 1555 Tensor block_shape = test::AsTensor<int32>({2, 3}); 1556 op.input_tensors[1] = &block_shape; 1557 INFER_ERROR("divisible", op, "[3,1,1,1];[2];[2,2]"); 1558 op.input_tensors[1] = nullptr; 1559 } 1560 } 1561 1562 TEST(ArrayOpsTest, SpaceToDepth_ShapeFn) { 1563 ShapeInferenceTestOp op("SpaceToDepth"); 1564 TF_ASSERT_OK(NodeDefBuilder("test", "SpaceToDepth") 1565 .Input("input", 0, DT_FLOAT) 1566 .Attr("block_size", 2) 1567 .Finalize(&op.node_def)); 1568 1569 INFER_OK(op, "[1,2,4,4]", "[d0_0,1,2,16]"); 1570 1571 // block_size not compatible with space 1572 INFER_ERROR("Dimension size must be evenly divisible by 2 but is 3", op, 1573 "[1,3,8,4]"); 1574 INFER_ERROR("Dimension size must be evenly divisible by 2 but is 5", op, 1575 "[1,2,5,4]"); 1576 1577 // Unknown depth --> Unknown depth. 1578 INFER_OK(op, "[1,2,4,?]", "[d0_0,1,2,?]"); 1579 } 1580 1581 TEST(ArrayOpsTest, DepthToSpace_ShapeFn) { 1582 ShapeInferenceTestOp op("DepthToSpace"); 1583 TF_ASSERT_OK(NodeDefBuilder("test", "DepthToSpace") 1584 .Input("input", 0, DT_FLOAT) 1585 .Attr("block_size", 2) 1586 .Finalize(&op.node_def)); 1587 1588 INFER_OK(op, "[1,1,2,16]", "[d0_0,2,4,4]"); 1589 1590 // Bad depth 1591 INFER_ERROR("Dimension size must be evenly divisible by 4 but is 15", op, 1592 "[1,1,2,15]"); 1593 1594 // Unknown depth --> Unknown depth. 1595 INFER_OK(op, "[1,2,4,?]", "[d0_0,4,8,?]"); 1596 1597 // Check another block size. 1598 TF_ASSERT_OK(NodeDefBuilder("test", "DepthToSpace") 1599 .Input("input", 0, DT_FLOAT) 1600 .Attr("block_size", 10) 1601 .Finalize(&op.node_def)); 1602 INFER_OK(op, "[1,1,2,200]", "[d0_0,10,20,2]"); 1603 } 1604 1605 TEST(ArrayOpsTest, Slice_ShapeFn) { 1606 ShapeInferenceTestOp op("Slice"); 1607 TF_ASSERT_OK(NodeDefBuilder("test", "Slice") 1608 .Input("input", 0, DT_FLOAT) 1609 .Input("begin", 1, DT_INT64) 1610 .Input("sizes", 2, DT_INT64) 1611 .Finalize(&op.node_def)); 1612 1613 // Known rank of input and shape of begin/sizes, but unknown values. 1614 // The best we know is the rank of the output. 1615 INFER_OK(op, "[2,3,4,5];[4];[4]", "[?,?,?,?]"); 1616 1617 // Unknown shape of begin/sizes, we still know the rank. 1618 INFER_OK(op, "[2,3,4,5];[?];[?]", "[?,?,?,?]"); 1619 // Unknown all around 1620 INFER_OK(op, "?;[?];[?]", "?"); 1621 // Can infer based on begin 1622 INFER_OK(op, "?;[4];[?]", "[?,?,?,?]"); 1623 1624 // Bad rank of begin, sizes 1625 INFER_ERROR("must be rank 1", op, "[2,3,4,5];[2,3];[3]"); 1626 INFER_ERROR("must be rank 1", op, "[2,3,4,5];[2];[3,4]"); 1627 // Length of begin doesn't match input rank 1628 INFER_ERROR("must be rank 2", op, "[2,3,4,5];[2];[2]"); 1629 1630 // Tests with known values. 1631 op.input_tensors.resize(3); 1632 Tensor begin = test::AsTensor<int32>({0, 1, 2, 1}); 1633 Tensor sizes = test::AsTensor<int32>({1, 2, 1, 3}); 1634 op.input_tensors[1] = &begin; 1635 op.input_tensors[2] = &sizes; 1636 INFER_OK(op, "[2,3,4,5];[4];[4]", "[1,2,1,3]"); 1637 1638 // -1 in sizes means "get the rest" 1639 sizes = test::AsTensor<int32>({-1, -1, 1, -1}); 1640 INFER_OK(op, "[2,3,4,5];[4];[4]", "[d0_0,2,1,4]"); 1641 1642 begin = test::AsTensor<int32>({0, 1, 2, 6}); 1643 sizes = test::AsTensor<int32>({-1, -1, -1, -1}); 1644 INFER_ERROR("Negative dimension size", op, "[2,3,4,5];[4];[4]"); 1645 1646 begin = test::AsTensor<int32>({0, 1, 2, 5}); 1647 sizes = test::AsTensor<int32>({-1, -1, -1, -2}); 1648 INFER_ERROR("cannot be < -1", op, "[2,3,4,5];[4];[4]"); 1649 } 1650 1651 TEST(ArrayOpsTest, StridedSlice_ShapeFn) { 1652 ShapeInferenceTestOp op("StridedSlice"); 1653 TF_ASSERT_OK(NodeDefBuilder("test", "StridedSlice") 1654 .Input("input", 0, DT_FLOAT) 1655 .Input("begin", 1, DT_INT32) 1656 .Input("end", 2, DT_INT32) 1657 .Input("strides", 3, DT_INT32) 1658 .Attr("shrink_axis_mask", 1) 1659 .Finalize(&op.node_def)); 1660 op.input_tensors.resize(4); 1661 Tensor strides = test::AsTensor<int32>({1}); 1662 op.input_tensors[3] = &strides; 1663 // Slicing on the 0-th dimension. 1664 INFER_OK(op, "[2,3,4,5];[1];[1];[1]", "[3,4,5]"); 1665 // Slicing on the 0-th dimension. This time some of the result dimension is 0. 1666 INFER_OK(op, "[2,0,3,4];[1];[1];[1]", "[0,3,4]"); 1667 } 1668 1669 TEST(ArrayOpsTest, StridedSliceGrad_ShapeFn) { 1670 ShapeInferenceTestOp op("StridedSliceGrad"); 1671 op.input_tensors.resize(5); 1672 INFER_OK(op, "?;?;?;?;?", "?"); 1673 INFER_OK(op, "[?];?;?;?;?", "?"); 1674 INFER_OK(op, "[4];?;?;?;?", "[?,?,?,?]"); 1675 1676 Tensor in_t = test::AsTensor<int32>({1, 2, 3, 4}); 1677 op.input_tensors[0] = &in_t; 1678 INFER_OK(op, "[4];?;?;?;?", "[1,2,3,4]"); 1679 } 1680 1681 TEST(ArrayOpsTest, UnchangedWithQuantizationScalars_ShapeFn) { 1682 for (const char* op_name : {"Dequantize", "FakeQuantWithMinMaxVars"}) { 1683 ShapeInferenceTestOp op(op_name); 1684 1685 INFER_OK(op, "?;?;?", "in0"); 1686 INFER_OK(op, "[1,?,3];[];[]", "in0"); 1687 1688 // Rank check scalars. 1689 INFER_ERROR("be rank 0", op, "[1,?,3];[1];[]"); 1690 INFER_ERROR("be rank 0", op, "[1,?,3];[];[1]"); 1691 } 1692 } 1693 1694 TEST(ArrayOpsTest, FakeQuantWithMinMaxVarsPerChannel) { 1695 ShapeInferenceTestOp op("FakeQuantWithMinMaxVarsPerChannel"); 1696 1697 INFER_OK(op, "?;?;?", "in0"); 1698 INFER_OK(op, "[?];?;?", "in0"); 1699 INFER_OK(op, "[1,?,3];[3];[3]", "in0"); 1700 INFER_OK(op, "[3];[3];[3]", "in0"); 1701 1702 // Rank check vectors. 1703 INFER_ERROR("be rank 1", op, "[1,?,3];[1];[]"); 1704 INFER_ERROR("be rank 1", op, "[1,?,3];[];[1]"); 1705 1706 // Vectors must match each other, and match last dim of input. 1707 INFER_ERROR("must be equal", op, "[1,?,3];[2];[?]"); 1708 INFER_ERROR("must be equal", op, "[1,?,3];[?];[2]"); 1709 INFER_ERROR("must be equal", op, "[1,?,?];[1];[2]"); 1710 INFER_ERROR("must be equal", op, "[5];[4];[?]"); 1711 } 1712 1713 TEST(ArrayOpsTest, FakeQuantWithMinMaxVarsPerChannelGradient) { 1714 ShapeInferenceTestOp op("FakeQuantWithMinMaxVarsPerChannelGradient"); 1715 1716 INFER_OK(op, "?;?;?;?", "in0;[?];[?]"); 1717 INFER_OK(op, "[3];[3];[3];[3]", "in0;in3;in3"); 1718 INFER_OK(op, "[1,3];[1,3];[3];[3]", "in0;in3;in3"); 1719 INFER_OK(op, "[1,2,3,4];[1,2,3,4];[4];[4]", "in0;in3;in3"); 1720 1721 // Rank check vectors. 1722 INFER_ERROR("be equal rank", op, "[1,?,3];[1,?,3];[3];[]"); 1723 INFER_ERROR("be rank 1", op, "[1,?,3];[1,?,3];[];[3]"); 1724 INFER_ERROR("be at least rank 1", op, "[];[];[1];[1]"); 1725 INFER_ERROR("be at most rank 4", op, "[1,2,3,4,5];[1,2,3,4,5];[1];[1]"); 1726 1727 // Vectors must match each other, and match last dim of input. 1728 INFER_ERROR("must be equal", op, "[1,3];[1,3];[2];[3]"); 1729 INFER_ERROR("must be equal", op, "[1,3];[1,3];[3];[2]"); 1730 } 1731 1732 TEST(ArrayOpsTest, QuantizedConcat_ShapeFn) { 1733 ShapeInferenceTestOp op("QuantizedConcat"); 1734 auto set_n = [&op](int n) { 1735 std::vector<NodeDefBuilder::NodeOut> src_list; 1736 std::vector<NodeDefBuilder::NodeOut> limit_list; 1737 for (int i = 0; i < n; ++i) { 1738 src_list.emplace_back("a", 0, DT_QUINT8); 1739 limit_list.emplace_back("b", 0, DT_FLOAT); 1740 } 1741 TF_ASSERT_OK(NodeDefBuilder("test", "QuantizedConcat") 1742 .Input({"concat_dim", 0, DT_INT32}) 1743 .Input(src_list) 1744 .Input(limit_list) 1745 .Input(limit_list) 1746 .Attr("N", n) 1747 .Finalize(&op.node_def)); 1748 }; 1749 1750 // Confirm dimension[0] of the input (the concat_dim) is a scalar. 1751 set_n(1); 1752 INFER_ERROR("Shape must be rank 0 but is rank 1", op, "[1];?;?;?"); 1753 1754 // Last 2*<N> are all scalars. 1755 set_n(2); 1756 INFER_ERROR("must be rank 0", op, "[];?;?;?;?;?;[1]"); 1757 INFER_ERROR("must be rank 0", op, "[];?;?;?;?;[1];?"); 1758 INFER_ERROR("must be rank 0", op, "[];?;?;?;[1];?;?"); 1759 INFER_ERROR("must be rank 0", op, "[];?;?;[1];?;?;?"); 1760 1761 // First is concat dim; next N must be compatible for concat. 1762 set_n(2); 1763 INFER_ERROR("must be rank 2", op, "[];[1,2];[1,2,3];?;?;?;?"); 1764 INFER_OK(op, "[];[1,2];[1,3];?;?;?;?", "[?,?];[];[]"); 1765 1766 // Test when the concat_dim tensor is known. The concatenated dimension is 1767 // summed across all input tensors, and other dimensions are merged. 1768 Tensor concat_dim_t; 1769 op.input_tensors.push_back(&concat_dim_t); 1770 set_n(2); 1771 concat_dim_t = test::AsScalar(0); // Sum dim 0, merge the other two dims. 1772 INFER_OK(op, "[];[100,2,?];[10,?,3];?;?;?;?", "[110,d1_1,d2_2];[];[]"); 1773 INFER_ERROR("Dimension 1 in both shapes must be equal, but are 5 and 3", op, 1774 "[];[100,2,5];[10,?,3];?;?;?;?"); 1775 // Note that other cases of concat are covered in the Concat tests. 1776 } 1777 1778 TEST(StateOpsTest, _ParallelConcatStart_ShapeFn) { 1779 ShapeInferenceTestOp op("_ParallelConcatStart"); 1780 TensorShape shape({1, 2, 3}); 1781 TensorShapeProto shape_proto; 1782 shape.AsProto(&shape_proto); 1783 TF_ASSERT_OK(NodeDefBuilder("test", "_ParallelConcatStart") 1784 .Attr("shape", shape_proto) 1785 .Attr("dtype", DT_FLOAT) 1786 .Finalize(&op.node_def)); 1787 INFER_OK(op, "", "[1,2,3]"); 1788 } 1789 1790 } // end namespace tensorflow 1791
