






import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import numpy as np
import pytest

from datasets.arrow_dataset import Dataset
from datasets.search import ElasticSearchIndex, FaissIndex, MissingIndex

from .utils import require_elasticsearch, require_faiss


lowerCamelCase__       		=  pytest.mark.integration




@require_faiss
class       snake_case__ (				lowercase_):

       '''simple docstring'''



       def        __lowercase    (				self     )			->   Dict:






            '''simple docstring'''





            __snake_case      :Tuple 							=    Dataset.from_dict({"""filename""": ["""my_name-train""" + """_""" + str(a__     ) for x in np.arange(30     ).tolist()]}     )
            return dset



       def        __lowercase    (				self     )			->   List[str]:






            '''simple docstring'''





            import faiss

            __snake_case      :Dataset 							=    self._create_dummy_dataset()
            __snake_case      :Optional[int] 							=    dset.map(
                lambda a__      ,    a__     : {"vecs": i * np.ones(5      ,    dtype=np.floataa     )}      ,    with_indices=a__      ,    keep_in_memory=a__     )
            __snake_case      :str 							=    dset.add_faiss_index("""vecs"""      ,    batch_size=1_00      ,    metric_type=faiss.METRIC_INNER_PRODUCT     )
            __snake_case							,			__snake_case      :Optional[Any] 							=    dset.get_nearest_examples("""vecs"""      ,    np.ones(5      ,    dtype=np.floataa     )     )
            self.assertEqual(examples["""filename"""][0]      ,    """my_name-train_29"""     )
            dset.drop_index("""vecs"""     )



       def        __lowercase    (				self     )			->   int:






            '''simple docstring'''





            import faiss

            __snake_case      :Dataset 							=    self._create_dummy_dataset()
            dset.add_faiss_index_from_external_arrays(
                external_arrays=np.ones((30, 5)     ) * np.arange(30     ).reshape(-1      ,    1     )      ,    index_name="""vecs"""      ,    batch_size=1_00      ,    metric_type=faiss.METRIC_INNER_PRODUCT      ,    )
            __snake_case							,			__snake_case      :Optional[Any] 							=    dset.get_nearest_examples("""vecs"""      ,    np.ones(5      ,    dtype=np.floataa     )     )
            self.assertEqual(examples["""filename"""][0]      ,    """my_name-train_29"""     )



       def        __lowercase    (				self     )			->   Tuple:






            '''simple docstring'''





            import faiss

            __snake_case      :Dataset 							=    self._create_dummy_dataset()
            dset.add_faiss_index_from_external_arrays(
                external_arrays=np.ones((30, 5)     ) * np.arange(30     ).reshape(-1      ,    1     )      ,    index_name="""vecs"""      ,    metric_type=faiss.METRIC_INNER_PRODUCT      ,    )

            # Setting delete=False and unlinking manually is not pretty... but it is required on Windows to
            # ensure somewhat stable behaviour. If we don't, we get PermissionErrors. This is an age-old issue.
            # see https://bugs.python.org/issue14243 and
            # https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file/23212515
            with tempfile.NamedTemporaryFile(delete=a__     ) as tmp_file:
                 dset.save_faiss_index("""vecs"""      ,    tmp_file.name     )
                 dset.load_faiss_index("""vecs2"""      ,    tmp_file.name     )
            os.unlink(tmp_file.name     )

            __snake_case							,			__snake_case      :List[Any] 							=    dset.get_nearest_examples("""vecs2"""      ,    np.ones(5      ,    dtype=np.floataa     )     )
            self.assertEqual(examples["""filename"""][0]      ,    """my_name-train_29"""     )



       def        __lowercase    (				self     )			->   Dict:






            '''simple docstring'''





            __snake_case      :Dataset 							=    self._create_dummy_dataset()
            dset.add_faiss_index_from_external_arrays(
                external_arrays=np.ones((30, 5)     ) * np.arange(30     ).reshape(-1      ,    1     )      ,    index_name="""vecs"""     )
            dset.drop_index("""vecs"""     )
            self.assertRaises(a__      ,    partial(dset.get_nearest_examples      ,    """vecs2"""      ,    np.ones(5      ,    dtype=np.floataa     )     )     )



       def        __lowercase    (				self     )			->   Any:






            '''simple docstring'''





            from elasticsearch import Elasticsearch

            __snake_case      :Dataset 							=    self._create_dummy_dataset()
            with patch("""elasticsearch.Elasticsearch.search"""     ) as mocked_search, patch(
                """elasticsearch.client.IndicesClient.create"""     ) as mocked_index_create, patch("""elasticsearch.helpers.streaming_bulk"""     ) as mocked_bulk:
                 __snake_case      :Any 							=    {"""acknowledged""": True}
                 mocked_bulk.return_value([(True, None)] * 30     )
                 __snake_case      :Dict 							=    {"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 29}]}}
                 __snake_case      :Optional[Any] 							=    Elasticsearch()

                 dset.add_elasticsearch_index("""filename"""      ,    es_client=a__     )
                 __snake_case							,			__snake_case      :Tuple 							=    dset.get_nearest_examples("""filename"""      ,    """my_name-train_29"""     )
                 self.assertEqual(examples["""filename"""][0]      ,    """my_name-train_29"""     )




@require_faiss
class       snake_case__ (				lowercase_):

       '''simple docstring'''



       def        __lowercase    (				self     )			->   List[Any]:






            '''simple docstring'''





            import faiss

            __snake_case      :Tuple 							=    FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT     )

            # add vectors
            index.add_vectors(np.eye(5      ,    dtype=np.floataa     )     )
            self.assertIsNotNone(index.faiss_index     )
            self.assertEqual(index.faiss_index.ntotal      ,    5     )
            index.add_vectors(np.zeros((5, 5)      ,    dtype=np.floataa     )     )
            self.assertEqual(index.faiss_index.ntotal      ,    10     )

            # single query
            __snake_case      :str 							=    np.zeros(5      ,    dtype=np.floataa     )
            __snake_case      :Tuple 							=    1
            __snake_case							,			__snake_case      :str 							=    index.search(a__     )
            self.assertRaises(a__      ,    index.search      ,    query.reshape(-1      ,    1     )     )
            self.assertGreater(scores[0]      ,    0     )
            self.assertEqual(indices[0]      ,    1     )

            # batched queries
            __snake_case      :List[str] 							=    np.eye(5      ,    dtype=np.floataa     )[::-1]
            __snake_case							,			__snake_case      :Optional[int] 							=    index.search_batch(a__     )
            self.assertRaises(a__      ,    index.search_batch      ,    queries[0]     )
            __snake_case      :Dict 							=    [scores[0] for scores in total_scores]
            __snake_case      :Optional[Any] 							=    [indices[0] for indices in total_indices]
            self.assertGreater(np.min(a__     )      ,    0     )
            self.assertListEqual([4, 3, 2, 1, 0]      ,    a__     )



       def        __lowercase    (				self     )			->   List[str]:






            '''simple docstring'''





            import faiss

            __snake_case      :Union[str, Any] 							=    FaissIndex(string_factory="""Flat"""     )
            index.add_vectors(np.eye(5      ,    dtype=np.floataa     )     )
            self.assertIsInstance(index.faiss_index      ,    faiss.IndexFlat     )
            __snake_case      :Tuple 							=    FaissIndex(string_factory="""LSH"""     )
            index.add_vectors(np.eye(5      ,    dtype=np.floataa     )     )
            self.assertIsInstance(index.faiss_index      ,    faiss.IndexLSH     )
            with self.assertRaises(a__     ):
                 __snake_case      :List[Any] 							=    FaissIndex(string_factory="""Flat"""      ,    custom_index=faiss.IndexFlat(5     )     )



       def        __lowercase    (				self     )			->   Tuple:






            '''simple docstring'''





            import faiss

            __snake_case      :Union[str, Any] 							=    faiss.IndexFlat(5     )
            __snake_case      :List[Any] 							=    FaissIndex(custom_index=a__     )
            index.add_vectors(np.eye(5      ,    dtype=np.floataa     )     )
            self.assertIsInstance(index.faiss_index      ,    faiss.IndexFlat     )



       def        __lowercase    (				self     )			->   Any:






            '''simple docstring'''





            import faiss

            __snake_case      :Optional[int] 							=    FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT     )
            index.add_vectors(np.eye(5      ,    dtype=np.floataa     )     )

            # Setting delete=False and unlinking manually is not pretty... but it is required on Windows to
            # ensure somewhat stable behaviour. If we don't, we get PermissionErrors. This is an age-old issue.
            # see https://bugs.python.org/issue14243 and
            # https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file/23212515
            with tempfile.NamedTemporaryFile(delete=a__     ) as tmp_file:
                 index.save(tmp_file.name     )
                 __snake_case      :int 							=    FaissIndex.load(tmp_file.name     )
            os.unlink(tmp_file.name     )

            __snake_case      :Tuple 							=    np.zeros(5      ,    dtype=np.floataa     )
            __snake_case      :List[Any] 							=    1
            __snake_case							,			__snake_case      :str 							=    index.search(a__     )
            self.assertGreater(scores[0]      ,    0     )
            self.assertEqual(indices[0]      ,    1     )






