






import random
import unittest

from torch.utils.data import BatchSampler, DataLoader, IterableDataset

from accelerate import Accelerator
from accelerate.data_loader import (
    BatchSamplerShard,
    DataLoaderDispatcher,
    DataLoaderShard,
    IterableDatasetShard,
    SkipBatchSampler,
    SkipDataLoader,
    skip_first_batches,
)







class 	_UpperCAmelCase  (  _A		):




      """simple docstring"""



      def __init__(   self    ,	_lowerCAmelCase=0.01    ,	_lowerCAmelCase=1_000		):



          '''simple docstring'''
          lowerCAmelCase__   :Tuple      							=				p_stop
          lowerCAmelCase__   :Tuple      							=				max_length



      def __iter__(   self		):



          '''simple docstring'''
          lowerCAmelCase__   :List[str]      							=				0
          lowerCAmelCase__   :Dict      							=				False
          while not stop and count < self.max_length:
              yield count
              count += 1
              lowerCAmelCase__   :int      							=				random.random() < self.p_stop







class 	_UpperCAmelCase  (  unittest.TestCase		):




      """simple docstring"""



      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase=False    ,	_lowerCAmelCase=True		):



          '''simple docstring'''
          lowerCAmelCase__   :Union[str, Any]      							=				[
              BatchSamplerShard(_lowerCAmelCase    ,	2    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)
              for i in range(2		)
          ]
          lowerCAmelCase__   :List[str]      							=				[list(_lowerCAmelCase		) for batch_sampler_shard in batch_sampler_shards]
          if not split_batches:
              self.assertListEqual([len(_lowerCAmelCase		) for shard in batch_sampler_shards]    ,	[len(_lowerCAmelCase		) for e in expected]		)
          self.assertListEqual(_lowerCAmelCase    ,	_lowerCAmelCase		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :List[Any]      							=				BatchSampler(range(24		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Dict      							=				[
              [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
              [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase		)

          lowerCAmelCase__   :Optional[Any]      							=				BatchSampler(range(24		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          # Expected shouldn't change
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase		)

          # Check the shards when the dataset is a round multiple of batch size but not total batch size.
          lowerCAmelCase__   :Any      							=				BatchSampler(range(21		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Optional[int]      							=				[
              [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
              [[3, 4, 5], [9, 10, 11], [15, 16, 17], [0, 1, 2]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase		)

          lowerCAmelCase__   :List[Any]      							=				BatchSampler(range(21		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Optional[int]      							=				[
              [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
              [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase		)

          # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
          # num_processes batch.
          lowerCAmelCase__   :Any      							=				BatchSampler(range(22		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Union[str, Any]      							=				[
              [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
              [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 0, 1]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase		)

          lowerCAmelCase__   :Any      							=				BatchSampler(range(22		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Dict      							=				[
              [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
              [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase		)

          # Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
          # num_processes batch.
          lowerCAmelCase__   :Optional[Any]      							=				BatchSampler(range(20		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :str      							=				[
              [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 0]],
              [[3, 4, 5], [9, 10, 11], [15, 16, 17], [1, 2, 3]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase		)

          lowerCAmelCase__   :List[Any]      							=				BatchSampler(range(20		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Union[str, Any]      							=				[
              [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
              [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase		)

          # Check the shards when the dataset is very small.
          lowerCAmelCase__   :Union[str, Any]      							=				BatchSampler(range(2		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :List[str]      							=				[[[0, 1, 0]], [[1, 0, 1]]]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase		)

          lowerCAmelCase__   :int      							=				BatchSampler(range(2		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :List[Any]      							=				[[], []]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Union[str, Any]      							=				BatchSampler(range(24		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Optional[int]      							=				[
              [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
              [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase		)

          lowerCAmelCase__   :Optional[Any]      							=				BatchSampler(range(24		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          # Expected shouldn't change
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase		)

          # Check the shards when the dataset is not a round multiple of batch size.
          lowerCAmelCase__   :Optional[Any]      							=				BatchSampler(range(22		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Tuple      							=				[
              [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
              [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [0, 1]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase		)

          lowerCAmelCase__   :int      							=				BatchSampler(range(22		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :List[Any]      							=				[
              [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
              [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase		)

          # Check the shards when the dataset is not a round multiple of batch size or num_processes.
          lowerCAmelCase__   :Optional[Any]      							=				BatchSampler(range(21		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Dict      							=				[
              [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 0]],
              [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [1, 2]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase		)

          lowerCAmelCase__   :Optional[int]      							=				BatchSampler(range(21		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Tuple      							=				[
              [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
              [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase		)

          # Check the shards when the dataset is very small.
          lowerCAmelCase__   :Tuple      							=				BatchSampler(range(2		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Optional[Any]      							=				[[[0, 1]], [[0, 1]]]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase		)

          lowerCAmelCase__   :str      							=				BatchSampler(range(2		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Union[str, Any]      							=				[[], []]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Any      							=				BatchSampler(range(24		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :str      							=				[
              [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
              [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21, 22, 23]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)

          lowerCAmelCase__   :List[Any]      							=				BatchSampler(range(24		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          # Expected shouldn't change
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)

          # Check the shards when the dataset is a round multiple of batch size but not total batch size.
          lowerCAmelCase__   :Optional[int]      							=				BatchSampler(range(21		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :int      							=				[
              [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
              [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)

          lowerCAmelCase__   :List[Any]      							=				BatchSampler(range(21		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Tuple      							=				[
              [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
              [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)

          # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
          # num_processes batch.
          lowerCAmelCase__   :Union[str, Any]      							=				BatchSampler(range(22		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :str      							=				[
              [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
              [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)

          lowerCAmelCase__   :Optional[Any]      							=				BatchSampler(range(22		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :List[str]      							=				[
              [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
              [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)

          # Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
          # num_processes batch.
          lowerCAmelCase__   :Optional[int]      							=				BatchSampler(range(20		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Optional[Any]      							=				[
              [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19]],
              [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)

          lowerCAmelCase__   :Union[str, Any]      							=				BatchSampler(range(20		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Any      							=				[
              [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
              [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)

          # Check the shards when the dataset is very small.
          lowerCAmelCase__   :str      							=				BatchSampler(range(2		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Dict      							=				[[[0, 1]], []]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)

          lowerCAmelCase__   :Union[str, Any]      							=				BatchSampler(range(2		)    ,	batch_size=3    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Optional[int]      							=				[[], []]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :str      							=				BatchSampler(range(24		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Tuple      							=				[
              [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
              [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19], [22, 23]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)

          lowerCAmelCase__   :List[Any]      							=				BatchSampler(range(24		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          # Expected shouldn't change
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)

          # Check the shards when the dataset is not a round multiple of batch size.
          lowerCAmelCase__   :Union[str, Any]      							=				BatchSampler(range(22		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :int      							=				[
              [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20, 21]],
              [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)

          lowerCAmelCase__   :Optional[Any]      							=				BatchSampler(range(22		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Optional[int]      							=				[
              [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
              [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)

          # Check the shards when the dataset is not a round multiple of batch size or num_processes.
          lowerCAmelCase__   :Dict      							=				BatchSampler(range(21		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :List[Any]      							=				[
              [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20]],
              [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)

          lowerCAmelCase__   :List[Any]      							=				BatchSampler(range(21		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Union[str, Any]      							=				[
              [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17]],
              [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
          ]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)

          # Check the shards when the dataset is very small.
          lowerCAmelCase__   :List[Any]      							=				BatchSampler(range(2		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :int      							=				[[[0, 1]], []]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)

          lowerCAmelCase__   :str      							=				BatchSampler(range(2		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Dict      							=				[[], []]
          self.check_batch_sampler_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	split_batches=_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Tuple      							=				[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
          lowerCAmelCase__   :Optional[int]      							=				[BatchSamplerShard(_lowerCAmelCase    ,	2    ,	_lowerCAmelCase    ,	even_batches=_lowerCAmelCase		) for i in range(2		)]

          self.assertEqual(len(batch_sampler_shards[0]		)    ,	3		)
          self.assertEqual(len(batch_sampler_shards[1]		)    ,	2		)

          self.assertListEqual(list(batch_sampler_shards[0]		)    ,	[[0, 1, 2], [5, 6, 7, 8], [12, 13]]		)
          self.assertListEqual(list(batch_sampler_shards[1]		)    ,	[[3, 4], [9, 10, 11]]		)
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase=False    ,	_lowerCAmelCase=2    ,	_lowerCAmelCase=False		):



          '''simple docstring'''
          random.seed(_lowerCAmelCase		)
          lowerCAmelCase__   :Union[str, Any]      							=				list(_lowerCAmelCase		)

          lowerCAmelCase__   :Tuple      							=				[
              IterableDatasetShard(
                  _lowerCAmelCase    ,	batch_size=_lowerCAmelCase    ,	drop_last=_lowerCAmelCase    ,	num_processes=_lowerCAmelCase    ,	process_index=_lowerCAmelCase    ,	split_batches=_lowerCAmelCase    ,	)
              for i in range(_lowerCAmelCase		)
          ]
          lowerCAmelCase__   :Optional[Any]      							=				[]
          for iterable_dataset_shard in iterable_dataset_shards:
              # Since our random iterable dataset will be... random... we need to use a seed to get reproducible results.
              random.seed(_lowerCAmelCase		)
              iterable_dataset_lists.append(list(_lowerCAmelCase		)		)

          lowerCAmelCase__   :List[str]      							=				batch_size // num_processes if split_batches else batch_size
          # All iterable dataset shard should have the same length, a round multiple of shard_batch_size
          lowerCAmelCase__   :Optional[int]      							=				iterable_dataset_lists[0]
          for l in iterable_dataset_lists[1:]:
              self.assertEqual(len(_lowerCAmelCase		)    ,	len(_lowerCAmelCase		)		)
              self.assertTrue(len(_lowerCAmelCase		) % shard_batch_size == 0		)

          lowerCAmelCase__   :Optional[Any]      							=				[]
          for idx in range(0    ,	len(_lowerCAmelCase		)    ,	_lowerCAmelCase		):
              for l in iterable_dataset_lists:
                  observed += l[idx : idx + shard_batch_size]

          if not drop_last:
              while len(_lowerCAmelCase		) < len(_lowerCAmelCase		):
                  reference += reference
          self.assertListEqual(_lowerCAmelCase    ,	reference[: len(_lowerCAmelCase		)]		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Tuple      							=				42
          lowerCAmelCase__   :Any      							=				RandomIterableDataset()

          self.check_iterable_dataset_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	batch_size=4    ,	drop_last=_lowerCAmelCase    ,	split_batches=_lowerCAmelCase		)
          self.check_iterable_dataset_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	batch_size=4    ,	drop_last=_lowerCAmelCase    ,	split_batches=_lowerCAmelCase		)
          self.check_iterable_dataset_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	batch_size=4    ,	drop_last=_lowerCAmelCase    ,	split_batches=_lowerCAmelCase		)
          self.check_iterable_dataset_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	batch_size=4    ,	drop_last=_lowerCAmelCase    ,	split_batches=_lowerCAmelCase		)

          # Edge case with a very small dataset
          lowerCAmelCase__   :List[str]      							=				RandomIterableDataset(max_length=2		)

          self.check_iterable_dataset_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	batch_size=4    ,	drop_last=_lowerCAmelCase    ,	split_batches=_lowerCAmelCase		)
          self.check_iterable_dataset_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	batch_size=4    ,	drop_last=_lowerCAmelCase    ,	split_batches=_lowerCAmelCase		)
          self.check_iterable_dataset_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	batch_size=4    ,	drop_last=_lowerCAmelCase    ,	split_batches=_lowerCAmelCase		)
          self.check_iterable_dataset_shards(_lowerCAmelCase    ,	_lowerCAmelCase    ,	batch_size=4    ,	drop_last=_lowerCAmelCase    ,	split_batches=_lowerCAmelCase		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Any      							=				BatchSampler(range(16		)    ,	batch_size=4    ,	drop_last=_lowerCAmelCase		)
          lowerCAmelCase__   :Optional[Any]      							=				SkipBatchSampler(_lowerCAmelCase    ,	2		)
          self.assertListEqual(list(_lowerCAmelCase		)    ,	[[8, 9, 10, 11], [12, 13, 14, 15]]		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Any      							=				SkipDataLoader(list(range(16		)		)    ,	batch_size=4    ,	skip_batches=2		)
          self.assertListEqual([t.tolist() for t in dataloader]    ,	[[8, 9, 10, 11], [12, 13, 14, 15]]		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Any      							=				DataLoader(list(range(16		)		)    ,	batch_size=4		)
          lowerCAmelCase__   :Optional[Any]      							=				skip_first_batches(_lowerCAmelCase    ,	num_batches=2		)
          self.assertListEqual([t.tolist() for t in new_dataloader]    ,	[[8, 9, 10, 11], [12, 13, 14, 15]]		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Optional[int]      							=				DataLoaderShard(list(range(16		)		)    ,	batch_size=4		)
          for idx, _ in enumerate(_lowerCAmelCase		):
              self.assertEqual(dataloader.end_of_dataloader    ,	idx == 3		)

          # Test it also works on the second iteration
          for idx, _ in enumerate(_lowerCAmelCase		):
              self.assertEqual(dataloader.end_of_dataloader    ,	idx == 3		)



      def 			snake_case_      (   self		):



          '''simple docstring'''
          Accelerator()
          lowerCAmelCase__   :Optional[Any]      							=				DataLoaderDispatcher(range(16		)    ,	batch_size=4		)
          for idx, _ in enumerate(_lowerCAmelCase		):
              self.assertEqual(dataloader.end_of_dataloader    ,	idx == 3		)

          # Test it also works on the second iteration
          for idx, _ in enumerate(_lowerCAmelCase		):
              self.assertEqual(dataloader.end_of_dataloader    ,	idx == 3		)









def 				snake_case__							(							UpperCAmelCase							:  float       ):

    if edge <= 0 or not isinstance(UpperCAmelCase    ,			UpperCAmelCase       ):
        raise ValueError("Length must be a positive."       )
    return 3 * ((2_5 + 1_0 * (5 ** (1 / 2))) ** (1 / 2)) * (edge**2)


def 				snake_case__							(							UpperCAmelCase							:  float       ):

    if edge <= 0 or not isinstance(UpperCAmelCase    ,			UpperCAmelCase       ):
        raise ValueError("Length must be a positive."       )
    return ((1_5 + (7 * (5 ** (1 / 2)))) / 4) * (edge**3)


if __name__ == "__main__":
 import doctest

 doctest.testmod()



