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    A_ (       a_ ):

  def __init__(		self :      int     ,							__SCREAMING_SNAKE_CASE :      Union[str, Any]=0.01     ,							__SCREAMING_SNAKE_CASE :      List[str]=10_00       ):
   __a       =  p_stop
   __a       =  max_length

  def __iter__(		self :      Dict       ):
   __a       =  0
   __a       =  False
   while not stop and count < self.max_length:
    yield count
    count += 1
    __a       =  random.random() < self.p_stop
class    A_ (       unittest.TestCase ):

  def 		_UpperCAmelCase   (		self :      Tuple     ,							__SCREAMING_SNAKE_CASE :      Any     ,							__SCREAMING_SNAKE_CASE :      Union[str, Any]     ,							__SCREAMING_SNAKE_CASE :      str=False     ,							__SCREAMING_SNAKE_CASE :      List[Any]=True       ):
   __a       =  [
       BatchSamplerShard(__SCREAMING_SNAKE_CASE     ,							2     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       )
       for i in range(2       )
   ]
   __a       =  [list(__SCREAMING_SNAKE_CASE       ) for batch_sampler_shard in batch_sampler_shards]
   if not split_batches:
    self.assertListEqual([len(__SCREAMING_SNAKE_CASE       ) for shard in batch_sampler_shards]     ,							[len(__SCREAMING_SNAKE_CASE       ) for e in expected]       )
   self.assertListEqual(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE       )

  def 		_UpperCAmelCase   (		self :      Optional[Any]       ):
   # Check the shards when the dataset is a round multiple of total batch size.
   __a       =  BatchSampler(range(24       )     ,							batch_size=3     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE       )

   __a       =  BatchSampler(range(24       )     ,							batch_size=3     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   # Expected shouldn't change
   self.check_batch_sampler_shards(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE       )

   # Check the shards when the dataset is a round multiple of batch size but not total batch size.
   __a       =  BatchSampler(range(21       )     ,							batch_size=3     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE       )

   __a       =  BatchSampler(range(21       )     ,							batch_size=3     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
       [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
   ]
   self.check_batch_sampler_shards(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE       )

   # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
   # num_processes batch.
   __a       =  BatchSampler(range(22       )     ,							batch_size=3     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE       )

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

   # Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
   # num_processes batch.
   __a       =  BatchSampler(range(20       )     ,							batch_size=3     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE       )

   __a       =  BatchSampler(range(20       )     ,							batch_size=3     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
       [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
   ]
   self.check_batch_sampler_shards(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE       )

   # Check the shards when the dataset is very small.
   __a       =  BatchSampler(range(2       )     ,							batch_size=3     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [[[0, 1, 0]], [[1, 0, 1]]]
   self.check_batch_sampler_shards(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE       )

   __a       =  BatchSampler(range(2       )     ,							batch_size=3     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [[], []]
   self.check_batch_sampler_shards(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE       )

  def 		_UpperCAmelCase   (		self :      Any       ):
   # Check the shards when the dataset is a round multiple of batch size.
   __a       =  BatchSampler(range(24       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE       )

   __a       =  BatchSampler(range(24       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   # Expected shouldn't change
   self.check_batch_sampler_shards(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE       )

   # Check the shards when the dataset is not a round multiple of batch size.
   __a       =  BatchSampler(range(22       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE       )

   __a       =  BatchSampler(range(22       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE       )

   # Check the shards when the dataset is not a round multiple of batch size or num_processes.
   __a       =  BatchSampler(range(21       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE       )

   __a       =  BatchSampler(range(21       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE       )

   # Check the shards when the dataset is very small.
   __a       =  BatchSampler(range(2       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [[[0, 1]], [[0, 1]]]
   self.check_batch_sampler_shards(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE       )

   __a       =  BatchSampler(range(2       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [[], []]
   self.check_batch_sampler_shards(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE       )

  def 		_UpperCAmelCase   (		self :      Dict       ):
   # Check the shards when the dataset is a round multiple of total batch size.
   __a       =  BatchSampler(range(24       )     ,							batch_size=3     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       )

   __a       =  BatchSampler(range(24       )     ,							batch_size=3     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   # Expected shouldn't change
   self.check_batch_sampler_shards(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       )

   # Check the shards when the dataset is a round multiple of batch size but not total batch size.
   __a       =  BatchSampler(range(21       )     ,							batch_size=3     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       )

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

   # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
   # num_processes batch.
   __a       =  BatchSampler(range(22       )     ,							batch_size=3     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       )

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

   # Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
   # num_processes batch.
   __a       =  BatchSampler(range(20       )     ,							batch_size=3     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       )

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

   # Check the shards when the dataset is very small.
   __a       =  BatchSampler(range(2       )     ,							batch_size=3     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [[[0, 1]], []]
   self.check_batch_sampler_shards(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       )

   __a       =  BatchSampler(range(2       )     ,							batch_size=3     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [[], []]
   self.check_batch_sampler_shards(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       )

  def 		_UpperCAmelCase   (		self :      Union[str, Any]       ):
   # Check the shards when the dataset is a round multiple of batch size.
   __a       =  BatchSampler(range(24       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       )

   __a       =  BatchSampler(range(24       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   # Expected shouldn't change
   self.check_batch_sampler_shards(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       )

   # Check the shards when the dataset is not a round multiple of batch size.
   __a       =  BatchSampler(range(22       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       )

   __a       =  BatchSampler(range(22       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       )

   # Check the shards when the dataset is not a round multiple of batch size or num_processes.
   __a       =  BatchSampler(range(21       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       )

   __a       =  BatchSampler(range(21       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [
       [[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(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       )

   # Check the shards when the dataset is very small.
   __a       =  BatchSampler(range(2       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [[[0, 1]], []]
   self.check_batch_sampler_shards(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       )

   __a       =  BatchSampler(range(2       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  [[], []]
   self.check_batch_sampler_shards(__SCREAMING_SNAKE_CASE     ,							__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       )

  def 		_UpperCAmelCase   (		self :      int       ):
   __a       =  [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
   __a       =  [BatchSamplerShard(__SCREAMING_SNAKE_CASE     ,							2     ,							__SCREAMING_SNAKE_CASE     ,							even_batches=__SCREAMING_SNAKE_CASE       ) 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 		_UpperCAmelCase   (		self :      Optional[Any]     ,							__SCREAMING_SNAKE_CASE :      Tuple     ,							__SCREAMING_SNAKE_CASE :      Any     ,							__SCREAMING_SNAKE_CASE :      Optional[int]     ,							__SCREAMING_SNAKE_CASE :      Optional[int]=False     ,							__SCREAMING_SNAKE_CASE :      Union[str, Any]=2     ,							__SCREAMING_SNAKE_CASE :      Tuple=False       ):
   random.seed(__SCREAMING_SNAKE_CASE       )
   __a       =  list(__SCREAMING_SNAKE_CASE       )

   __a       =  [
       IterableDatasetShard(
           __SCREAMING_SNAKE_CASE     ,							batch_size=__SCREAMING_SNAKE_CASE     ,							drop_last=__SCREAMING_SNAKE_CASE     ,							num_processes=__SCREAMING_SNAKE_CASE     ,							process_index=__SCREAMING_SNAKE_CASE     ,							split_batches=__SCREAMING_SNAKE_CASE     ,							)
       for i in range(__SCREAMING_SNAKE_CASE       )
   ]
   __a       =  []
   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(__SCREAMING_SNAKE_CASE       )
    iterable_dataset_lists.append(list(__SCREAMING_SNAKE_CASE       )       )

   __a       =  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
   __a       =  iterable_dataset_lists[0]
   for l in iterable_dataset_lists[1:]:
    self.assertEqual(len(__SCREAMING_SNAKE_CASE       )     ,							len(__SCREAMING_SNAKE_CASE       )       )
    self.assertTrue(len(__SCREAMING_SNAKE_CASE       ) % shard_batch_size == 0       )

   __a       =  []
   for idx in range(0     ,							len(__SCREAMING_SNAKE_CASE       )     ,							__SCREAMING_SNAKE_CASE       ):
    for l in iterable_dataset_lists:
     observed += l[idx : idx + shard_batch_size]

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

  def 		_UpperCAmelCase   (		self :      Tuple       ):
   __a       =  42
   __a       =  RandomIterableDataset()

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

   # Edge case with a very small dataset
   __a       =  RandomIterableDataset(max_length=2       )

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

  def 		_UpperCAmelCase   (		self :      int       ):
   __a       =  BatchSampler(range(16       )     ,							batch_size=4     ,							drop_last=__SCREAMING_SNAKE_CASE       )
   __a       =  SkipBatchSampler(__SCREAMING_SNAKE_CASE     ,							2       )
   self.assertListEqual(list(__SCREAMING_SNAKE_CASE       )     ,							[[8, 9, 10, 11], [12, 13, 14, 15]]       )

  def 		_UpperCAmelCase   (		self :      Dict       ):
   __a       =  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 		_UpperCAmelCase   (		self :      Tuple       ):
   __a       =  DataLoader(list(range(16       )       )     ,							batch_size=4       )
   __a       =  skip_first_batches(__SCREAMING_SNAKE_CASE     ,							num_batches=2       )
   self.assertListEqual([t.tolist() for t in new_dataloader]     ,							[[8, 9, 10, 11], [12, 13, 14, 15]]       )

  def 		_UpperCAmelCase   (		self :      int       ):
   __a       =  DataLoaderShard(list(range(16       )       )     ,							batch_size=4       )
   for idx, _ in enumerate(__SCREAMING_SNAKE_CASE       ):
    self.assertEqual(dataloader.end_of_dataloader     ,							idx == 3       )

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

  def 		_UpperCAmelCase   (		self :      Dict       ):
   Accelerator()
   __a       =  DataLoaderDispatcher(range(16       )     ,							batch_size=4       )
   for idx, _ in enumerate(__SCREAMING_SNAKE_CASE       ):
    self.assertEqual(dataloader.end_of_dataloader     ,							idx == 3       )

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



import unittest

from transformers import AutoTokenizer, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow


if is_flax_available():
  import jax.numpy as jnp

  from transformers import FlaxXLMRobertaModel
@require_sentencepiece
@require_tokenizers
@require_flax
class    A_ (       unittest.TestCase ):

  @slow
  def 		_UpperCAmelCase   (		self :      Dict       ):
   __a       =  FlaxXLMRobertaModel.from_pretrained("xlm-roberta-base"       )
   __a       =  AutoTokenizer.from_pretrained("xlm-roberta-base"       )
   __a       =  "The dog is cute and lives in the garden house"
   __a       =  jnp.array([tokenizer.encode(__SCREAMING_SNAKE_CASE       )]       )

   __a       =  (1, 12, 7_68)  # batch_size, sequence_length, embedding_vector_dim
   __a       =  jnp.array(
       [[-0.01_01, 0.12_18, -0.08_03, 0.08_01, 0.13_27, 0.07_76, -0.12_15, 0.23_83, 0.33_38, 0.31_06, 0.03_00, 0.02_52]]       )

   __a       =  model(__SCREAMING_SNAKE_CASE       )["last_hidden_state"]
   self.assertEqual(output.shape     ,							__SCREAMING_SNAKE_CASE       )
   # compare the actual values for a slice of last dim
   self.assertTrue(jnp.allclose(output[:, :, -1]     ,							__SCREAMING_SNAKE_CASE     ,							atol=1E-3       )       )


