| 29 | namespace { |
| 30 | |
| 31 | TEST(ParallelBatch, Batch) { |
| 32 | using test::function::NDef; |
| 33 | GrapplerItem item; |
| 34 | item.graph = test::function::GDef( |
| 35 | {NDef("start", "Const", {}, {{"value", 0}, {"dtype", DT_INT32}}), |
| 36 | NDef("stop", "Const", {}, {{"value", 10}, {"dtype", DT_INT32}}), |
| 37 | NDef("step", "Const", {}, {{"value", 1}, {"dtype", DT_INT32}}), |
| 38 | NDef("range", "RangeDataset", {"start", "stop", "step"}, {}), |
| 39 | NDef("batch_size", "Const", {}, {{"value", 5}, {"dtype", DT_INT32}}), |
| 40 | NDef("drop_remainder", "Const", {}, |
| 41 | {{"value", false}, {"dtype", DT_BOOL}}), |
| 42 | NDef("batch", "BatchDatasetV2", |
| 43 | {"range", "batch_size", "drop_remainder"}, {})}); |
| 44 | |
| 45 | ParallelBatch optimizer; |
| 46 | GraphDef output; |
| 47 | TF_ASSERT_OK(optimizer.Optimize(nullptr, item, &output)); |
| 48 | EXPECT_TRUE(graph_utils::ContainsGraphNodeWithName("batch", output)); |
| 49 | int index = graph_utils::FindGraphNodeWithName("batch", output); |
| 50 | EXPECT_TRUE(output.node(index).attr().at("parallel_copy").b()); |
| 51 | } |
| 52 | |
| 53 | TEST(ParallelBatch, PaddedBatch) { |
| 54 | using test::function::NDef; |
nothing calls this directly
no test coverage detected