code: def test_sharded_and_parallel_iter_same_num_batches():
      """Tests that a sharded data iterator with just a single shard produces as many shards as an iterator directly
      using the same dataset."""
      batch_size = 2
      num_batches_per_bucket = 10
      buckets = data_io.define_parallel_buckets(100, 100, 10, True, 1.0)
      bucket_counts = [
          batch_size * num_batches_per_bucket for _ in buckets
      ]
      num_batches = num_batches_per_bucket * len(buckets)
      bucket_batch_sizes = data_io.define_bucket_batch_sizes(
          buckets,
          batch_size,
          batch_type=C.BATCH_TYPE_SENTENCE,
          data_target_average_len=[None] * len(buckets),
      )
  
      dataset = data_io.ParallelDataSet(
          *_get_random_bucketed_data(
              buckets,
              min_count=0,
              max_count=5,
              bucket_counts=bucket_counts,
          )
      )
  
      with TemporaryDirectory() as work_dir:
          shard_fname = os.path.join(work_dir, "shard1")
          dataset.save(shard_fname)
          shard_fnames = [shard_fname]
  
          it_sharded = data_io.ShardedParallelSampleIter(
              shard_fnames, buckets, batch_size, bucket_batch_sizes
          )
  
          it_parallel = data_io.ParallelSampleIter(
              dataset, buckets, batch_size, bucket_batch_sizes
          )
  
          num_batches_seen = 0
          while it_parallel.iter_next():
              assert it_sharded.iter_next()
              it_parallel.next()
              it_sharded.next()
              num_batches_seen += 1
          assert num_batches_seen == num_batches
  
          print("Resetting...")
          it_sharded.reset()
          it_parallel.reset()
  
          num_batches_seen = 0
          while it_parallel.iter_next():
              assert it_sharded.iter_next()
              it_parallel.next()
              it_sharded.next()
  
              num_batches_seen += 1
  
          assert num_batches_seen == num_batches
  
imports: import os
  import random
  from tempfile import TemporaryDirectory
  from typing import Optional, List, Tuple
  import numpy
  import pytest
  import torch
  from sockeye import constants
  from sockeye import data_io
  from sockeye import utils
  from sockeye import vocab
  from sockeye.test_utils import tmp_digits_dataset
  from sockeye.utils import SockeyeError, get_tokens, seed_rngs
label_desc: Use image node_id 28 for calling a global function with example usage: test_sharded_and_parallel_iter_same_num_batches() without return types