




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					(  _lowerCamelCase      ):






  def __init__(  self		:	Dict      ,						_A		:	str=0.01      ,						_A		:	int=1_000      )		->    Optional[Any]:






       """simple docstring"""







       lowercase					:   Any								= p_stop
       lowercase					:   List[Any]								= max_length






  def __iter__(  self		:	int      )		->    Union[str, Any]:






       """simple docstring"""







       lowercase					:   Optional[int]								= 0
       lowercase					:   Dict								= False
       while not stop and count < self.max_length:
            yield count
            count += 1
            lowercase					:   int								= random.random() < self.p_stop


class 						_A					(  unittest.TestCase      ):






  def 					__a			(  self		:	Any      ,						_A		:	Union[str, Any]      ,						_A		:	List[str]      ,						_A		:	int=False      ,						_A		:	List[Any]=True      )		->    int:






       """simple docstring"""







       lowercase					:   Dict								= [
           BatchSamplerShard(lowerCAmelCase_      ,						2      ,						lowerCAmelCase_      ,						split_batches=lowerCAmelCase_      ,						even_batches=lowerCAmelCase_      )
           for i in range(2      )
       ]
       lowercase					:   Optional[int]								= [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 					__a			(  self		:	int      )		->    int:






       """simple docstring"""







       lowercase					:   str								= BatchSampler(range(24      )      ,						batch_size=3      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   int								= [
           [[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_      )

       lowercase					:   int								= 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.
       lowercase					:   str								= BatchSampler(range(21      )      ,						batch_size=3      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   Optional[Any]								= [
           [[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_      )

       lowercase					:   List[str]								= BatchSampler(range(21      )      ,						batch_size=3      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   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_      )

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

       lowercase					:   int								= BatchSampler(range(22      )      ,						batch_size=3      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   Optional[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 not a round multiple of batch size but and has not a multiple of
       # num_processes batch.
       lowercase					:   Union[str, Any]								= BatchSampler(range(20      )      ,						batch_size=3      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   Dict								= [
           [[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_      )

       lowercase					:   str								= BatchSampler(range(20      )      ,						batch_size=3      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   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.
       lowercase					:   Optional[Any]								= BatchSampler(range(2      )      ,						batch_size=3      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   List[Any]								= [[[0, 1, 0]], [[1, 0, 1]]]
       self.check_batch_sampler_shards(lowerCAmelCase_      ,						lowerCAmelCase_      )

       lowercase					:   List[str]								= BatchSampler(range(2      )      ,						batch_size=3      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   Any								= [[], []]
       self.check_batch_sampler_shards(lowerCAmelCase_      ,						lowerCAmelCase_      )






  def 					__a			(  self		:	Union[str, Any]      )		->    Dict:






       """simple docstring"""







       lowercase					:   Optional[Any]								= BatchSampler(range(24      )      ,						batch_size=4      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   List[str]								= [
           [[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_      )

       lowercase					:   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.
       lowercase					:   Dict								= BatchSampler(range(22      )      ,						batch_size=4      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   List[str]								= [
           [[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_      )

       lowercase					:   int								= BatchSampler(range(22      )      ,						batch_size=4      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   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_      )

       # Check the shards when the dataset is not a round multiple of batch size or num_processes.
       lowercase					:   List[str]								= BatchSampler(range(21      )      ,						batch_size=4      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   str								= [
           [[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_      )

       lowercase					:   int								= BatchSampler(range(21      )      ,						batch_size=4      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   Optional[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 very small.
       lowercase					:   Any								= BatchSampler(range(2      )      ,						batch_size=4      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   Optional[Any]								= [[[0, 1]], [[0, 1]]]
       self.check_batch_sampler_shards(lowerCAmelCase_      ,						lowerCAmelCase_      ,						split_batches=lowerCAmelCase_      )

       lowercase					:   str								= BatchSampler(range(2      )      ,						batch_size=4      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   str								= [[], []]
       self.check_batch_sampler_shards(lowerCAmelCase_      ,						lowerCAmelCase_      ,						split_batches=lowerCAmelCase_      )






  def 					__a			(  self		:	str      )		->    Union[str, Any]:






       """simple docstring"""







       lowercase					:   Dict								= BatchSampler(range(24      )      ,						batch_size=3      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   List[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_      )

       lowercase					:   str								= 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.
       lowercase					:   Union[str, Any]								= BatchSampler(range(21      )      ,						batch_size=3      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   Any								= [
           [[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_      )

       lowercase					:   Dict								= BatchSampler(range(21      )      ,						batch_size=3      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   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 not a round multiple of batch size but has a multiple of
       # num_processes batch.
       lowercase					:   Optional[int]								= BatchSampler(range(22      )      ,						batch_size=3      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   Any								= [
           [[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_      )

       lowercase					:   int								= BatchSampler(range(22      )      ,						batch_size=3      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   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.
       lowercase					:   Optional[Any]								= BatchSampler(range(20      )      ,						batch_size=3      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   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_      )

       lowercase					:   Union[str, Any]								= BatchSampler(range(20      )      ,						batch_size=3      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   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_      ,						even_batches=lowerCAmelCase_      )

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

       lowercase					:   str								= BatchSampler(range(2      )      ,						batch_size=3      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   Optional[Any]								= [[], []]
       self.check_batch_sampler_shards(lowerCAmelCase_      ,						lowerCAmelCase_      ,						even_batches=lowerCAmelCase_      )






  def 					__a			(  self		:	Dict      )		->    str:






       """simple docstring"""







       lowercase					:   List[Any]								= BatchSampler(range(24      )      ,						batch_size=4      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   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_      )

       lowercase					:   int								= 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.
       lowercase					:   List[Any]								= BatchSampler(range(22      )      ,						batch_size=4      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   Optional[Any]								= [
           [[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_      )

       lowercase					:   Optional[Any]								= BatchSampler(range(22      )      ,						batch_size=4      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   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.
       lowercase					:   str								= BatchSampler(range(21      )      ,						batch_size=4      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   List[str]								= [
           [[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_      )

       lowercase					:   Optional[Any]								= BatchSampler(range(21      )      ,						batch_size=4      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   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.
       lowercase					:   Any								= BatchSampler(range(2      )      ,						batch_size=4      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   Any								= [[[0, 1]], []]
       self.check_batch_sampler_shards(lowerCAmelCase_      ,						lowerCAmelCase_      ,						split_batches=lowerCAmelCase_      ,						even_batches=lowerCAmelCase_      )

       lowercase					:   Any								= BatchSampler(range(2      )      ,						batch_size=4      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   Optional[Any]								= [[], []]
       self.check_batch_sampler_shards(lowerCAmelCase_      ,						lowerCAmelCase_      ,						split_batches=lowerCAmelCase_      ,						even_batches=lowerCAmelCase_      )






  def 					__a			(  self		:	List[str]      )		->    Dict:






       """simple docstring"""







       lowercase					:   Optional[Any]								= [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
       lowercase					:   str								= [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 					__a			(  self		:	Union[str, Any]      ,						_A		:	List[Any]      ,						_A		:	Tuple      ,						_A		:	Tuple      ,						_A		:	Tuple=False      ,						_A		:	Tuple=2      ,						_A		:	List[str]=False      )		->    Optional[Any]:






       """simple docstring"""







       random.seed(lowerCAmelCase_      )
       lowercase					:   Optional[Any]								= list(lowerCAmelCase_      )

       lowercase					:   Optional[Any]								= [
           IterableDatasetShard(
               lowerCAmelCase_      ,						batch_size=lowerCAmelCase_      ,						drop_last=lowerCAmelCase_      ,						num_processes=lowerCAmelCase_      ,						process_index=lowerCAmelCase_      ,						split_batches=lowerCAmelCase_      ,						)
           for i in range(lowerCAmelCase_      )
       ]
       lowercase					:   Tuple								= []
       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_      )      )

       lowercase					:   Tuple								= 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
       lowercase					:   List[str]								= 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      )

       lowercase					:   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 					__a			(  self		:	Tuple      )		->    List[Any]:






       """simple docstring"""







       lowercase					:   Any								= 42
       lowercase					:   List[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
       lowercase					:   List[Any]								= 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 					__a			(  self		:	Any      )		->    List[Any]:






       """simple docstring"""







       lowercase					:   List[str]								= BatchSampler(range(16      )      ,						batch_size=4      ,						drop_last=lowerCAmelCase_      )
       lowercase					:   Optional[int]								= SkipBatchSampler(lowerCAmelCase_      ,						2      )
       self.assertListEqual(list(lowerCAmelCase_      )      ,						[[8, 9, 10, 11], [12, 13, 14, 15]]      )






  def 					__a			(  self		:	Union[str, Any]      )		->    int:






       """simple docstring"""







       lowercase					:   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 					__a			(  self		:	Tuple      )		->    List[Any]:






       """simple docstring"""







       lowercase					:   int								= DataLoader(list(range(16      )      )      ,						batch_size=4      )
       lowercase					:   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 					__a			(  self		:	Optional[Any]      )		->    int:






       """simple docstring"""







       lowercase					:   Tuple								= 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 					__a			(  self		:	Optional[Any]      )		->    Optional[Any]:






       """simple docstring"""







       Accelerator()
       lowercase					:   Union[str, 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      )

'''simple docstring'''







import tensorflow as tf

from ...tf_utils import shape_list


class 						__UpperCAmelCase      (					tf.keras.layers.Layer						):


							def __init__(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_=1    ,   lowerCAmelCase_=False    ,   **lowerCAmelCase_  ):





								"""simple docstring"""






								super().__init__(**lowerCAmelCase_  )

								_snake_case   						=		vocab_size
								_snake_case   						=		d_embed
								_snake_case   						=		d_proj

								_snake_case   						=		cutoffs + [vocab_size]
								_snake_case   						=		[0] + self.cutoffs
								_snake_case   						=		div_val

								_snake_case   						=		self.cutoffs[0]
								_snake_case   						=		len(self.cutoffs  ) - 1
								_snake_case   						=		self.shortlist_size + self.n_clusters
								_snake_case   						=		keep_order

								_snake_case   						=		[]
								_snake_case   						=		[]


							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								if self.n_clusters > 0:
									_snake_case   						=		self.add_weight(
									    shape=(self.n_clusters, self.d_embed)    ,   initializer='zeros'    ,   trainable=lowerCAmelCase_    ,   name='cluster_weight'  )
									_snake_case   						=		self.add_weight(
									    shape=(self.n_clusters,)    ,   initializer='zeros'    ,   trainable=lowerCAmelCase_    ,   name='cluster_bias'  )

								if self.div_val == 1:
									for i in range(len(self.cutoffs  )  ):
										if self.d_proj != self.d_embed:
											_snake_case   						=		self.add_weight(
											    shape=(self.d_embed, self.d_proj)    ,   initializer='zeros'    ,   trainable=lowerCAmelCase_    ,   name=F'out_projs_._{i}'    ,   )
											self.out_projs.append(lowerCAmelCase_  )
										else:
											self.out_projs.append(lowerCAmelCase_  )
										_snake_case   						=		self.add_weight(
										    shape=(self.vocab_size, self.d_embed)    ,   initializer='zeros'    ,   trainable=lowerCAmelCase_    ,   name=F'out_layers_._{i}_._weight'    ,   )
										_snake_case   						=		self.add_weight(
										    shape=(self.vocab_size,)    ,   initializer='zeros'    ,   trainable=lowerCAmelCase_    ,   name=F'out_layers_._{i}_._bias'    ,   )
										self.out_layers.append((weight, bias)  )
								else:
									for i in range(len(self.cutoffs  )  ):
										_snake_case						, _snake_case   						=		self.cutoff_ends[i], self.cutoff_ends[i + 1]
										_snake_case   						=		self.d_embed // (self.div_val**i)

										_snake_case   						=		self.add_weight(
										    shape=(d_emb_i, self.d_proj)    ,   initializer='zeros'    ,   trainable=lowerCAmelCase_    ,   name=F'out_projs_._{i}'  )
										self.out_projs.append(lowerCAmelCase_  )
										_snake_case   						=		self.add_weight(
										    shape=(r_idx - l_idx, d_emb_i)    ,   initializer='zeros'    ,   trainable=lowerCAmelCase_    ,   name=F'out_layers_._{i}_._weight'    ,   )
										_snake_case   						=		self.add_weight(
										    shape=(r_idx - l_idx,)    ,   initializer='zeros'    ,   trainable=lowerCAmelCase_    ,   name=F'out_layers_._{i}_._bias'    ,   )
										self.out_layers.append((weight, bias)  )
								super().build(lowerCAmelCase_  )


							@staticmethod
							def  lowerCamelCase				(       lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_=None  ):





								"""simple docstring"""






								_snake_case   						=		x
								if proj is not None:
									_snake_case   						=		tf.einsum('ibd,ed->ibe'    ,   lowerCAmelCase_    ,   lowerCAmelCase_  )
								return tf.einsum('ibd,nd->ibn'    ,   lowerCAmelCase_    ,   lowerCAmelCase_  ) + b


							@staticmethod
							def  lowerCamelCase				(       lowerCAmelCase_    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								_snake_case   						=		shape_list(lowerCAmelCase_  )
								_snake_case   						=		tf.range(lp_size[0]    ,   dtype=target.dtype  )
								_snake_case   						=		tf.stack([r, target]    ,   1  )
								return tf.gather_nd(lowerCAmelCase_    ,   lowerCAmelCase_  )




							def  lowerCamelCase				(       self    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_=True    ,   lowerCAmelCase_=False  ):





								"""simple docstring"""






								_snake_case   						=		0
								if self.n_clusters == 0:
									_snake_case   						=		self._logit(lowerCAmelCase_    ,   self.out_layers[0][0]    ,   self.out_layers[0][1]    ,   self.out_projs[0]  )
									if target is not None:
										_snake_case   						=		tf.nn.sparse_softmax_cross_entropy_with_logits(labels=lowerCAmelCase_    ,   logits=lowerCAmelCase_  )
									_snake_case   						=		tf.nn.log_softmax(lowerCAmelCase_    ,   axis=-1  )
								else:
									_snake_case   						=		shape_list(lowerCAmelCase_  )
									_snake_case   						=		[]
									_snake_case   						=		tf.zeros(hidden_sizes[:2]  )
									for i in range(len(self.cutoffs  )  ):
										_snake_case						, _snake_case   						=		self.cutoff_ends[i], self.cutoff_ends[i + 1]
										if target is not None:
											_snake_case   						=		(target >= l_idx) & (target < r_idx)
											_snake_case   						=		tf.where(lowerCAmelCase_  )
											_snake_case   						=		tf.boolean_mask(lowerCAmelCase_    ,   lowerCAmelCase_  ) - l_idx

										if self.div_val == 1:
											_snake_case   						=		self.out_layers[0][0][l_idx:r_idx]
											_snake_case   						=		self.out_layers[0][1][l_idx:r_idx]
										else:
											_snake_case   						=		self.out_layers[i][0]
											_snake_case   						=		self.out_layers[i][1]

										if i == 0:
											_snake_case   						=		tf.concat([cur_W, self.cluster_weight]    ,   0  )
											_snake_case   						=		tf.concat([cur_b, self.cluster_bias]    ,   0  )

											_snake_case   						=		self._logit(lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   self.out_projs[0]  )
											_snake_case   						=		tf.nn.log_softmax(lowerCAmelCase_  )
											out.append(head_logprob[..., : self.cutoffs[0]]  )
											if target is not None:
												_snake_case   						=		tf.boolean_mask(lowerCAmelCase_    ,   lowerCAmelCase_  )
												_snake_case   						=		self._gather_logprob(lowerCAmelCase_    ,   lowerCAmelCase_  )
										else:
											_snake_case   						=		self._logit(lowerCAmelCase_    ,   lowerCAmelCase_    ,   lowerCAmelCase_    ,   self.out_projs[i]  )
											_snake_case   						=		tf.nn.log_softmax(lowerCAmelCase_  )
											_snake_case   						=		self.cutoffs[0] + i - 1  # No probability for the head cluster
											_snake_case   						=		head_logprob[..., cluster_prob_idx, None] + tail_logprob
											out.append(lowerCAmelCase_  )
											if target is not None:
												_snake_case   						=		tf.boolean_mask(lowerCAmelCase_    ,   lowerCAmelCase_  )
												_snake_case   						=		tf.boolean_mask(lowerCAmelCase_    ,   lowerCAmelCase_  )
												_snake_case   						=		self._gather_logprob(lowerCAmelCase_    ,   lowerCAmelCase_  )
												cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
										if target is not None:
											loss += tf.scatter_nd(lowerCAmelCase_    ,   -cur_logprob    ,   shape_list(lowerCAmelCase_  )  )
									_snake_case   						=		tf.concat(lowerCAmelCase_    ,   axis=-1  )

								if target is not None:
									if return_mean:
										_snake_case   						=		tf.reduce_mean(lowerCAmelCase_  )
									# Add the training-time loss value to the layer using `self.add_loss()`.
									self.add_loss(lowerCAmelCase_  )

									# Log the loss as a metric (we could log arbitrary metrics,
									# including different metrics for training and inference.
									self.add_metric(lowerCAmelCase_    ,   name=self.name    ,   aggregation='mean' if return_mean else ''  )

								return out

