[tf.data] Add AssertCardinalityDataset to list of passthrough datasets in autosharding

PiperOrigin-RevId: 307150560 Change-Id: I1b270211d70ac2eddff0f6c3dfb575e58dfc4f7f
2020-04-17 18:30:43 -07:00 · 2020-04-17 18:30:43 -07:00 · a8acb1d990
commit a8acb1d990
parent 77efb99516
1 changed files with 2 additions and 1 deletions
--- a/tensorflow/core/grappler/optimizers/data/auto_shard.cc
+++ b/tensorflow/core/grappler/optimizers/data/auto_shard.cc
@ -57,8 +57,9 @@ constexpr std::array<const char*, 2> kMultipleInputsDatasetOps = {
    "ZipDataset"
 };

-constexpr std::array<const char*, 29> kPassThroughOps = {
+constexpr std::array<const char*, 30> kPassThroughOps = {
    "_Retval",
+    "AssertCardinalityDataset",
    "AssertNextDataset",
    "BatchDataset",
    "BatchDatasetV2",