


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 lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 def __init__(			self	,   _UpperCAmelCase=0.01	,   _UpperCAmelCase=10_00      ):
       snake_case_ =   p_stop
       snake_case_ =   max_length





 def __iter__(			self      ):
       snake_case_ =   0
       snake_case_ =   False
       while not stop and count < self.max_length:
             yield count
             count += 1
             snake_case_ =   random.random() < self.p_stop







class lowerCAmelCase_		(				unittest.TestCase ):

 '''simple docstring'''

 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase=False	,   _UpperCAmelCase=True      ):
       snake_case_ =   [
           BatchSamplerShard(_UpperCAmelCase	,   2	,   _UpperCAmelCase	,   split_batches=_UpperCAmelCase	,   even_batches=_UpperCAmelCase      )
           for i in range(2      )
       ]
       snake_case_ =   [list(_UpperCAmelCase      ) for batch_sampler_shard in batch_sampler_shards]
       if not split_batches:
             self.assertListEqual([len(_UpperCAmelCase      ) for shard in batch_sampler_shards]	,   [len(_UpperCAmelCase      ) for e in expected]      )
       self.assertListEqual(_UpperCAmelCase	,   _UpperCAmelCase      )



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

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

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

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

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

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

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

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

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

       snake_case_ =   BatchSampler(range(2      )	,   batch_size=3	,   drop_last=_UpperCAmelCase      )
       snake_case_ =   [[], []]
       self.check_batch_sampler_shards(_UpperCAmelCase	,   _UpperCAmelCase      )



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

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

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

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

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

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

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

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



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

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

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

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

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

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

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

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

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

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



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

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

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

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

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

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

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

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



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
       snake_case_ =   [BatchSamplerShard(_UpperCAmelCase	,   2	,   _UpperCAmelCase	,   even_batches=_UpperCAmelCase      ) 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	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase=False	,   _UpperCAmelCase=2	,   _UpperCAmelCase=False      ):
       random.seed(_UpperCAmelCase      )
       snake_case_ =   list(_UpperCAmelCase      )

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

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

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

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



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   42
       snake_case_ =   RandomIterableDataset()

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

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

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



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



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



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

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





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

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



from typing import List, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase    =					logging.get_logger(__name__)

UpperCAmelCase    =					{
    """huggingface/informer-tourism-monthly""": (
        """https://huggingface.co/huggingface/informer-tourism-monthly/resolve/main/config.json"""
    ),
    # See all Informer models at https://huggingface.co/models?filter=informer
}







class lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 __snake_case				=						"informer"
 __snake_case				=						{
     "hidden_size": "d_model",
     "num_attention_heads": "encoder_attention_heads",
     "num_hidden_layers": "encoder_layers",
 }



 def __init__(			self	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = "student_t"	,   _UpperCAmelCase = "nll"	,   _UpperCAmelCase = 1	,   _UpperCAmelCase = None	,   _UpperCAmelCase = "mean"	,   _UpperCAmelCase = 0	,   _UpperCAmelCase = 0	,   _UpperCAmelCase = 0	,   _UpperCAmelCase = 0	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = 64	,   _UpperCAmelCase = 32	,   _UpperCAmelCase = 32	,   _UpperCAmelCase = 2	,   _UpperCAmelCase = 2	,   _UpperCAmelCase = 2	,   _UpperCAmelCase = 2	,   _UpperCAmelCase = True	,   _UpperCAmelCase = "gelu"	,   _UpperCAmelCase = 0.05	,   _UpperCAmelCase = 0.1	,   _UpperCAmelCase = 0.1	,   _UpperCAmelCase = 0.1	,   _UpperCAmelCase = 0.1	,   _UpperCAmelCase = 1_00	,   _UpperCAmelCase = 0.02	,   _UpperCAmelCase=True	,   _UpperCAmelCase = "prob"	,   _UpperCAmelCase = 5	,   _UpperCAmelCase = True	,   **_UpperCAmelCase	,   ):
       # time series specific configuration
       snake_case_ =   prediction_length
       snake_case_ =   context_length or prediction_length
       snake_case_ =   distribution_output
       snake_case_ =   loss
       snake_case_ =   input_size
       snake_case_ =   num_time_features
       snake_case_ =   lags_sequence if lags_sequence is not None else [1, 2, 3, 4, 5, 6, 7]
       snake_case_ =   scaling
       snake_case_ =   num_dynamic_real_features
       snake_case_ =   num_static_real_features
       snake_case_ =   num_static_categorical_features

       # set cardinality
       if cardinality and num_static_categorical_features > 0:
             if len(_UpperCAmelCase      ) != num_static_categorical_features:
                   raise ValueError(
                       '''The cardinality should be a list of the same length as `num_static_categorical_features`'''      )
             snake_case_ =   cardinality
       else:
             snake_case_ =   [0]

       # set embedding_dimension
       if embedding_dimension and num_static_categorical_features > 0:
             if len(_UpperCAmelCase      ) != num_static_categorical_features:
                   raise ValueError(
                       '''The embedding dimension should be a list of the same length as `num_static_categorical_features`'''      )
             snake_case_ =   embedding_dimension
       else:
             snake_case_ =   [min(50	,   (cat + 1) // 2      ) for cat in self.cardinality]

       snake_case_ =   num_parallel_samples

       # Transformer architecture configuration
       snake_case_ =   input_size * len(self.lags_sequence      ) + self._number_of_features
       snake_case_ =   d_model
       snake_case_ =   encoder_attention_heads
       snake_case_ =   decoder_attention_heads
       snake_case_ =   encoder_ffn_dim
       snake_case_ =   decoder_ffn_dim
       snake_case_ =   encoder_layers
       snake_case_ =   decoder_layers

       snake_case_ =   dropout
       snake_case_ =   attention_dropout
       snake_case_ =   activation_dropout
       snake_case_ =   encoder_layerdrop
       snake_case_ =   decoder_layerdrop

       snake_case_ =   activation_function
       snake_case_ =   init_std

       snake_case_ =   use_cache

       # Informer
       snake_case_ =   attention_type
       snake_case_ =   sampling_factor
       snake_case_ =   distil

       super().__init__(is_encoder_decoder=_UpperCAmelCase	,   **_UpperCAmelCase      )





 @property
 def 		UpperCamelCase__  (			self      ):
       return (
           sum(self.embedding_dimension      )
           + self.num_dynamic_real_features
           + self.num_time_features
           + self.num_static_real_features
           + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
       )