diff --git a/tensorflow/core/kernels/data/shuffle_dataset_op_test.cc b/tensorflow/core/kernels/data/shuffle_dataset_op_test.cc
index dbf8c9a5a70..38b93f13808 100644
--- a/tensorflow/core/kernels/data/shuffle_dataset_op_test.cc
+++ b/tensorflow/core/kernels/data/shuffle_dataset_op_test.cc
@@ -16,30 +16,41 @@ namespace tensorflow {
 namespace data {
 namespace {
 
-constexpr char kNodeName[] = "shuffle_dataset";
-constexpr char kOpName[] = "ShuffleDataset";
+constexpr char kShuffleNodeName[] = "shuffle_dataset";
+constexpr char kShuffleOpName[] = "ShuffleDataset";
+constexpr char kShuffleAndRepeatNodeName[] = "shuffle_and_repeat_dataset";
+constexpr char kShuffleAndRepeatOpName[] = "ShuffleAndRepeatDataset";
 
 class ShuffleDatasetOpTest : public DatasetOpsTestBase {
  protected:
-  // Creates a new `ShuffleDataset` op kernel
-  Status CreateShuffleDatasetOpKernel(
-      bool reshuffle_each_iteration, const DataTypeVector& output_types,
+  // Creates a new `ShuffleDataset`/`ShuffleAndRepeatDataset` op kernel
+  Status CreateDatasetOpKernel(
+      int64 count, bool reshuffle_each_iteration,
+      const DataTypeVector& output_types,
       const std::vector<PartialTensorShape>& output_shapes,
       std::unique_ptr<OpKernel>* shuffle_dataset_kernel) {
-    NodeDef node_def = test::function::NDef(
-        kNodeName, kOpName, {"input_dataset", "buffer_size", "seed", "seed2"},
-        {{"reshuffle_each_iteration", reshuffle_each_iteration},
-         {"output_types", output_types},
-         {"output_shapes", output_shapes}});
+    NodeDef node_def;
+    if (count == 1) {
+      node_def = test::function::NDef(
+          kShuffleNodeName, kShuffleOpName,
+          {"input_dataset", "buffer_size", "seed", "seed2"},
+          {{"reshuffle_each_iteration", reshuffle_each_iteration},
+           {"output_types", output_types},
+           {"output_shapes", output_shapes}});
+    } else {
+      node_def = test::function::NDef(
+          kShuffleAndRepeatNodeName, kShuffleAndRepeatOpName,
+          {"input_dataset", "buffer_size", "seed", "seed2", "count"},
+          {{"output_types", output_types}, {"output_shapes", output_shapes}});
+    }
     TF_RETURN_IF_ERROR(CreateOpKernel(node_def, shuffle_dataset_kernel));
     return Status::OK();
   }
 
-  // Creates a new `ShuffleDataset` op kernel context.
-  Status CreateShuffleDatasetContext(
-      OpKernel* const op_kernel,
-      gtl::InlinedVector<TensorValue, 4>* const inputs,
-      std::unique_ptr<OpKernelContext>* context) {
+  // Creates a new `ShuffleDataset`/`ShuffleAndRepeatDataset` op kernel context.
+  Status CreateDatasetContext(OpKernel* const op_kernel,
+                              gtl::InlinedVector<TensorValue, 4>* const inputs,
+                              std::unique_ptr<OpKernelContext>* context) {
     TF_RETURN_IF_ERROR(CheckOpKernelInput(*op_kernel, *inputs));
     TF_RETURN_IF_ERROR(CreateOpKernelContext(op_kernel, inputs, context));
     return Status::OK();
@@ -57,8 +68,10 @@ struct TestCase {
   Tensor buffer_size;
   Tensor seed;
   Tensor seed2;
+  Tensor count;
   bool reshuffle_each_iteration;
-  std::vector<Tensor> expected_outputs;
+  std::vector<Tensor> expected_shuffle_outputs;
+  std::vector<Tensor> expected_reshuffle_outputs;
   DataTypeVector expected_output_dtypes;
   std::vector<PartialTensorShape> expected_output_shapes;
   int64 expected_cardinality;
@@ -76,7 +89,7 @@ std::vector<Tensor> ConvertToTensorVec(std::vector<T> values) {
   return tensors;
 }
 
-// Test case 1: normal case with reshuffle_each_iteration = false
+// Test case 1: test shuffle_dataset with reshuffle_each_iteration = false.
 TestCase TestCase1() {
   return {
       /*range_data_param*/ {0, 10, 1},
@@ -84,16 +97,19 @@ TestCase TestCase1() {
       DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {3}),
       /*seed*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
       /*seed2*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {2}),
+      /*count*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
       /*reshuffle_each_iteration*/ false,
-      /*expected_outputs*/
-      ConvertToTensorVec<int64>({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}),
+      /*expected_shuffle_outputs*/
+      ConvertToTensorVec<int64>({2, 3, 0, 5, 6, 4, 7, 8, 9, 1}),
+      /*expected_reshuffle_outputs*/
+      ConvertToTensorVec<int64>({2, 3, 0, 5, 6, 4, 7, 8, 9, 1}),
       /*expected_output_dtypes*/ {DT_INT64},
       /*expected_output_shapes*/ {PartialTensorShape({})},
       /*expected_cardinality*/ 10,
       /*breakpoints*/ {0, 1, 9}};
 }
 
-// Test case 2: normal case with reshuffle_each_iteration = true
+// Test case 2: test shuffle_dataset with reshuffle_each_iteration = true.
 TestCase TestCase2() {
   return {
       /*range_data_param*/ {0, 10, 1},
@@ -101,26 +117,73 @@ TestCase TestCase2() {
       DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {10}),
       /*seed*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
       /*seed2*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {2}),
+      /*count*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
       /*reshuffle_each_iteration*/ true,
-      /*expected_outputs*/
-      ConvertToTensorVec<int64>({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}),
+      /*expected_shuffle_outputs*/
+      ConvertToTensorVec<int64>({2, 6, 1, 3, 9, 5, 0, 8, 7, 4}),
+      /*expected_reshuffle_outputs*/
+      ConvertToTensorVec<int64>({1, 6, 0, 5, 2, 7, 4, 3, 9, 8}),
       /*expected_output_dtypes*/ {DT_INT64},
       /*expected_output_shapes*/ {PartialTensorShape({})},
       /*expected_cardinality*/ 10,
       /*breakpoints*/ {0, 1, 9}};
 }
 
-// Test case 3: special case with buffer_size = 1 &
-// reshuffle_each_iteration = true
+// Test case 3: similar with the test case 2 but a smaller buffer size than
+// the input dataset.
 TestCase TestCase3() {
+  return {
+      /*range_data_param*/ {0, 10, 1},
+      /*buffer_size*/
+      DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {2}),
+      /*seed*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
+      /*seed2*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {2}),
+      /*count*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
+      /*reshuffle_each_iteration*/ true,
+      /*expected_shuffle_outputs*/
+      ConvertToTensorVec<int64>({0, 2, 1, 3, 5, 6, 4, 7, 8, 9}),
+      /*expected_reshuffle_outputs*/
+      ConvertToTensorVec<int64>({1, 0, 2, 3, 4, 5, 6, 7, 9, 8}),
+      /*expected_output_dtypes*/ {DT_INT64},
+      /*expected_output_shapes*/ {PartialTensorShape({})},
+      /*expected_cardinality*/ 10,
+      /*breakpoints*/ {0, 1, 9}};
+}
+
+// Test case 4: similar with the test case 2 but has different seeds.
+TestCase TestCase4() {
+  return {
+      /*range_data_param*/ {0, 10, 1},
+      /*buffer_size*/
+      DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {10}),
+      /*seed*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {2}),
+      /*seed2*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {2}),
+      /*count*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
+      /*reshuffle_each_iteration*/ true,
+      /*expected_shuffle_outputs*/
+      ConvertToTensorVec<int64>({3, 0, 8, 1, 5, 4, 7, 2, 6, 9}),
+      /*expected_reshuffle_outputs*/
+      ConvertToTensorVec<int64>({4, 6, 9, 0, 1, 8, 2, 7, 3, 5}),
+      /*expected_output_dtypes*/ {DT_INT64},
+      /*expected_output_shapes*/ {PartialTensorShape({})},
+      /*expected_cardinality*/ 10,
+      /*breakpoints*/ {0, 1, 9}};
+}
+
+// Test case 5: test shuffle_dataset with buffer_size = 1 &
+// reshuffle_each_iteration = true.
+TestCase TestCase5() {
   return {
       /*range_data_param*/ {0, 10, 1},
       /*buffer_size*/
       DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
       /*seed*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
       /*seed2*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {2}),
+      /*count*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
       /*reshuffle_each_iteration*/ true,
-      /*expected_outputs*/
+      /*expected_shuffle_outputs*/
+      ConvertToTensorVec<int64>({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}),
+      /*expected_reshuffle_outputs*/
       ConvertToTensorVec<int64>({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}),
       /*expected_output_dtypes*/ {DT_INT64},
       /*expected_output_shapes*/ {PartialTensorShape({})},
@@ -128,21 +191,125 @@ TestCase TestCase3() {
       /*breakpoints*/ {0, 1, 9}};
 }
 
-TestCase InvalidBufferSizeTestCase() {
+// Test case 6: test shuffle_dataset with an empty input dataset.
+TestCase TestCase6() {
+  return {
+      /*range_data_param*/ {0, 0, 1},
+      /*buffer_size*/
+      DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {10}),
+      /*seed*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
+      /*seed2*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {2}),
+      /*count*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
+      /*reshuffle_each_iteration*/ true,
+      /*expected_shuffle_outputs*/
+      ConvertToTensorVec<int64>({}),
+      /*expected_reshuffle_outputs*/
+      ConvertToTensorVec<int64>({}),
+      /*expected_output_dtypes*/ {DT_INT64},
+      /*expected_output_shapes*/ {PartialTensorShape({})},
+      /*expected_cardinality*/ 0,
+      /*breakpoints*/ {0, 1, 9}};
+}
+
+// Test case 7: test shuffle_and_repeat_dataset with buffer_size = 10 &
+// count = 2.
+TestCase TestCase7() {
+  return {
+      /*range_data_param*/ {0, 10, 1},
+      /*buffer_size*/
+      DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {10}),
+      /*seed*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
+      /*seed2*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {2}),
+      /*count*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {2}),
+      /*reshuffle_each_iteration*/ false,
+      /*expected_shuffle_outputs*/
+      ConvertToTensorVec<int64>(
+          {9, 0, 8, 6, 1, 3, 7, 2, 4, 5, 4, 3, 0, 5, 8, 2, 6, 9, 7, 1}),
+      /*expected_reshuffle_outputs*/
+      ConvertToTensorVec<int64>(
+          {9, 0, 8, 6, 1, 3, 7, 2, 4, 5, 4, 3, 0, 5, 8, 2, 6, 9, 7, 1}),
+      /*expected_output_dtypes*/ {DT_INT64},
+      /*expected_output_shapes*/ {PartialTensorShape({})},
+      /*expected_cardinality*/ 20,
+      /*breakpoints*/ {0, 5, 22}};
+}
+
+// Test case 8: test shuffle_and_repeat_dataset with buffer_size = 10 &
+// count = -1
+TestCase TestCase8() {
+  return {
+      /*range_data_param*/ {0, 3, 1},
+      /*buffer_size*/
+      DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {10}),
+      /*seed*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
+      /*seed2*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {2}),
+      /*count*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {-1}),
+      /*reshuffle_each_iteration*/ false,
+      /*expected_shuffle_outputs*/
+      ConvertToTensorVec<int64>(
+          {2, 0, 1, 2, 0, 1, 1, 2, 0, 1, 0, 2, 2, 0, 1, 1, 0, 2, 2, 1, 0}),
+      /*expected_reshuffle_outputs*/
+      ConvertToTensorVec<int64>(
+          {2, 0, 1, 2, 0, 1, 1, 2, 0, 1, 0, 2, 2, 0, 1, 1, 0, 2, 2, 1, 0}),
+      /*expected_output_dtypes*/ {DT_INT64},
+      /*expected_output_shapes*/ {PartialTensorShape({})},
+      /*expected_cardinality*/ kInfiniteCardinality,
+      /*breakpoints*/ {0, 5, 20}};
+}
+
+TestCase InvalidBufferSizeTestCaseForShuffleDataset() {
   return {
       /*range_data_param*/ {0, 10, 1},
       /*buffer_size*/
       DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {-1}),
       /*seed*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
       /*seed2*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {2}),
+      /*count*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
       /*reshuffle_each_iteration*/ true,
-      /*expected_outputs*/ ConvertToTensorVec<int64>({}),
+      /*expected_shuffle_outputs*/ ConvertToTensorVec<int64>({}),
+      /*expected_reshuffle_outputs*/ ConvertToTensorVec<int64>({}),
       /*expected_output_dtypes*/ {DT_INT64},
       /*expected_output_shapes*/ {PartialTensorShape({})},
-      /*expected_cardinality*/ 10,
+      /*expected_cardinality*/ 0,
       /*breakpoints*/ {0, 1, 9}};
 }
 
+TestCase InvalidBufferSizeTestCaseForShuffleAndRepeatDataset() {
+  return {
+      /*range_data_param*/ {0, 10, 1},
+      /*buffer_size*/
+      DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {-1}),
+      /*seed*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
+      /*seed2*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {2}),
+      /*count*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {2}),
+      /*reshuffle_each_iteration*/ true,
+      /*expected_shuffle_outputs*/ ConvertToTensorVec<int64>({}),
+      /*expected_reshuffle_outputs*/ ConvertToTensorVec<int64>({}),
+      /*expected_output_dtypes*/ {DT_INT64},
+      /*expected_output_shapes*/ {PartialTensorShape({})},
+      /*expected_cardinality*/ 0,
+      /*breakpoints*/ {0, 1, 9}};
+}
+
+TestCase InvalidCountTestCaseForShuffleAndRepeatDataset() {
+  return {
+      /*range_data_param*/ {0, 3, 1},
+      /*buffer_size*/
+      DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {10}),
+      /*seed*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {1}),
+      /*seed2*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {2}),
+      /*count*/ DatasetOpsTestBase::CreateTensor<int64>(TensorShape({}), {0}),
+      /*reshuffle_each_iteration*/ false,
+      /*expected_shuffle_outputs*/
+      ConvertToTensorVec<int64>({}),
+      /*expected_reshuffle_outputs*/
+      ConvertToTensorVec<int64>({}),
+      /*expected_output_dtypes*/ {DT_INT64},
+      /*expected_output_shapes*/ {PartialTensorShape({})},
+      /*expected_cardinality*/ 0,
+      /*breakpoints*/ {0, 5, 20}};
+}
+
 class ParameterizedShuffleDatasetOpTest
     : public ShuffleDatasetOpTest,
       public ::testing::WithParamInterface<TestCase> {};
@@ -153,10 +320,13 @@ TEST_P(ParameterizedShuffleDatasetOpTest, GetNext) {
   TF_ASSERT_OK(InitThreadPool(thread_num));
   TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));
 
-  std::unique_ptr<OpKernel> shuffle_dataset_kernel;
-  TF_ASSERT_OK(CreateShuffleDatasetOpKernel(
-      test_case.reshuffle_each_iteration, test_case.expected_output_dtypes,
-      test_case.expected_output_shapes, &shuffle_dataset_kernel));
+  Tensor count = test_case.count;
+  int64 count_value = count.flat<int64>()(0);
+  std::unique_ptr<OpKernel> dataset_kernel;
+  TF_ASSERT_OK(
+      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
+                            test_case.expected_output_dtypes,
+                            test_case.expected_output_shapes, &dataset_kernel));
 
   DatasetBase* range_dataset;
   TF_ASSERT_OK(CreateRangeDataset<int64>(
@@ -170,49 +340,78 @@ TEST_P(ParameterizedShuffleDatasetOpTest, GetNext) {
   Tensor seed2 = test_case.seed2;
   gtl::InlinedVector<TensorValue, 4> inputs(
       {&range_dataset_tensor, &buffer_size, &seed, &seed2});
+  if (count_value != 1) inputs.push_back(&count);
 
-  std::unique_ptr<OpKernelContext> shuffle_dataset_context;
-  TF_ASSERT_OK(CreateShuffleDatasetContext(shuffle_dataset_kernel.get(),
-                                           &inputs, &shuffle_dataset_context));
-  DatasetBase* shuffle_dataset;
-  TF_ASSERT_OK(CreateDataset(shuffle_dataset_kernel.get(),
-                             shuffle_dataset_context.get(), &shuffle_dataset));
-  core::ScopedUnref scoped_unref_shuffle_dataset(shuffle_dataset);
+  std::unique_ptr<OpKernelContext> dataset_context;
+  TF_ASSERT_OK(
+      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
+  DatasetBase* dataset;
+  TF_ASSERT_OK(
+      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
+  core::ScopedUnref scoped_unref_dataset(dataset);
 
   std::unique_ptr<IteratorContext> iterator_ctx;
-  TF_ASSERT_OK(
-      CreateIteratorContext(shuffle_dataset_context.get(), &iterator_ctx));
+  TF_ASSERT_OK(CreateIteratorContext(dataset_context.get(), &iterator_ctx));
   std::unique_ptr<IteratorBase> iterator;
   TF_ASSERT_OK(
-      shuffle_dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));
+      dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));
 
   bool end_of_sequence = false;
-  std::vector<Tensor> out_tensors;
+  std::vector<Tensor> shuffled_out_tensors;
   while (!end_of_sequence) {
     std::vector<Tensor> next;
     TF_EXPECT_OK(
         iterator->GetNext(iterator_ctx.get(), &next, &end_of_sequence));
-    out_tensors.insert(out_tensors.end(), next.begin(), next.end());
+    shuffled_out_tensors.insert(shuffled_out_tensors.end(), next.begin(),
+                                next.end());
+    // For the forever-repeat case, we test only a finite number of steps of
+    // the infinite sequence.
+    if (count_value == -1 && shuffled_out_tensors.size() ==
+                                 test_case.expected_shuffle_outputs.size()) {
+      break;
+    }
   }
 
-  // When `buffer_size = 1`, the output sequence of `ShuffleDataset` will be in
-  // order, so we need to consider the element sequence when evaluating the
-  // result for this case.
-  bool expect_items_equal = test_case.buffer_size.flat<int64>()(0) > 1;
-  TF_EXPECT_OK(ExpectEqual(out_tensors, test_case.expected_outputs,
-                           /*expect_items_equal*/ expect_items_equal));
+  // Reshuffle the dataset.
+  end_of_sequence = false;
+  TF_ASSERT_OK(
+      dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));
+  std::vector<Tensor> reshuffled_out_tensors;
+  while (!end_of_sequence) {
+    std::vector<Tensor> next;
+    TF_EXPECT_OK(
+        iterator->GetNext(iterator_ctx.get(), &next, &end_of_sequence));
+    reshuffled_out_tensors.insert(reshuffled_out_tensors.end(), next.begin(),
+                                  next.end());
+    // For the forever-repeat case, we test only a finite number of steps of
+    // the infinite sequence.
+    if (count_value == -1 && reshuffled_out_tensors.size() ==
+                                 test_case.expected_shuffle_outputs.size()) {
+      break;
+    }
+  }
+
+  TF_EXPECT_OK(ExpectEqual(shuffled_out_tensors,
+                           test_case.expected_shuffle_outputs,
+                           /*compare_order*/ true));
+  TF_EXPECT_OK(ExpectEqual(reshuffled_out_tensors,
+                           test_case.expected_reshuffle_outputs,
+                           /*compare_order*/ true));
 }
 
