message: shard the dataset in map_and_batch_test
diff: @@ -412,6 +412,7 @@ class MapAndBatchDatasetCheckpointTest(checkpoint_test_base.CheckpointTestBase,
  num_outputs_drop_remainder = total_outputs // batch_size
  num_outputs_keep_remainder = int(math.ceil(total_outputs / batch_size))
  num_parallel_batches = 2
  + num_shards = 3
  def build_ds(range_start, drop_remainder=False):
  @@ -419,7 +420,8 @@ class MapAndBatchDatasetCheckpointTest(checkpoint_test_base.CheckpointTestBase,
  return math_ops.square(x)
  return dataset_ops.Dataset.range(
  - range_start, range_start + range_size).repeat(num_repeats).apply(
  + range_start, range_start + range_size).shard(
  + num_shards=num_shards, index=0).repeat(num_repeats).apply(
  batching.map_and_batch(
  map_func=_map_fn,
  batch_size=batch_size,
  @@ -438,6 +440,7 @@ class MapAndBatchDatasetCheckpointTest(checkpoint_test_base.CheckpointTestBase,
  num_outputs_drop_remainder = total_outputs // batch_size
  num_outputs_keep_remainder = int(math.ceil(total_outputs / batch_size))
  num_parallel_calls = 7
  + num_shards = 3
  def build_ds(range_start, drop_remainder=False):
  @@ -445,7 +448,8 @@ class MapAndBatchDatasetCheckpointTest(checkpoint_test_base.CheckpointTestBase,
  return math_ops.square(x)
  return dataset_ops.Dataset.range(
  - range_start, range_start + range_size).repeat(num_repeats).apply(
  + range_start, range_start + range_size).shard(
  + num_shards=num_shards, index=0).repeat(num_repeats).apply(
  batching.map_and_batch(
  map_func=_map_fn,
  batch_size=batch_size,
  