| 24 | namespace tensorflow { |
| 25 | |
| 26 | TEST(NNOpsTest, TopK_ShapeFn) { |
| 27 | ShapeInferenceTestOp op("TopK"); |
| 28 | auto set_k = [&op](int k) { |
| 29 | TF_ASSERT_OK(NodeDefBuilder("test", "Pack") |
| 30 | .Input({{"a", 0, DT_FLOAT}}) |
| 31 | .Attr("k", k) |
| 32 | .Finalize(&op.node_def)); |
| 33 | }; |
| 34 | |
| 35 | set_k(20); |
| 36 | // With known input, each output is an unknown shape. |
| 37 | INFER_OK(op, "?", "?;?"); |
| 38 | // With vector input, each output is [k]. |
| 39 | INFER_OK(op, "[20]", "[20];[20]"); |
| 40 | INFER_OK(op, "[21]", "[20];[20]"); |
| 41 | |
| 42 | // With input rank 3, each output is the two first 2 dims of input, plus k. |
| 43 | INFER_OK(op, "[1,?,21]", "[d0_0,d0_1,20];[d0_0,d0_1,20]"); |
| 44 | // With input rank 4, each output is the two first 3 dims of input, plus k. |
| 45 | INFER_OK(op, "[1,?,21,?]", "[d0_0,d0_1,d0_2,20];[d0_0,d0_1,d0_2,20]"); |
| 46 | |
| 47 | INFER_ERROR("Shape must be at least rank 1 but is rank 0", op, "[]"); |
| 48 | INFER_ERROR("input must have last dimension >= k = 20 but is 1", op, "[1]"); |
| 49 | INFER_ERROR("input must have last dimension >= k = 20 but is 4", op, |
| 50 | "[1,2,3,4]"); |
| 51 | set_k(-1); |
| 52 | INFER_ERROR("Need k >= 0, got -1", op, "[1,2,3,4]"); |
| 53 | } |
| 54 | |
| 55 | TEST(NNOpsTest, TopKV2_ShapeFn) { |
| 56 | ShapeInferenceTestOp op("TopKV2"); |