-TEST_F(ShuffleDatasetOpTest, DatasetNodeName) {
+TEST_P(ParameterizedShuffleDatasetOpTest, DatasetNodeName) {
   int thread_num = 2, cpu_num = 2;
-  TestCase test_case = TestCase1();
+  TestCase test_case = GetParam();
   TF_ASSERT_OK(InitThreadPool(thread_num));
   TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));
 
-  std::unique_ptr<OpKernel> shuffle_dataset_kernel;
-  TF_ASSERT_OK(CreateShuffleDatasetOpKernel(
-      test_case.reshuffle_each_iteration, test_case.expected_output_dtypes,
-      test_case.expected_output_shapes, &shuffle_dataset_kernel));
+  Tensor count = test_case.count;
+  int64 count_value = count.flat<int64>()(0);
+  std::unique_ptr<OpKernel> dataset_kernel;
+  TF_ASSERT_OK(
+      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
+                            test_case.expected_output_dtypes,
+                            test_case.expected_output_shapes, &dataset_kernel));
 
   DatasetBase* range_dataset;
   TF_ASSERT_OK(CreateRangeDataset<int64>(
@@ -226,28 +425,36 @@ TEST_F(ShuffleDatasetOpTest, DatasetNodeName) {
   Tensor seed2 = test_case.seed2;
   gtl::InlinedVector<TensorValue, 4> inputs(
       {&range_dataset_tensor, &buffer_size, &seed, &seed2});
+  if (count_value != 1) inputs.push_back(&count);
 
-  std::unique_ptr<OpKernelContext> shuffle_dataset_context;
-  TF_ASSERT_OK(CreateShuffleDatasetContext(shuffle_dataset_kernel.get(),
-                                           &inputs, &shuffle_dataset_context));
-  DatasetBase* shuffle_dataset;
-  TF_ASSERT_OK(CreateDataset(shuffle_dataset_kernel.get(),
-                             shuffle_dataset_context.get(), &shuffle_dataset));
-  core::ScopedUnref scoped_unref_shuffle_dataset(shuffle_dataset);
+  std::unique_ptr<OpKernelContext> dataset_context;
+  TF_ASSERT_OK(
+      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
+  DatasetBase* dataset;
+  TF_ASSERT_OK(
+      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
+  core::ScopedUnref scoped_unref_dataset(dataset);
 
-  EXPECT_EQ(shuffle_dataset->node_name(), kNodeName);
+  if (count_value == 1) {
+    EXPECT_EQ(dataset->node_name(), kShuffleNodeName);
+  } else {
+    EXPECT_EQ(dataset->node_name(), kShuffleAndRepeatNodeName);
+  }
 }
 
-TEST_F(ShuffleDatasetOpTest, DatasetTypeString) {
+TEST_P(ParameterizedShuffleDatasetOpTest, DatasetTypeString) {
   int thread_num = 2, cpu_num = 2;
-  TestCase test_case = TestCase1();
+  TestCase test_case = GetParam();
   TF_ASSERT_OK(InitThreadPool(thread_num));
   TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));
 
-  std::unique_ptr<OpKernel> shuffle_dataset_kernel;
-  TF_ASSERT_OK(CreateShuffleDatasetOpKernel(
-      test_case.reshuffle_each_iteration, test_case.expected_output_dtypes,
-      test_case.expected_output_shapes, &shuffle_dataset_kernel));
+  Tensor count = test_case.count;
+  int64 count_value = count.flat<int64>()(0);
+  std::unique_ptr<OpKernel> dataset_kernel;
+  TF_ASSERT_OK(
+      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
+                            test_case.expected_output_dtypes,
+                            test_case.expected_output_shapes, &dataset_kernel));
 
   DatasetBase* range_dataset;
   TF_ASSERT_OK(CreateRangeDataset<int64>(
@@ -261,16 +468,21 @@ TEST_F(ShuffleDatasetOpTest, DatasetTypeString) {
   Tensor seed2 = test_case.seed2;
   gtl::InlinedVector<TensorValue, 4> inputs(
       {&range_dataset_tensor, &buffer_size, &seed, &seed2});
+  if (count_value != 1) inputs.push_back(&count);
 
-  std::unique_ptr<OpKernelContext> shuffle_dataset_context;
-  TF_ASSERT_OK(CreateShuffleDatasetContext(shuffle_dataset_kernel.get(),
-                                           &inputs, &shuffle_dataset_context));
-  DatasetBase* shuffle_dataset;
-  TF_ASSERT_OK(CreateDataset(shuffle_dataset_kernel.get(),
-                             shuffle_dataset_context.get(), &shuffle_dataset));
-  core::ScopedUnref scoped_unref_shuffle_dataset(shuffle_dataset);
+  std::unique_ptr<OpKernelContext> dataset_context;
+  TF_ASSERT_OK(
+      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
+  DatasetBase* dataset;
+  TF_ASSERT_OK(
+      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
+  core::ScopedUnref scoped_unref_dataset(dataset);
 
-  EXPECT_EQ(shuffle_dataset->type_string(), kOpName);
+  if (count_value == 1) {
+    EXPECT_EQ(dataset->type_string(), kShuffleOpName);
+  } else {
+    EXPECT_EQ(dataset->type_string(), kShuffleAndRepeatOpName);
+  }
 }
 
 TEST_P(ParameterizedShuffleDatasetOpTest, DatasetOutputDtypes) {
@@ -279,10 +491,13 @@ TEST_P(ParameterizedShuffleDatasetOpTest, DatasetOutputDtypes) {
   TF_ASSERT_OK(InitThreadPool(thread_num));
   TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));
 
-  std::unique_ptr<OpKernel> shuffle_dataset_kernel;
-  TF_ASSERT_OK(CreateShuffleDatasetOpKernel(
-      test_case.reshuffle_each_iteration, test_case.expected_output_dtypes,
-      test_case.expected_output_shapes, &shuffle_dataset_kernel));
+  Tensor count = test_case.count;
+  int64 count_value = count.flat<int64>()(0);
+  std::unique_ptr<OpKernel> dataset_kernel;
+  TF_ASSERT_OK(
+      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
+                            test_case.expected_output_dtypes,
+                            test_case.expected_output_shapes, &dataset_kernel));
 
   DatasetBase* range_dataset;
   TF_ASSERT_OK(CreateRangeDataset<int64>(
@@ -296,16 +511,17 @@ TEST_P(ParameterizedShuffleDatasetOpTest, DatasetOutputDtypes) {
   Tensor seed2 = test_case.seed2;
   gtl::InlinedVector<TensorValue, 4> inputs(
       {&range_dataset_tensor, &buffer_size, &seed, &seed2});
+  if (count_value != 1) inputs.push_back(&count);
 
-  std::unique_ptr<OpKernelContext> shuffle_dataset_context;
-  TF_ASSERT_OK(CreateShuffleDatasetContext(shuffle_dataset_kernel.get(),
-                                           &inputs, &shuffle_dataset_context));
-  DatasetBase* shuffle_dataset;
-  TF_ASSERT_OK(CreateDataset(shuffle_dataset_kernel.get(),
-                             shuffle_dataset_context.get(), &shuffle_dataset));
-  core::ScopedUnref scoped_unref_shuffle_dataset(shuffle_dataset);
+  std::unique_ptr<OpKernelContext> dataset_context;
+  TF_ASSERT_OK(
+      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
+  DatasetBase* dataset;
+  TF_ASSERT_OK(
+      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
+  core::ScopedUnref scoped_unref_dataset(dataset);
 
-  TF_EXPECT_OK(VerifyTypesMatch(shuffle_dataset->output_dtypes(),
+  TF_EXPECT_OK(VerifyTypesMatch(dataset->output_dtypes(),
                                 test_case.expected_output_dtypes));
 }
 
@@ -315,10 +531,13 @@ TEST_P(ParameterizedShuffleDatasetOpTest, DatasetOutputShapes) {
   TF_ASSERT_OK(InitThreadPool(thread_num));
   TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));
 
-  std::unique_ptr<OpKernel> shuffle_dataset_kernel;
-  TF_ASSERT_OK(CreateShuffleDatasetOpKernel(
-      test_case.reshuffle_each_iteration, test_case.expected_output_dtypes,
-      test_case.expected_output_shapes, &shuffle_dataset_kernel));
+  Tensor count = test_case.count;
+  int64 count_value = count.flat<int64>()(0);
+  std::unique_ptr<OpKernel> dataset_kernel;
+  TF_ASSERT_OK(
+      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
+                            test_case.expected_output_dtypes,
+                            test_case.expected_output_shapes, &dataset_kernel));
 
   DatasetBase* range_dataset;
   TF_ASSERT_OK(CreateRangeDataset<int64>(
@@ -332,16 +551,17 @@ TEST_P(ParameterizedShuffleDatasetOpTest, DatasetOutputShapes) {
   Tensor seed2 = test_case.seed2;
   gtl::InlinedVector<TensorValue, 4> inputs(
       {&range_dataset_tensor, &buffer_size, &seed, &seed2});
+  if (count_value != 1) inputs.push_back(&count);
 
-  std::unique_ptr<OpKernelContext> shuffle_dataset_context;
-  TF_ASSERT_OK(CreateShuffleDatasetContext(shuffle_dataset_kernel.get(),
-                                           &inputs, &shuffle_dataset_context));
-  DatasetBase* shuffle_dataset;
-  TF_ASSERT_OK(CreateDataset(shuffle_dataset_kernel.get(),
-                             shuffle_dataset_context.get(), &shuffle_dataset));
-  core::ScopedUnref scoped_unref_shuffle_dataset(shuffle_dataset);
+  std::unique_ptr<OpKernelContext> dataset_context;
+  TF_ASSERT_OK(
+      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
+  DatasetBase* dataset;
+  TF_ASSERT_OK(
+      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
+  core::ScopedUnref scoped_unref_dataset(dataset);
 
-  TF_EXPECT_OK(VerifyShapesCompatible(shuffle_dataset->output_shapes(),
+  TF_EXPECT_OK(VerifyShapesCompatible(dataset->output_shapes(),
                                       test_case.expected_output_shapes));
 }
 
@@ -351,10 +571,13 @@ TEST_P(ParameterizedShuffleDatasetOpTest, Cardinality) {
   TF_ASSERT_OK(InitThreadPool(thread_num));
   TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));
 
-  std::unique_ptr<OpKernel> shuffle_dataset_kernel;
-  TF_ASSERT_OK(CreateShuffleDatasetOpKernel(
-      test_case.reshuffle_each_iteration, test_case.expected_output_dtypes,
-      test_case.expected_output_shapes, &shuffle_dataset_kernel));
+  Tensor count = test_case.count;
+  int64 count_value = count.flat<int64>()(0);
+  std::unique_ptr<OpKernel> dataset_kernel;
+  TF_ASSERT_OK(
+      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
+                            test_case.expected_output_dtypes,
+                            test_case.expected_output_shapes, &dataset_kernel));
 
   DatasetBase* range_dataset;
   TF_ASSERT_OK(CreateRangeDataset<int64>(
@@ -368,17 +591,17 @@ TEST_P(ParameterizedShuffleDatasetOpTest, Cardinality) {
   Tensor seed2 = test_case.seed2;
   gtl::InlinedVector<TensorValue, 4> inputs(
       {&range_dataset_tensor, &buffer_size, &seed, &seed2});
+  if (count_value != 1) inputs.push_back(&count);
 
-  std::unique_ptr<OpKernelContext> shuffle_dataset_context;
-  TF_ASSERT_OK(CreateShuffleDatasetContext(shuffle_dataset_kernel.get(),
-                                           &inputs, &shuffle_dataset_context));
-  DatasetBase* shuffle_dataset;
-  TF_ASSERT_OK(CreateDataset(shuffle_dataset_kernel.get(),
-                             shuffle_dataset_context.get(), &shuffle_dataset));
-  core::ScopedUnref scoped_unref_shuffle_dataset(shuffle_dataset);
+  std::unique_ptr<OpKernelContext> dataset_context;
+  TF_ASSERT_OK(
+      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
+  DatasetBase* dataset;
+  TF_ASSERT_OK(
+      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
+  core::ScopedUnref scoped_unref_dataset(dataset);
 
-  TF_EXPECT_OK(VerifyShapesCompatible(shuffle_dataset->output_shapes(),
-                                      test_case.expected_output_shapes));
+  EXPECT_EQ(dataset->Cardinality(), test_case.expected_cardinality);
 }
 
 TEST_P(ParameterizedShuffleDatasetOpTest, DatasetSave) {
@@ -387,10 +610,13 @@ TEST_P(ParameterizedShuffleDatasetOpTest, DatasetSave) {
   TF_ASSERT_OK(InitThreadPool(thread_num));
   TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));
 
-  std::unique_ptr<OpKernel> shuffle_dataset_kernel;
-  TF_ASSERT_OK(CreateShuffleDatasetOpKernel(
-      test_case.reshuffle_each_iteration, test_case.expected_output_dtypes,
-      test_case.expected_output_shapes, &shuffle_dataset_kernel));
+  Tensor count = test_case.count;
+  int64 count_value = count.flat<int64>()(0);
+  std::unique_ptr<OpKernel> dataset_kernel;
+  TF_ASSERT_OK(
+      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
+                            test_case.expected_output_dtypes,
+                            test_case.expected_output_shapes, &dataset_kernel));
 
   DatasetBase* range_dataset;
   TF_ASSERT_OK(CreateRangeDataset<int64>(
@@ -404,20 +630,21 @@ TEST_P(ParameterizedShuffleDatasetOpTest, DatasetSave) {
   Tensor seed2 = test_case.seed2;
   gtl::InlinedVector<TensorValue, 4> inputs(
       {&range_dataset_tensor, &buffer_size, &seed, &seed2});
+  if (count_value != 1) inputs.push_back(&count);
 
-  std::unique_ptr<OpKernelContext> shuffle_dataset_context;
-  TF_ASSERT_OK(CreateShuffleDatasetContext(shuffle_dataset_kernel.get(),
-                                           &inputs, &shuffle_dataset_context));
-  DatasetBase* shuffle_dataset;
-  TF_ASSERT_OK(CreateDataset(shuffle_dataset_kernel.get(),
-                             shuffle_dataset_context.get(), &shuffle_dataset));
-  core::ScopedUnref scoped_unref_shuffle_dataset(shuffle_dataset);
+  std::unique_ptr<OpKernelContext> dataset_context;
+  TF_ASSERT_OK(
+      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
+  DatasetBase* dataset;
+  TF_ASSERT_OK(
+      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
+  core::ScopedUnref scoped_unref_dataset(dataset);
 
   std::unique_ptr<SerializationContext> serialization_context;
   TF_ASSERT_OK(CreateSerializationContext(&serialization_context));
   VariantTensorData data;
   VariantTensorDataWriter writer(&data);
-  TF_ASSERT_OK(shuffle_dataset->Save(serialization_context.get(), &writer));
+  TF_ASSERT_OK(dataset->Save(serialization_context.get(), &writer));
   TF_ASSERT_OK(writer.Flush());
 }
 
@@ -427,10 +654,13 @@ TEST_P(ParameterizedShuffleDatasetOpTest, IteratorOutputDtypes) {
   TF_ASSERT_OK(InitThreadPool(thread_num));
   TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));
 
-  std::unique_ptr<OpKernel> shuffle_dataset_kernel;
-  TF_ASSERT_OK(CreateShuffleDatasetOpKernel(
-      test_case.reshuffle_each_iteration, test_case.expected_output_dtypes,
-      test_case.expected_output_shapes, &shuffle_dataset_kernel));
+  Tensor count = test_case.count;
+  int64 count_value = count.flat<int64>()(0);
+  std::unique_ptr<OpKernel> dataset_kernel;
+  TF_ASSERT_OK(
+      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
+                            test_case.expected_output_dtypes,
+                            test_case.expected_output_shapes, &dataset_kernel));
 
   DatasetBase* range_dataset;
   TF_ASSERT_OK(CreateRangeDataset<int64>(
@@ -444,21 +674,21 @@ TEST_P(ParameterizedShuffleDatasetOpTest, IteratorOutputDtypes) {
   Tensor seed2 = test_case.seed2;
   gtl::InlinedVector<TensorValue, 4> inputs(
       {&range_dataset_tensor, &buffer_size, &seed, &seed2});
+  if (count_value != 1) inputs.push_back(&count);
 
-  std::unique_ptr<OpKernelContext> shuffle_dataset_context;
-  TF_ASSERT_OK(CreateShuffleDatasetContext(shuffle_dataset_kernel.get(),
-                                           &inputs, &shuffle_dataset_context));
-  DatasetBase* shuffle_dataset;
-  TF_ASSERT_OK(CreateDataset(shuffle_dataset_kernel.get(),
-                             shuffle_dataset_context.get(), &shuffle_dataset));
-  core::ScopedUnref scoped_unref_shuffle_dataset(shuffle_dataset);
+  std::unique_ptr<OpKernelContext> dataset_context;
+  TF_ASSERT_OK(
+      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
+  DatasetBase* dataset;
+  TF_ASSERT_OK(
+      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
+  core::ScopedUnref scoped_unref_dataset(dataset);
 
   std::unique_ptr<IteratorContext> iterator_ctx;
-  TF_ASSERT_OK(
-      CreateIteratorContext(shuffle_dataset_context.get(), &iterator_ctx));
+  TF_ASSERT_OK(CreateIteratorContext(dataset_context.get(), &iterator_ctx));
   std::unique_ptr<IteratorBase> iterator;
   TF_ASSERT_OK(
-      shuffle_dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));
+      dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));
 
   TF_EXPECT_OK(VerifyTypesMatch(iterator->output_dtypes(),
                                 test_case.expected_output_dtypes));
@@ -470,10 +700,13 @@ TEST_P(ParameterizedShuffleDatasetOpTest, IteratorOutputShapes) {
   TF_ASSERT_OK(InitThreadPool(thread_num));
   TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));
 
-  std::unique_ptr<OpKernel> shuffle_dataset_kernel;
-  TF_ASSERT_OK(CreateShuffleDatasetOpKernel(
-      test_case.reshuffle_each_iteration, test_case.expected_output_dtypes,
-      test_case.expected_output_shapes, &shuffle_dataset_kernel));
+  Tensor count = test_case.count;
+  int64 count_value = count.flat<int64>()(0);
+  std::unique_ptr<OpKernel> dataset_kernel;
+  TF_ASSERT_OK(
+      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
+                            test_case.expected_output_dtypes,
+                            test_case.expected_output_shapes, &dataset_kernel));
 
   DatasetBase* range_dataset;
   TF_ASSERT_OK(CreateRangeDataset<int64>(
@@ -487,36 +720,39 @@ TEST_P(ParameterizedShuffleDatasetOpTest, IteratorOutputShapes) {
   Tensor seed2 = test_case.seed2;
   gtl::InlinedVector<TensorValue, 4> inputs(
       {&range_dataset_tensor, &buffer_size, &seed, &seed2});
+  if (count_value != 1) inputs.push_back(&count);
 
-  std::unique_ptr<OpKernelContext> shuffle_dataset_context;
-  TF_ASSERT_OK(CreateShuffleDatasetContext(shuffle_dataset_kernel.get(),
-                                           &inputs, &shuffle_dataset_context));
-  DatasetBase* shuffle_dataset;
-  TF_ASSERT_OK(CreateDataset(shuffle_dataset_kernel.get(),
-                             shuffle_dataset_context.get(), &shuffle_dataset));
-  core::ScopedUnref scoped_unref_shuffle_dataset(shuffle_dataset);
+  std::unique_ptr<OpKernelContext> dataset_context;
+  TF_ASSERT_OK(
+      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
+  DatasetBase* dataset;
+  TF_ASSERT_OK(
+      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
+  core::ScopedUnref scoped_unref_dataset(dataset);
 
   std::unique_ptr<IteratorContext> iterator_ctx;
-  TF_ASSERT_OK(
-      CreateIteratorContext(shuffle_dataset_context.get(), &iterator_ctx));
+  TF_ASSERT_OK(CreateIteratorContext(dataset_context.get(), &iterator_ctx));
   std::unique_ptr<IteratorBase> iterator;
   TF_ASSERT_OK(
-      shuffle_dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));
+      dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));
 
   TF_EXPECT_OK(VerifyShapesCompatible(iterator->output_shapes(),
                                       test_case.expected_output_shapes));
 }
 
-TEST_F(ShuffleDatasetOpTest, IteratorOutputPrefix) {
+TEST_P(ParameterizedShuffleDatasetOpTest, IteratorOutputPrefix) {
   int thread_num = 2, cpu_num = 2;
-  TestCase test_case = TestCase1();
+  TestCase test_case = GetParam();
   TF_ASSERT_OK(InitThreadPool(thread_num));
   TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));
 
-  std::unique_ptr<OpKernel> shuffle_dataset_kernel;
-  TF_ASSERT_OK(CreateShuffleDatasetOpKernel(
-      test_case.reshuffle_each_iteration, test_case.expected_output_dtypes,
-      test_case.expected_output_shapes, &shuffle_dataset_kernel));
+  Tensor count = test_case.count;
+  int64 count_value = count.flat<int64>()(0);
+  std::unique_ptr<OpKernel> dataset_kernel;
+  TF_ASSERT_OK(
+      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
+                            test_case.expected_output_dtypes,
+                            test_case.expected_output_shapes, &dataset_kernel));
 
   DatasetBase* range_dataset;
   TF_ASSERT_OK(CreateRangeDataset<int64>(
@@ -530,23 +766,27 @@ TEST_F(ShuffleDatasetOpTest, IteratorOutputPrefix) {
   Tensor seed2 = test_case.seed2;
   gtl::InlinedVector<TensorValue, 4> inputs(
       {&range_dataset_tensor, &buffer_size, &seed, &seed2});
+  if (count_value != 1) inputs.push_back(&count);
 
-  std::unique_ptr<OpKernelContext> shuffle_dataset_context;
-  TF_ASSERT_OK(CreateShuffleDatasetContext(shuffle_dataset_kernel.get(),
-                                           &inputs, &shuffle_dataset_context));
-  DatasetBase* shuffle_dataset;
-  TF_ASSERT_OK(CreateDataset(shuffle_dataset_kernel.get(),
-                             shuffle_dataset_context.get(), &shuffle_dataset));
-  core::ScopedUnref scoped_unref_shuffle_dataset(shuffle_dataset);
+  std::unique_ptr<OpKernelContext> dataset_context;
+  TF_ASSERT_OK(
+      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
+  DatasetBase* dataset;
+  TF_ASSERT_OK(
+      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
+  core::ScopedUnref scoped_unref_dataset(dataset);
 
   std::unique_ptr<IteratorContext> iterator_ctx;
-  TF_ASSERT_OK(
-      CreateIteratorContext(shuffle_dataset_context.get(), &iterator_ctx));
+  TF_ASSERT_OK(CreateIteratorContext(dataset_context.get(), &iterator_ctx));
   std::unique_ptr<IteratorBase> iterator;
   TF_ASSERT_OK(
-      shuffle_dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));
+      dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));
 
-  EXPECT_EQ(iterator->prefix(), "Iterator::Shuffle");
+  if (count_value == 1) {
+    EXPECT_EQ(iterator->prefix(), "Iterator::Shuffle");
+  } else {
+    EXPECT_EQ(iterator->prefix(), "Iterator::ShuffleAndRepeat");
+  }
 }
 
 TEST_P(ParameterizedShuffleDatasetOpTest, Roundtrip) {
@@ -555,10 +795,13 @@ TEST_P(ParameterizedShuffleDatasetOpTest, Roundtrip) {
   TF_ASSERT_OK(InitThreadPool(thread_num));
   TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));
 
-  std::unique_ptr<OpKernel> shuffle_dataset_kernel;
-  TF_ASSERT_OK(CreateShuffleDatasetOpKernel(
-      test_case.reshuffle_each_iteration, test_case.expected_output_dtypes,
-      test_case.expected_output_shapes, &shuffle_dataset_kernel));
+  Tensor count = test_case.count;
+  int64 count_value = count.flat<int64>()(0);
+  std::unique_ptr<OpKernel> dataset_kernel;
+  TF_ASSERT_OK(
+      CreateDatasetOpKernel(count_value, test_case.reshuffle_each_iteration,
+                            test_case.expected_output_dtypes,
+                            test_case.expected_output_shapes, &dataset_kernel));
 
   DatasetBase* range_dataset;
   TF_ASSERT_OK(CreateRangeDataset<int64>(
@@ -572,21 +815,21 @@ TEST_P(ParameterizedShuffleDatasetOpTest, Roundtrip) {
   Tensor seed2 = test_case.seed2;
   gtl::InlinedVector<TensorValue, 4> inputs(
       {&range_dataset_tensor, &buffer_size, &seed, &seed2});
+  if (count_value != 1) inputs.push_back(&count);
 
-  std::unique_ptr<OpKernelContext> shuffle_dataset_context;
-  TF_ASSERT_OK(CreateShuffleDatasetContext(shuffle_dataset_kernel.get(),
-                                           &inputs, &shuffle_dataset_context));
-  DatasetBase* shuffle_dataset;
-  TF_ASSERT_OK(CreateDataset(shuffle_dataset_kernel.get(),
-                             shuffle_dataset_context.get(), &shuffle_dataset));
-  core::ScopedUnref scoped_unref_shuffle_dataset(shuffle_dataset);
+  std::unique_ptr<OpKernelContext> dataset_context;
+  TF_ASSERT_OK(
+      CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
+  DatasetBase* dataset;
+  TF_ASSERT_OK(
+      CreateDataset(dataset_kernel.get(), dataset_context.get(), &dataset));
+  core::ScopedUnref scoped_unref_dataset(dataset);
 
   std::unique_ptr<IteratorContext> iterator_ctx;
-  TF_ASSERT_OK(
-      CreateIteratorContext(shuffle_dataset_context.get(), &iterator_ctx));
+  TF_ASSERT_OK(CreateIteratorContext(dataset_context.get(), &iterator_ctx));
   std::unique_ptr<IteratorBase> iterator;
   TF_ASSERT_OK(
-      shuffle_dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));
+      dataset->MakeIterator(iterator_ctx.get(), "Iterator", &iterator));
 
   std::unique_ptr<SerializationContext> serialization_ctx;
   TF_ASSERT_OK(CreateSerializationContext(&serialization_ctx));
@@ -602,7 +845,7 @@ TEST_P(ParameterizedShuffleDatasetOpTest, Roundtrip) {
     TF_EXPECT_OK(writer.Flush());
     VariantTensorDataReader reader(&data);
     TF_EXPECT_OK(RestoreIterator(iterator_ctx.get(), &reader, "Iterator",
-                                 *shuffle_dataset, &iterator));
+                                 *dataset, &iterator));
 
     while (cur_iteration <= breakpoint) {
       std::vector<Tensor> next;
@@ -613,51 +856,58 @@ TEST_P(ParameterizedShuffleDatasetOpTest, Roundtrip) {
     }
   }
 
-  // When `buffer_size = 1`, the output sequence of `ShuffleDataset` will be in
-  // order, so we need to consider the element sequence when evaluating the
-  // result for this case.
-  bool expect_items_equal = test_case.buffer_size.flat<int64>()(0) > 1;
-  TF_EXPECT_OK(ExpectEqual(out_tensors, test_case.expected_outputs,
-                           /*expect_items_equal*/ expect_items_equal));
+  TF_EXPECT_OK(ExpectEqual(out_tensors, test_case.expected_shuffle_outputs,
+                           /*compare_order*/ true));
 }
 
 INSTANTIATE_TEST_SUITE_P(ShuffleDatasetOpTest,
                          ParameterizedShuffleDatasetOpTest,
                          ::testing::ValuesIn(std::vector<TestCase>(
-                             {TestCase1(), TestCase2(), TestCase3()})));
+                             {TestCase1(), TestCase2(), TestCase3(),
+                              TestCase4(), TestCase5(), TestCase6(),
+                              TestCase7(), TestCase8()})));
 
-TEST_F(ShuffleDatasetOpTest, InvalidBufferSize) {
+TEST_F(ShuffleDatasetOpTest, InvalidArguments) {
   int thread_num = 2, cpu_num = 2;
-  TestCase test_case = InvalidBufferSizeTestCase();
+  std::vector<TestCase> test_cases = {
+      InvalidBufferSizeTestCaseForShuffleDataset(),
+      InvalidBufferSizeTestCaseForShuffleAndRepeatDataset(),
+      InvalidCountTestCaseForShuffleAndRepeatDataset()};
   TF_ASSERT_OK(InitThreadPool(thread_num));
   TF_ASSERT_OK(InitFunctionLibraryRuntime({}, cpu_num));
 
-  std::unique_ptr<OpKernel> shuffle_dataset_kernel;
-  TF_ASSERT_OK(CreateShuffleDatasetOpKernel(
-      test_case.reshuffle_each_iteration, test_case.expected_output_dtypes,
-      test_case.expected_output_shapes, &shuffle_dataset_kernel));
+  for (const auto& test_case : test_cases) {
+    Tensor count = test_case.count;
+    int64 count_value = count.flat<int64>()(0);
+    std::unique_ptr<OpKernel> dataset_kernel;
+    TF_ASSERT_OK(CreateDatasetOpKernel(
+        count_value, test_case.reshuffle_each_iteration,
+        test_case.expected_output_dtypes, test_case.expected_output_shapes,
+        &dataset_kernel));
 
-  DatasetBase* range_dataset;
-  TF_ASSERT_OK(CreateRangeDataset<int64>(
-      test_case.range_data_param.start, test_case.range_data_param.end,
-      test_case.range_data_param.step, "range", &range_dataset));
-  Tensor range_dataset_tensor(DT_VARIANT, TensorShape({}));
-  TF_ASSERT_OK(
-      StoreDatasetInVariantTensor(range_dataset, &range_dataset_tensor));
-  Tensor buffer_size = test_case.buffer_size;
-  Tensor seed = test_case.seed;
-  Tensor seed2 = test_case.seed2;
-  gtl::InlinedVector<TensorValue, 4> inputs(
-      {&range_dataset_tensor, &buffer_size, &seed, &seed2});
+    DatasetBase* range_dataset;
+    TF_ASSERT_OK(CreateRangeDataset<int64>(
+        test_case.range_data_param.start, test_case.range_data_param.end,
+        test_case.range_data_param.step, "range", &range_dataset));
+    Tensor range_dataset_tensor(DT_VARIANT, TensorShape({}));
+    TF_ASSERT_OK(
+        StoreDatasetInVariantTensor(range_dataset, &range_dataset_tensor));
+    Tensor buffer_size = test_case.buffer_size;
+    Tensor seed = test_case.seed;
+    Tensor seed2 = test_case.seed2;
+    gtl::InlinedVector<TensorValue, 4> inputs(
+        {&range_dataset_tensor, &buffer_size, &seed, &seed2});
+    if (count_value != 1) inputs.push_back(&count);
 
-  std::unique_ptr<OpKernelContext> shuffle_dataset_context;
-  TF_ASSERT_OK(CreateShuffleDatasetContext(shuffle_dataset_kernel.get(),
-                                           &inputs, &shuffle_dataset_context));
-  DatasetBase* shuffle_dataset;
-  EXPECT_EQ(CreateDataset(shuffle_dataset_kernel.get(),
-                          shuffle_dataset_context.get(), &shuffle_dataset)
-                .code(),
-            tensorflow::error::INVALID_ARGUMENT);
+    std::unique_ptr<OpKernelContext> dataset_context;
+    TF_ASSERT_OK(
+        CreateDatasetContext(dataset_kernel.get(), &inputs, &dataset_context));
+    DatasetBase* shuffle_dataset;
+    EXPECT_EQ(CreateDataset(dataset_kernel.get(), dataset_context.get(),
+                            &shuffle_dataset)
+                  .code(),
+              tensorflow::error::INVALID_ARGUMENT);
+  }
 }
 
 }  // namespace