@require_faiss
def 	UpperCamelCase	(							snake_case__  : Union[str, Any]	):




     '''simple docstring'''






     import faiss

     __snake_case      :str 							=    FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT	)
     index.add_vectors(np.eye(5						,dtype=np.floataa	)	)

     __snake_case      :List[str] 							=    """index.faiss"""
     __snake_case      :List[str] 							=    f'''mock://{index_name}'''
     index.save(snake_case__						,storage_options=mockfs.storage_options	)
     __snake_case      :Dict 							=    FaissIndex.load(snake_case__						,storage_options=mockfs.storage_options	)

     __snake_case      :Any 							=    np.zeros(5						,dtype=np.floataa	)
     __snake_case      :List[str] 							=    1
     __snake_case							,			__snake_case      :List[str] 							=    index.search(snake_case__	)
     assert scores[0] > 0
     assert indices[0] == 1




@require_elasticsearch
class       snake_case__ (				lowercase_):

       '''simple docstring'''



       def        __lowercase    (				self     )			->   List[Any]:






            '''simple docstring'''





            from elasticsearch import Elasticsearch

            with patch("""elasticsearch.Elasticsearch.search"""     ) as mocked_search, patch(
                """elasticsearch.client.IndicesClient.create"""     ) as mocked_index_create, patch("""elasticsearch.helpers.streaming_bulk"""     ) as mocked_bulk:
                 __snake_case      :Union[str, Any] 							=    Elasticsearch()
                 __snake_case      :Tuple 							=    {"""acknowledged""": True}
                 __snake_case      :Union[str, Any] 							=    ElasticSearchIndex(es_client=a__     )
                 mocked_bulk.return_value([(True, None)] * 3     )
                 index.add_documents(["""foo""", """bar""", """foobar"""]     )

                 # single query
                 __snake_case      :Union[str, Any] 							=    """foo"""
                 __snake_case      :int 							=    {"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 0}]}}
                 __snake_case							,			__snake_case      :Any 							=    index.search(a__     )
                 self.assertEqual(scores[0]      ,    1     )
                 self.assertEqual(indices[0]      ,    0     )

                 # single query with timeout
                 __snake_case      :Dict 							=    """foo"""
                 __snake_case      :Any 							=    {"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 0}]}}
                 __snake_case							,			__snake_case      :Dict 							=    index.search(a__      ,    request_timeout=30     )
                 self.assertEqual(scores[0]      ,    1     )
                 self.assertEqual(indices[0]      ,    0     )

                 # batched queries
                 __snake_case      :Tuple 							=    ["""foo""", """bar""", """foobar"""]
                 __snake_case      :List[str] 							=    {"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 1}]}}
                 __snake_case							,			__snake_case      :str 							=    index.search_batch(a__     )
                 __snake_case      :Optional[int] 							=    [scores[0] for scores in total_scores]
                 __snake_case      :List[Any] 							=    [indices[0] for indices in total_indices]
                 self.assertGreater(np.min(a__     )      ,    0     )
                 self.assertListEqual([1, 1, 1]      ,    a__     )

                 # batched queries with timeout
                 __snake_case      :Optional[int] 							=    ["""foo""", """bar""", """foobar"""]
                 __snake_case      :Optional[int] 							=    {"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 1}]}}
                 __snake_case							,			__snake_case      :Optional[Any] 							=    index.search_batch(a__      ,    request_timeout=30     )
                 __snake_case      :Union[str, Any] 							=    [scores[0] for scores in total_scores]
                 __snake_case      :List[Any] 							=    [indices[0] for indices in total_indices]
                 self.assertGreater(np.min(a__     )      ,    0     )
                 self.assertListEqual([1, 1, 1]      ,    a__     )










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       snake_case__ (				lowercase_):

       '''simple docstring'''



       def __init__(				self      ,    a__=0.01      ,    a__=10_00     )			->   List[Any]:






            '''simple docstring'''





            __snake_case      :int 							=    p_stop
            __snake_case      :List[Any] 							=    max_length



       def __iter__(				self     )			->   Optional[int]:






            '''simple docstring'''





            __snake_case      :str 							=    0
            __snake_case      :Optional[Any] 							=    False
            while not stop and count < self.max_length:
                 yield count
                 count += 1
                 __snake_case      :str 							=    random.random() < self.p_stop




class       snake_case__ (				unittest.TestCase):

       '''simple docstring'''



       def        __lowercase    (				self      ,    a__      ,    a__      ,    a__=False      ,    a__=True     )			->   List[Any]:






            '''simple docstring'''





            __snake_case      :Optional[Any] 							=    [
                BatchSamplerShard(a__      ,    2      ,    a__      ,    split_batches=a__      ,    even_batches=a__     )
                for i in range(2     )
            ]
            __snake_case      :Union[str, Any] 							=    [list(a__     ) for batch_sampler_shard in batch_sampler_shards]
            if not split_batches:
                 self.assertListEqual([len(a__     ) for shard in batch_sampler_shards]      ,    [len(a__     ) for e in expected]     )
            self.assertListEqual(a__      ,    a__     )



       def        __lowercase    (				self     )			->   Any:






            '''simple docstring'''





            __snake_case      :Optional[Any] 							=    BatchSampler(range(24     )      ,    batch_size=3      ,    drop_last=a__     )
            __snake_case      :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(a__      ,    a__     )

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

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

            __snake_case      :List[Any] 							=    BatchSampler(range(21     )      ,    batch_size=3      ,    drop_last=a__     )
            __snake_case      :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(a__      ,    a__     )

            # Check the shards when the dataset is not a round multiple of batch size but has a multiple of
            # num_processes batch.
            __snake_case      :Optional[Any] 							=    BatchSampler(range(22     )      ,    batch_size=3      ,    drop_last=a__     )
            __snake_case      :List[str] 							=    [
                [[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(a__      ,    a__     )

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

            # 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      :List[str] 							=    BatchSampler(range(20     )      ,    batch_size=3      ,    drop_last=a__     )
            __snake_case      :int 							=    [
                [[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(a__      ,    a__     )

            __snake_case      :Optional[Any] 							=    BatchSampler(range(20     )      ,    batch_size=3      ,    drop_last=a__     )
            __snake_case      :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(a__      ,    a__     )

            # Check the shards when the dataset is very small.
            __snake_case      :Dict 							=    BatchSampler(range(2     )      ,    batch_size=3      ,    drop_last=a__     )
            __snake_case      :Optional[Any] 							=    [[[0, 1, 0]], [[1, 0, 1]]]
            self.check_batch_sampler_shards(a__      ,    a__     )

            __snake_case      :Union[str, Any] 							=    BatchSampler(range(2     )      ,    batch_size=3      ,    drop_last=a__     )
            __snake_case      :List[str] 							=    [[], []]
            self.check_batch_sampler_shards(a__      ,    a__     )



       def        __lowercase    (				self     )			->   Dict:






            '''simple docstring'''





            __snake_case      :Optional[int] 							=    BatchSampler(range(24     )      ,    batch_size=4      ,    drop_last=a__     )
            __snake_case      :Any 							=    [
                [[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(a__      ,    a__      ,    split_batches=a__     )

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

            # Check the shards when the dataset is not a round multiple of batch size.
            __snake_case      :List[Any] 							=    BatchSampler(range(22     )      ,    batch_size=4      ,    drop_last=a__     )
            __snake_case      :Optional[int] 							=    [
                [[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(a__      ,    a__      ,    split_batches=a__     )

            __snake_case      :List[Any] 							=    BatchSampler(range(22     )      ,    batch_size=4      ,    drop_last=a__     )
            __snake_case      :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(a__      ,    a__      ,    split_batches=a__     )

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

            __snake_case      :Any 							=    BatchSampler(range(21     )      ,    batch_size=4      ,    drop_last=a__     )
            __snake_case      :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(a__      ,    a__      ,    split_batches=a__     )

            # Check the shards when the dataset is very small.
            __snake_case      :str 							=    BatchSampler(range(2     )      ,    batch_size=4      ,    drop_last=a__     )
            __snake_case      :Dict 							=    [[[0, 1]], [[0, 1]]]
            self.check_batch_sampler_shards(a__      ,    a__      ,    split_batches=a__     )

            __snake_case      :Optional[Any] 							=    BatchSampler(range(2     )      ,    batch_size=4      ,    drop_last=a__     )
            __snake_case      :List[Any] 							=    [[], []]
            self.check_batch_sampler_shards(a__      ,    a__      ,    split_batches=a__     )



       def        __lowercase    (				self     )			->   List[Any]:






            '''simple docstring'''





            __snake_case      :Tuple 							=    BatchSampler(range(24     )      ,    batch_size=3      ,    drop_last=a__     )
            __snake_case      :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(a__      ,    a__      ,    even_batches=a__     )

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

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

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

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

            __snake_case      :Tuple 							=    BatchSampler(range(22     )      ,    batch_size=3      ,    drop_last=a__     )
            __snake_case      :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(a__      ,    a__      ,    even_batches=a__     )

            # 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      :Dict 							=    BatchSampler(range(20     )      ,    batch_size=3      ,    drop_last=a__     )
            __snake_case      :int 							=    [
                [[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(a__      ,    a__      ,    even_batches=a__     )

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

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

            __snake_case      :List[str] 							=    BatchSampler(range(2     )      ,    batch_size=3      ,    drop_last=a__     )
            __snake_case      :int 							=    [[], []]
            self.check_batch_sampler_shards(a__      ,    a__      ,    even_batches=a__     )



       def        __lowercase    (				self     )			->   Dict:






            '''simple docstring'''





            __snake_case      :int 							=    BatchSampler(range(24     )      ,    batch_size=4      ,    drop_last=a__     )
            __snake_case      :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(a__      ,    a__      ,    split_batches=a__      ,    even_batches=a__     )

            __snake_case      :Dict 							=    BatchSampler(range(24     )      ,    batch_size=4      ,    drop_last=a__     )
            # Expected shouldn't change
            self.check_batch_sampler_shards(a__      ,    a__      ,    split_batches=a__      ,    even_batches=a__     )

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

            __snake_case      :Dict 							=    BatchSampler(range(22     )      ,    batch_size=4      ,    drop_last=a__     )
            __snake_case      :Dict 							=    [
                [[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(a__      ,    a__      ,    split_batches=a__      ,    even_batches=a__     )

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

            __snake_case      :int 							=    BatchSampler(range(21     )      ,    batch_size=4      ,    drop_last=a__     )
            __snake_case      :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(a__      ,    a__      ,    split_batches=a__      ,    even_batches=a__     )

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

            __snake_case      :Any 							=    BatchSampler(range(2     )      ,    batch_size=4      ,    drop_last=a__     )
            __snake_case      :List[str] 							=    [[], []]
            self.check_batch_sampler_shards(a__      ,    a__      ,    split_batches=a__      ,    even_batches=a__     )



       def        __lowercase    (				self     )			->   Optional[Any]:






            '''simple docstring'''





            __snake_case      :str 							=    [[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
            __snake_case      :List[str] 							=    [BatchSamplerShard(a__      ,    2      ,    a__      ,    even_batches=a__     ) 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        __lowercase    (				self      ,    a__      ,    a__      ,    a__      ,    a__=False      ,    a__=2      ,    a__=False     )			->   List[str]:






            '''simple docstring'''





            random.seed(a__     )
            __snake_case      :Optional[int] 							=    list(a__     )

            __snake_case      :Union[str, Any] 							=    [
                IterableDatasetShard(
                    a__      ,    batch_size=a__      ,    drop_last=a__      ,    num_processes=a__      ,    process_index=a__      ,    split_batches=a__      ,    )
                for i in range(a__     )
            ]
            __snake_case      :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(a__     )
                 iterable_dataset_lists.append(list(a__     )     )

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

            __snake_case      :int 							=    []
            for idx in range(0      ,    len(a__     )      ,    a__     ):
                 for l in iterable_dataset_lists:
                      observed += l[idx : idx + shard_batch_size]

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



       def        __lowercase    (				self     )			->   Any:






            '''simple docstring'''





            __snake_case      :int 							=    42
            __snake_case      :Tuple 							=    RandomIterableDataset()

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

            # Edge case with a very small dataset
            __snake_case      :Optional[int] 							=    RandomIterableDataset(max_length=2     )

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



       def        __lowercase    (				self     )			->   Dict:






            '''simple docstring'''





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



       def        __lowercase    (				self     )			->   str:






            '''simple docstring'''





            __snake_case      :str 							=    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        __lowercase    (				self     )			->   Dict:






            '''simple docstring'''





            __snake_case      :Union[str, Any] 							=    DataLoader(list(range(16     )     )      ,    batch_size=4     )
            __snake_case      :Dict 							=    skip_first_batches(a__      ,    num_batches=2     )
            self.assertListEqual([t.tolist() for t in new_dataloader]      ,    [[8, 9, 10, 11], [12, 13, 14, 15]]     )



       def        __lowercase    (				self     )			->   Optional[Any]:






            '''simple docstring'''





            __snake_case      :Any 							=    DataLoaderShard(list(range(16     )     )      ,    batch_size=4     )
            for idx, _ in enumerate(a__     ):
                 self.assertEqual(dataloader.end_of_dataloader      ,    idx == 3     )

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



       def        __lowercase    (				self     )			->   Any:






            '''simple docstring'''





            Accelerator()
            __snake_case      :Union[str, Any] 							=    DataLoaderDispatcher(range(16     )      ,    batch_size=4     )
            for idx, _ in enumerate(a__     ):
                 self.assertEqual(dataloader.end_of_dataloader      ,    idx == 3     )

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


