def _test_objects():   ]  +def _test_eager_objects():+  return [

 def testDatasetCompression(self, element):     dataset = dataset.map(lambda x: compression_ops.uncompress(x, element_spec))     self.assertDatasetProduces(dataset, [element]) +  @combinations.generate(combinations.times(

 def _test_objects():   ]  +def _test_eager_objects():+  return [

 from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import combinations from tensorflow.python.framework import dtypes+from tensorflow.python.framework import errors

 def testSimpleString(self):         (["foo", "bar", "baz", "baz", "bar", "foo"], ["foo", "bar", "baz"]),     ]) +  def _checkDatasetRaises(self, dtype, test_cases, error):+    """Test whether the dataset raises the appropriate errors+    while generating the outputs.++    Args:+      dtype: The actual `dtype` of the elements in each test case.+      test_cases: A list of lists. The dataset will be created from the list items.+      error: The expected error to be raised when a corrupted item in encountered.+    """++    current_test_case = []+    dataset = dataset_ops.Dataset.from_generator(lambda: current_test_case,+                                                 dtype).apply(unique.unique())++    for test_case in test_cases:+      current_test_case = test_case+      with self.assertRaises(error):+        _ = self.getDatasetOutput(dataset)

 def testSimpleString(self):         (["foo", "bar", "baz", "baz", "bar", "foo"], ["foo", "bar", "baz"]),     ]) +  def _checkDatasetRaises(self, dtype, test_cases, error):+    """Test whether the dataset raises the appropriate errors+    while generating the outputs.++    Args:+      dtype: The actual `dtype` of the elements in each test case.

 def testSimpleString(self):         (["foo", "bar", "baz", "baz", "bar", "foo"], ["foo", "bar", "baz"]),     ]) +  def _checkDatasetRaises(self, dtype, test_cases, error):+    """Test whether the dataset raises the appropriate errors+    while generating the outputs.++    Args:+      dtype: The actual `dtype` of the elements in each test case.+      test_cases: A list of lists. The dataset will be created from the list items.+      error: The expected error to be raised when a corrupted item in encountered.+    """++    current_test_case = []+    dataset = dataset_ops.Dataset.from_generator(lambda: current_test_case,+                                                 dtype).apply(unique.unique())++    for test_case in test_cases:+      current_test_case = test_case+      with self.assertRaises(error):+        _ = self.getDatasetOutput(dataset)++  @combinations.generate(test_base.graph_only_combinations())+  def testStringTypeMismatch(self):+    """Should raise InternalError when element type doesn't match+    with dtypes.string."""++    test_cases = [+        ["hello", 1],+        ["hello", "hello", "world", 3],+        ["hello", 1, 1],+        ["hello", "world", 1, 2],+        [1, "hello"],+        [1, 2, "hello"],+        [1, 3, "hello", "world"],+        [1, 1, "hello", "hello"]+    ]+    self._checkDatasetRaises(dtype=dtypes.string, test_cases=test_cases,+                             error=errors.InternalError)++  @combinations.generate(test_base.graph_only_combinations())+  def testInt32TypeMismatch(self):+    """Should raise InvalidArgumentError when element type doesn't+    match with dtypes.int32"""++    test_cases = [+        [1, "foo"],+        [1, 2, "bar"],+        [1, 3, "foo", "bar"],+        [1, 4, "foo", "foo"],+        ["bar", 1],+        ["bar", "foo", 2],+        ["bar", "bar", "foo", 3],+        ["foo", 1, 1],+        ["bar", "bar", 1, 1],+    ]+    self._checkDatasetRaises(dtype=dtypes.int32, test_cases=test_cases,+                             error=errors.InvalidArgumentError)++  @combinations.generate(test_base.graph_only_combinations())+  def testInt64TypeMismatch(self):+    """Should raise InvalidArgumentError when element type doesn't+    match with dtypes.int64."""++    test_cases = [+        [2, "hello"],+        [3, 2, "hello"],+        [5, 3, "hello", "world"],+        [6, 7, "hello", "hello"],+        ["hello", 6],+        ["hello", "world", 8],+        ["hello", "hello", "world", 8],+        ["hello", 9, 9],+        ["hello", "world", 10, 10],+    ]+    self._checkDatasetRaises(dtype=dtypes.int64, test_cases=test_cases,+                             error=errors.InvalidArgumentError)++  @combinations.generate(test_base.graph_only_combinations())+  def testUnsupportedTypes(self):+    """Should raise TypeError when element type doesn't match with the+    dtypes.int64, dtypes.int32 or dtypes.string (supported types)."""++    sample_unsupported_types = [dtypes.bool, dtypes.double, dtypes.complex64,+                                dtypes.float32, dtypes.float64, dtypes.qint16, dtypes.qint32]+    current_test_case = []

 def testSimpleString(self):         (["foo", "bar", "baz", "baz", "bar", "foo"], ["foo", "bar", "baz"]),     ]) +  def _checkDatasetRaises(self, dtype, test_cases, error):+    """Test whether the dataset raises the appropriate errors+    while generating the outputs.++    Args:+      dtype: The actual `dtype` of the elements in each test case.+      test_cases: A list of lists. The dataset will be created from the list items.+      error: The expected error to be raised when a corrupted item in encountered.+    """++    current_test_case = []+    dataset = dataset_ops.Dataset.from_generator(lambda: current_test_case,+                                                 dtype).apply(unique.unique())++    for test_case in test_cases:+      current_test_case = test_case+      with self.assertRaises(error):+        _ = self.getDatasetOutput(dataset)++  @combinations.generate(test_base.graph_only_combinations())+  def testStringTypeMismatch(self):+    """Should raise InternalError when element type doesn't match+    with dtypes.string."""++    test_cases = [+        ["hello", 1],+        ["hello", "hello", "world", 3],+        ["hello", 1, 1],+        ["hello", "world", 1, 2],+        [1, "hello"],+        [1, 2, "hello"],+        [1, 3, "hello", "world"],+        [1, 1, "hello", "hello"]+    ]+    self._checkDatasetRaises(dtype=dtypes.string, test_cases=test_cases,+                             error=errors.InternalError)++  @combinations.generate(test_base.graph_only_combinations())+  def testInt32TypeMismatch(self):+    """Should raise InvalidArgumentError when element type doesn't+    match with dtypes.int32"""++    test_cases = [+        [1, "foo"],+        [1, 2, "bar"],+        [1, 3, "foo", "bar"],+        [1, 4, "foo", "foo"],+        ["bar", 1],+        ["bar", "foo", 2],+        ["bar", "bar", "foo", 3],+        ["foo", 1, 1],+        ["bar", "bar", 1, 1],+    ]+    self._checkDatasetRaises(dtype=dtypes.int32, test_cases=test_cases,+                             error=errors.InvalidArgumentError)

 def testSimpleString(self):         (["foo", "bar", "baz", "baz", "bar", "foo"], ["foo", "bar", "baz"]),     ]) +  def _checkDatasetRaises(self, dtype, test_cases, error):+    """Test whether the dataset raises the appropriate errors+    while generating the outputs.++    Args:+      dtype: The actual `dtype` of the elements in each test case.+      test_cases: A list of lists. The dataset will be created from the list items.+      error: The expected error to be raised when a corrupted item in encountered.+    """++    current_test_case = []+    dataset = dataset_ops.Dataset.from_generator(lambda: current_test_case,+                                                 dtype).apply(unique.unique())++    for test_case in test_cases:+      current_test_case = test_case+      with self.assertRaises(error):+        _ = self.getDatasetOutput(dataset)++  @combinations.generate(test_base.graph_only_combinations())+  def testStringTypeMismatch(self):+    """Should raise InternalError when element type doesn't match+    with dtypes.string."""++    test_cases = [+        ["hello", 1],+        ["hello", "hello", "world", 3],+        ["hello", 1, 1],+        ["hello", "world", 1, 2],+        [1, "hello"],+        [1, 2, "hello"],+        [1, 3, "hello", "world"],+        [1, 1, "hello", "hello"]+    ]+    self._checkDatasetRaises(dtype=dtypes.string, test_cases=test_cases,+                             error=errors.InternalError)++  @combinations.generate(test_base.graph_only_combinations())+  def testInt32TypeMismatch(self):+    """Should raise InvalidArgumentError when element type doesn't+    match with dtypes.int32"""++    test_cases = [+        [1, "foo"],+        [1, 2, "bar"],+        [1, 3, "foo", "bar"],+        [1, 4, "foo", "foo"],+        ["bar", 1],+        ["bar", "foo", 2],+        ["bar", "bar", "foo", 3],+        ["foo", 1, 1],+        ["bar", "bar", 1, 1],+    ]+    self._checkDatasetRaises(dtype=dtypes.int32, test_cases=test_cases,+                             error=errors.InvalidArgumentError)++  @combinations.generate(test_base.graph_only_combinations())+  def testInt64TypeMismatch(self):

 def testSimpleString(self):         (["foo", "bar", "baz", "baz", "bar", "foo"], ["foo", "bar", "baz"]),     ]) +  def _checkDatasetRaises(self, dtype, test_cases, error):+    """Test whether the dataset raises the appropriate errors+    while generating the outputs.++    Args:+      dtype: The actual `dtype` of the elements in each test case.+      test_cases: A list of lists. The dataset will be created from the list items.+      error: The expected error to be raised when a corrupted item in encountered.

 def testSimpleString(self):         (["foo", "bar", "baz", "baz", "bar", "foo"], ["foo", "bar", "baz"]),     ]) +  @combinations.generate(test_base.graph_only_combinations())+  def testTypeMismatch(self):++    # Placeholder values are needed to fill in the expected array with dummy value so that,+    # when the dataset generates the element and observes that there is a type mismatch,+    # it raises the proper error and not an OutOfRangeError which occurs when it is unable+    # to fetch an element to compare from the expected array in the first place.+    string_placeholder = ""+    int32_placeholder = 0+    int64_placeholder = 0

 def testSimpleString(self):         (["foo", "bar", "baz", "baz", "bar", "foo"], ["foo", "bar", "baz"]),     ]) +  @combinations.generate(test_base.graph_only_combinations())+  def testTypeMismatch(self):++    # raises InternalError when dtypes don't match.+    with self.assertRaises(errors.InternalError):+      self._testSimpleHelper(dtypes.string, [+          (["hello", 1, 2, 1], ["hello"]),+          (["hello", "world", 1], ["hello", "world"]),+          (["hello", "hello", "world", 1, 2], ["hello", "world"]),+          (["hello", "world", 1, 1, 2], ["hello", "world"]),+          ([1, 2, "hello"], ["hello"]),+          ([1, 1, 2, 3, 3, "hello"], ["hello"]),+      ])++      self._testSimpleHelper(dtypes.int32, [+          ([1, "hello", "world"], [1]),+          ([1, 2, 1, "hello", "hello", "world"], [1, 2]),+          (["hello", 1, 2], [1, 2]),+          (["hello", 1, 1, 2, 3, 3], [1, 2, 3]),+      ])++      self._testSimpleHelper(dtypes.int64, [+          ([2, 3, "hello", "world"], [2, 3]),+          ([2, 3, 3, "hello", "hello", "world"], [2, 3]),+          (["hello", 2, 2], [2]),+          (["hello", "hello", 1, 1, 2, 3], [1, 2, 3]),

 def testSimpleString(self):         (["foo", "bar", "baz", "baz", "bar", "foo"], ["foo", "bar", "baz"]),     ]) +  @combinations.generate(test_base.graph_only_combinations())+  def testTypeMismatch(self):++    # raises InternalError when dtypes don't match.+    with self.assertRaises(errors.InternalError):+      self._testSimpleHelper(dtypes.string, [+          (["hello", 1, 2, 1], ["hello"]),+          (["hello", "world", 1], ["hello", "world"]),+          (["hello", "hello", "world", 1, 2], ["hello", "world"]),+          (["hello", "world", 1, 1, 2], ["hello", "world"]),+          ([1, 2, "hello"], ["hello"]),+          ([1, 1, 2, 3, 3, "hello"], ["hello"]),+      ])++      self._testSimpleHelper(dtypes.int32, [+          ([1, "hello", "world"], [1]),+          ([1, 2, 1, "hello", "hello", "world"], [1, 2]),+          (["hello", 1, 2], [1, 2]),+          (["hello", 1, 1, 2, 3, 3], [1, 2, 3]),+      ])++      self._testSimpleHelper(dtypes.int64, [+          ([2, 3, "hello", "world"], [2, 3]),+          ([2, 3, 3, "hello", "hello", "world"], [2, 3]),+          (["hello", 2, 2], [2]),

 def testSimpleString(self):         (["foo", "bar", "baz", "baz", "bar", "foo"], ["foo", "bar", "baz"]),     ]) +  @combinations.generate(test_base.graph_only_combinations())+  def testTypeMismatch(self):++    # raises InternalError when dtypes don't match.+    with self.assertRaises(errors.InternalError):+      self._testSimpleHelper(dtypes.string, [+          (["hello", 1, 2, 1], ["hello"]),+          (["hello", "world", 1], ["hello", "world"]),+          (["hello", "hello", "world", 1, 2], ["hello", "world"]),+          (["hello", "world", 1, 1, 2], ["hello", "world"]),+          ([1, 2, "hello"], ["hello"]),+          ([1, 1, 2, 3, 3, "hello"], ["hello"]),+      ])++      self._testSimpleHelper(dtypes.int32, [+          ([1, "hello", "world"], [1]),+          ([1, 2, 1, "hello", "hello", "world"], [1, 2]),+          (["hello", 1, 2], [1, 2]),+          (["hello", 1, 1, 2, 3, 3], [1, 2, 3]),

 def testSimpleString(self):         (["foo", "bar", "baz", "baz", "bar", "foo"], ["foo", "bar", "baz"]),     ]) +  @combinations.generate(test_base.graph_only_combinations())+  def testTypeMismatch(self):++    # raises InternalError when dtypes don't match.+    with self.assertRaises(errors.InternalError):+      self._testSimpleHelper(dtypes.string, [+          (["hello", 1, 2, 1], ["hello"]),+          (["hello", "world", 1], ["hello", "world"]),+          (["hello", "hello", "world", 1, 2], ["hello", "world"]),+          (["hello", "world", 1, 1, 2], ["hello", "world"]),+          ([1, 2, "hello"], ["hello"]),

 def testSimpleString(self):         (["foo", "bar", "baz", "baz", "bar", "foo"], ["foo", "bar", "baz"]),     ]) +  @combinations.generate(test_base.graph_only_combinations())+  def testTypeMismatch(self):++    # raises InternalError when dtypes don't match.+    with self.assertRaises(errors.InternalError):+      self._testSimpleHelper(dtypes.string, [+          (["hello", 1, 2, 1], ["hello"]),+          (["hello", "world", 1], ["hello", "world"]),+          (["hello", "hello", "world", 1, 2], ["hello", "world"]),+          (["hello", "world", 1, 1, 2], ["hello", "world"]),+          ([1, 2, "hello"], ["hello"]),+          ([1, 1, 2, 3, 3, "hello"], ["hello"]),+      ])++      self._testSimpleHelper(dtypes.int32, [+          ([1, "hello", "world"], [1]),+          ([1, 2, 1, "hello", "hello", "world"], [1, 2]),+          (["hello", 1, 2], [1, 2]),

 def testSimpleString(self):         (["foo", "bar", "baz", "baz", "bar", "foo"], ["foo", "bar", "baz"]),     ]) +  @combinations.generate(test_base.graph_only_combinations())+  def testTypeMismatch(self):++    # raises InternalError when dtypes don't match.+    with self.assertRaises(errors.InternalError):+      self._testSimpleHelper(dtypes.string, [+          (["hello", 1, 2, 1], ["hello"]),+          (["hello", "world", 1], ["hello", "world"]),+          (["hello", "hello", "world", 1, 2], ["hello", "world"]),+          (["hello", "world", 1, 1, 2], ["hello", "world"]),+          ([1, 2, "hello"], ["hello"]),+          ([1, 1, 2, 3, 3, "hello"], ["hello"]),

 def testSimpleString(self):         (["foo", "bar", "baz", "baz", "bar", "foo"], ["foo", "bar", "baz"]),     ]) +  @combinations.generate(test_base.graph_only_combinations())+  def testTypeMismatch(self):++    # raises InternalError when dtypes don't match.+    # NOTE: Generating the following expected outputs can be considered/taken up as an

 def _testSimpleHelper(self, dtype, test_cases):     for test_case, expected in test_cases:       current_test_case = test_case       self.assertDatasetProduces(dataset, [-          compat.as_bytes(element) if dtype == dtypes.string else element+          compat.as_bytes(+              element) if dtype == dtypes.string else element

 class ShardDatasetOp::Dataset : public DatasetBase {         return Status::OK();       } -      std::vector<Tensor> result;-      do {-        result.clear();-        TF_RETURN_IF_ERROR(input_impl_->GetNext(ctx, &result, end_of_sequence));-        if (*end_of_sequence) {-          input_impl_.reset();-          return Status::OK();-        }-      } while ((next_index_++ % dataset()->num_shards_) != dataset()->index_);+      int num_to_skip = (dataset()->index_ - next_index_) %+                        dataset()->num_shards_;+      if (num_to_skip < 0) {+        num_to_skip += dataset()->num_shards_;+      }+      int num_skipped;+      TF_RETURN_IF_ERROR(input_impl_->Skip(ctx, num_to_skip, end_of_sequence,+                                           &num_skipped));+      next_index_ += num_skipped;

 Status RecordReader::ReadRecord(uint64* offset, tstring* record) {   return Status::OK(); } -Status RecordReader::SkipRecords(uint64* offset, int num_to_skip) {+Status RecordReader::SkipRecords(uint64* offset, int num_to_skip,+                                 int* num_skipped) {   TF_RETURN_IF_ERROR(PositionInputStream(*offset));    Status s;   tstring record;+  *num_skipped = 0;   for (int i = 0; i < num_to_skip; ++i) {     s = ReadChecksummed(*offset, sizeof(uint64), &record);     if (!s.ok()) {       last_read_failed_ = true;       return s;     }     const uint64 length = core::DecodeFixed64(record.data());-    input_stream_->SkipNBytes(length + kFooterSize);++    // Skip data+    s = input_stream_->SkipNBytes(length + kFooterSize);+    if (!s.ok()) {+      last_read_failed_ = true;+      if (errors::IsOutOfRange(s)) {+        s = errors::DataLoss("truncated record at ", *offset);

 Status RecordReader::ReadRecord(uint64* offset, tstring* record) {   return Status::OK(); } +Status RecordReader::SkipRecords(uint64* offset, int num_to_skip) {

 TEST(RecordReaderWriterTest, TestBasics) {   } } +TEST(RecordReaderWriterTest, TestSkip) {

 Status RecordReader::ReadRecord(uint64* offset, tstring* record) {   return Status::OK(); } +Status RecordReader::SkipRecords(uint64* offset, int num_to_skip) {+  TF_RETURN_IF_ERROR(PositionInputStream(*offset));++  Status s;+  tstring record;+  for (int i = 0; i < num_to_skip; ++i) {+    s = ReadChecksummed(*offset, sizeof(uint64), &record);+    if (!s.ok()) {+      last_read_failed_ = true;+      return s;+    }+    const uint64 length = core::DecodeFixed64(record.data());+    input_stream_->SkipNBytes(length + kFooterSize);+    *offset += kHeaderSize + length + kFooterSize;+    DCHECK_EQ(*offset, input_stream_->Tell());

 Status DatasetBaseIterator::GetNext(IteratorContext* ctx,   return s; } +Status DatasetBaseIterator::Skip(IteratorContext* ctx, int num_to_skip,+                                 bool* end_of_sequence, int* num_skipped) {+  profiler::TraceMe activity([&] { return BuildTraceMeName(); },+                             profiler::TraceMeLevel::kInfo);+  DVLOG(3) << prefix() << " Skip enter";+  RecordStart(ctx, /*stop_output=*/true);+  Status s = SkipInternal(ctx, num_to_skip, end_of_sequence, num_skipped);+  if (s.ok() && !*end_of_sequence) RecordElement(ctx);

 class ShardDatasetOp::Dataset : public DatasetBase {         return Status::OK();       } -      std::vector<Tensor> result;-      do {-        result.clear();-        TF_RETURN_IF_ERROR(input_impl_->GetNext(ctx, &result, end_of_sequence));-        if (*end_of_sequence) {-          input_impl_.reset();-          return Status::OK();-        }-      } while ((next_index_++ % dataset()->num_shards_) != dataset()->index_);+      int num_to_skip = (dataset()->index_ - next_index_) %+                        dataset()->num_shards_;+      if (num_to_skip < 0) {+        num_to_skip += dataset()->num_shards_;+      }+      int num_skipped;+      TF_RETURN_IF_ERROR(input_impl_->Skip(ctx, num_to_skip, end_of_sequence,+                                           &num_skipped));+      next_index_ += num_skipped;

 class ShardDatasetOp::Dataset : public DatasetBase {         return Status::OK();       } -      std::vector<Tensor> result;-      do {-        result.clear();-        TF_RETURN_IF_ERROR(input_impl_->GetNext(ctx, &result, end_of_sequence));-        if (*end_of_sequence) {-          input_impl_.reset();-          return Status::OK();-        }-      } while ((next_index_++ % dataset()->num_shards_) != dataset()->index_);+      int num_to_skip = (dataset()->index_ - next_index_) %+                        dataset()->num_shards_;+      if (num_to_skip < 0) {+        num_to_skip += dataset()->num_shards_;+      }

 import itertools+import pickle+import os.path+from googleapiclient.discovery import build+from google_auth_oauthlib.flow import InstalledAppFlow+from google.auth.transport.requests import Request++# If modifying these scopes, delete the file token.pickle.+SCOPES = ['https://www.googleapis.com/auth/spreadsheets.readonly']++SPREADSHEET_ID = '1c2Kvc3y1GrF9EFqAKvPoRU8NmFV_exCrAeCTrmDct2Y'+ELOS_RANGE = 'TEST!A2:ZZ12'

 class OptimizationOptions(options.OptionsBase):       "cardinality preserved transformations, e.g. dataset.map(...).take(3)"

 class OptimizationOptions(options.OptionsBase):       "cardinality preserved transformations, e.g. dataset.map(...).take(3)"       "will be optimized to dataset.take(3).map(...). For now this"       "optimization will move `skip`, `shard` and `take` to the front of"-      "`cache`, `map` and `prefetch`. And notice this optimization is only"-      "for performance, it will not affect the output of the dataset."-      "However, it will influence the cache to the file, for the unused"-      "data will no longer be saved after this optimization."-      "If None, defaults to False.")+      "`map` and `prefetch`. And notice this optimization is only for"+      "performance, it will not affect the output of the dataset."

 class OptimizationOptions(options.OptionsBase):       "cardinality preserved transformations, e.g. dataset.map(...).take(3)"       "will be optimized to dataset.take(3).map(...). For now this"       "optimization will move `skip`, `shard` and `take` to the front of"-      "`cache`, `map` and `prefetch`. And notice this optimization is only"-      "for performance, it will not affect the output of the dataset."-      "However, it will influence the cache to the file, for the unused"-      "data will no longer be saved after this optimization."-      "If None, defaults to False.")+      "`map` and `prefetch`. And notice this optimization is only for"

 class OptimizationOptions(options.OptionsBase):       "cardinality preserved transformations, e.g. dataset.map(...).take(3)"       "will be optimized to dataset.take(3).map(...). For now this"       "optimization will move `skip`, `shard` and `take` to the front of"-      "`cache`, `map` and `prefetch`. And notice this optimization is only"-      "for performance, it will not affect the output of the dataset."-      "However, it will influence the cache to the file, for the unused"-      "data will no longer be saved after this optimization."-      "If None, defaults to False.")+      "`map` and `prefetch`. And notice this optimization is only for"+      "performance, it will not affect the output of the dataset."

+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.++Licensed under the Apache License, Version 2.0 (the "License");+you may not use this file except in compliance with the License.+You may obtain a copy of the License at++    http://www.apache.org/licenses/LICENSE-2.0++Unless required by applicable law or agreed to in writing, software+distributed under the License is distributed on an "AS IS" BASIS,+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+See the License for the specific language governing permissions and+limitations under the License.+==============================================================================*/++#include "tensorflow/core/grappler/optimizers/data/hoist_discard.h"++#include "absl/container/flat_hash_set.h"+#include "tensorflow/core/framework/attr_value.pb.h"+#include "tensorflow/core/framework/node_def.pb.h"+#include "tensorflow/core/grappler/clusters/cluster.h"+#include "tensorflow/core/grappler/grappler_item.h"+#include "tensorflow/core/grappler/mutable_graph_view.h"+#include "tensorflow/core/grappler/op_types.h"+#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h"+#include "tensorflow/core/grappler/optimizers/data/function_utils.h"+#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"+#include "tensorflow/core/grappler/utils.h"+#include "tensorflow/core/platform/protobuf.h"++namespace tensorflow {+namespace grappler {+namespace {++const std::unordered_set<string> kDataDiscarding = {+    "ShardDataset", "SkipDataset", "TakeDataset",+};++const std::unordered_set<string> kCardinalityPreserving = {+    "CacheDataset", "CacheDatasetV2", "PrefetchDataset",+    "MapDataset", "ParallelMapDataset", "ParallelMapDatasetV2",+};++bool IsDataDiscarding(const NodeDef& node) {+  auto iter = kDataDiscarding.find(node.op());+  if (iter == kDataDiscarding.end()) {+    return false;+  }+  return true;+}++bool IsCardinalityPreserving(const NodeDef& node) {+  auto iter = kCardinalityPreserving.find(node.op());+  if (iter == kCardinalityPreserving.end()) {+    return false;+  }+  auto attr_iter = node.attr().find("preserve_cardinality");

 class OptimizationOptions(options.OptionsBase):       "Whether to fuse filter dataset that predicts random_uniform < rate into "       "a sampling dataset. If None, defaults to False.") +  hoist_discard = options.create_option(+      name="hoist_discard",+      ty=bool,+      docstring=+      "Whether to hoist ops that will discard data (such as skip, take, shard) "+      "out of unary cardinality preserved transformations. "

+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.++Licensed under the Apache License, Version 2.0 (the "License");+you may not use this file except in compliance with the License.+You may obtain a copy of the License at++    http://www.apache.org/licenses/LICENSE-2.0++Unless required by applicable law or agreed to in writing, software+distributed under the License is distributed on an "AS IS" BASIS,+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+See the License for the specific language governing permissions and+limitations under the License.+==============================================================================*/++#include "tensorflow/core/grappler/optimizers/data/hoist_discard.h"++#include "absl/container/flat_hash_set.h"+#include "tensorflow/core/framework/attr_value.pb.h"+#include "tensorflow/core/framework/node_def.pb.h"+#include "tensorflow/core/grappler/clusters/cluster.h"+#include "tensorflow/core/grappler/grappler_item.h"+#include "tensorflow/core/grappler/mutable_graph_view.h"+#include "tensorflow/core/grappler/op_types.h"+#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h"+#include "tensorflow/core/grappler/optimizers/data/function_utils.h"+#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"+#include "tensorflow/core/grappler/utils.h"+#include "tensorflow/core/platform/protobuf.h"++namespace tensorflow {+namespace grappler {+namespace {++const std::unordered_set<string> kDataDiscarding = {+    "ShardDataset", "SkipDataset", "TakeDataset",+};++const std::unordered_set<string> kCardinalityPreserving = {+    "CacheDataset", "CacheDatasetV2", "PrefetchDataset",+    "MapDataset", "ParallelMapDataset", "ParallelMapDatasetV2",+};++bool IsDataDiscarding(const NodeDef& node) {+  auto iter = kDataDiscarding.find(node.op());+  if (iter == kDataDiscarding.end()) {+    return false;+  }+  return true;+}++bool IsCardinalityPreserving(const NodeDef& node) {+  auto iter = kCardinalityPreserving.find(node.op());+  if (iter == kCardinalityPreserving.end()) {+    return false;+  }+  auto attr_iter = node.attr().find("preserve_cardinality");+  if (attr_iter != node.attr().end() && !attr_iter->second.b()) {+    return false;+  }+  return true;+}++}  // namepsace++Status HoistDiscard::OptimizeAndCollectStats(Cluster* cluster,+                                             const GrapplerItem& item,+                                             GraphDef* output,+                                             OptimizationStats* stats) {+  *output = item.graph;+  MutableGraphView graph(output);+  bool updated;+  do {+    updated = false;+    for (int i = 0; i < graph.graph()->node_size(); i++) {+      auto node = graph.graph()->mutable_node(i);+      if (IsDataDiscarding(*node)) {+        NodeDef* start = node;+        NodeDef* start_parent = graph_utils::GetInputNode(*start, graph);+        while (IsCardinalityPreserving(*start_parent)) {+          start = start_parent;+          start_parent = graph_utils::GetInputNode(*start, graph);+        }+        if (start->name() == node->name()) {+          continue;+        }+        auto parent = graph_utils::GetInputNode(*node, graph);+        TF_RETURN_IF_ERROR(graph.UpdateFanouts(node->name(), parent->name()));+        if (!absl::StartsWith(node->name(), "hoist_discard/")) {

+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.++Licensed under the Apache License, Version 2.0 (the "License");+you may not use this file except in compliance with the License.+You may obtain a copy of the License at++    http://www.apache.org/licenses/LICENSE-2.0++Unless required by applicable law or agreed to in writing, software+distributed under the License is distributed on an "AS IS" BASIS,+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+See the License for the specific language governing permissions and+limitations under the License.+==============================================================================*/++#ifndef TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_HOIST_DATA_DISCARDING_OPS_H_+#define TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_HOIST_DATA_DISCARDING_OPS_H_++#include "tensorflow/core/grappler/optimizers/data/optimizer_base.h"++namespace tensorflow {+namespace grappler {++// This optimization hoists the data discarding ops (such as `skip`, `take` and+//  `shard`) to avoid unnecessary computation.+class HoistDiscard : public TFDataOptimizerBase {+ public:+  HoistDiscard() = default;+  ~HoistDiscard() override = default;++  string name() const override { return "hoist_discard"; };++  bool UsesFunctionLibrary() const override { return false; }++  Status Init(+      const tensorflow::RewriterConfig_CustomGraphOptimizer* config) override {+    return Status::OK();+  }++  Status OptimizeAndCollectStats(Cluster* cluster, const GrapplerItem& item,+                                 GraphDef* output,+                                 OptimizationStats* stats) override;++  void Feedback(Cluster* cluster, const GrapplerItem& item,+                const GraphDef& optimize_output, double result) override;+};++}  // namespace grappler+}  // namespace tensorflow++#endif  // TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_HOIST_DATA_DISCARDING_OPS_H_+

+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.+# ==============================================================================+"""Tests for the `HoistDiscard` rewrite."""+from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized++from tensorflow.python.data.experimental.ops import testing+from tensorflow.python.data.kernel_tests import test_base+from tensorflow.python.data.ops import dataset_ops+from tensorflow.python.framework import combinations+from tensorflow.python.platform import test+++class HoistDiscardTest(test_base.DatasetTestBase, parameterized.TestCase):++  @combinations.generate(combinations.combine(tf_api_version=2,+                                              mode=["eager", "graph"]))+  def testSimpleHoistingV2(self):+    dataset = dataset_ops.Dataset.range(100)+    dataset = dataset.apply(+        testing.assert_next(["FiniteSkip", "FiniteTake", "Shard",+                             "ParallelMap", "MemoryCacheImpl"]))+    dataset = dataset.map(+        lambda x: x + 1, num_parallel_calls=10)+    dataset = dataset.skip(10)+    dataset = dataset.cache()+    dataset = dataset.take(50)+    dataset = dataset.shard(2, 0)+    options = dataset_ops.Options()+    options.experimental_optimization.apply_default_optimizations = False+    options.experimental_optimization.hoist_discard = True+    dataset = dataset.with_options(options)+    self.assertDatasetProduces(dataset, range(11, 61, 2))++  @combinations.generate(combinations.combine(tf_api_version=1,+                                              mode=["eager", "graph"]))+  def testSimpleHoistingV1(self):+    dataset = dataset_ops.Dataset.range(100)+    dataset = dataset.apply(+        testing.assert_next(["ParallelMap", "FiniteSkip", "FiniteTake",

+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.++Licensed under the Apache License, Version 2.0 (the "License");+you may not use this file except in compliance with the License.+You may obtain a copy of the License at++    http://www.apache.org/licenses/LICENSE-2.0++Unless required by applicable law or agreed to in writing, software+distributed under the License is distributed on an "AS IS" BASIS,+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+See the License for the specific language governing permissions and+limitations under the License.+==============================================================================*/++#ifndef TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_HOIST_DATA_DISCARDING_OPS_H_+#define TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_HOIST_DATA_DISCARDING_OPS_H_++#include "tensorflow/core/grappler/optimizers/data/optimizer_base.h"++namespace tensorflow {+namespace grappler {++// This optimization hoists the data discarding ops (such as `skip`, `take` and+//  `shard`) to avoid unnecessary computation.

+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.+# ==============================================================================+"""Tests for the `HoistDiscard` rewrite."""+from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized++from tensorflow.python.data.experimental.ops import testing+from tensorflow.python.data.kernel_tests import test_base+from tensorflow.python.data.ops import dataset_ops+from tensorflow.python.framework import combinations+from tensorflow.python.platform import test+++class HoistDiscardTest(test_base.DatasetTestBase, parameterized.TestCase):

+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.++Licensed under the Apache License, Version 2.0 (the "License");+you may not use this file except in compliance with the License.+You may obtain a copy of the License at++    http://www.apache.org/licenses/LICENSE-2.0++Unless required by applicable law or agreed to in writing, software+distributed under the License is distributed on an "AS IS" BASIS,+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+See the License for the specific language governing permissions and+limitations under the License.+==============================================================================*/++#ifndef TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_DATA_HOIST_DATA_DISCARDING_OPS_H_