






'''simple docstring'''

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






 """simple docstring"""





 def __init__(							self			,     __lowerCAmelCase=0.01			,     __lowerCAmelCase=10_00 ):
       UpperCamelCase_				:   Any					    =							p_stop
       UpperCamelCase_				:   Dict					    =							max_length





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






class     A    (      unittest.TestCase						):






 """simple docstring"""





 def   _UpperCAmelCase       (							self			,     __lowerCAmelCase			,     __lowerCAmelCase			,     __lowerCAmelCase=False			,     __lowerCAmelCase=True ):
       UpperCamelCase_				:   str					    =							[
           BatchSamplerShard(__lowerCAmelCase			,     2			,     __lowerCAmelCase			,     split_batches=__lowerCAmelCase			,     even_batches=__lowerCAmelCase )
           for i in range(2 )
       ]
       UpperCamelCase_				:   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   _UpperCAmelCase       (							self ):
       # Check the shards when the dataset is a round multiple of total batch size.
       UpperCamelCase_				:   Optional[Any]					    =							BatchSampler(range(24 )			,     batch_size=3			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   Optional[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 )

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

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

       UpperCamelCase_				:   Optional[Any]					    =							BatchSampler(range(21 )			,     batch_size=3			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   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 )

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

       UpperCamelCase_				:   Optional[Any]					    =							BatchSampler(range(22 )			,     batch_size=3			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   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 )

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

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

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

       UpperCamelCase_				:   int					    =							BatchSampler(range(2 )			,     batch_size=3			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   Dict					    =							[[], []]
       self.check_batch_sampler_shards(__lowerCAmelCase			,     __lowerCAmelCase )
 def   _UpperCAmelCase       (							self ):
       # Check the shards when the dataset is a round multiple of batch size.
       UpperCamelCase_				:   Union[str, Any]					    =							BatchSampler(range(24 )			,     batch_size=4			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   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 )

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

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

       UpperCamelCase_				:   Optional[Any]					    =							BatchSampler(range(22 )			,     batch_size=4			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   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 not a round multiple of batch size or num_processes.
       UpperCamelCase_				:   Any					    =							BatchSampler(range(21 )			,     batch_size=4			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   Tuple					    =							[
           [[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 )

       UpperCamelCase_				:   Any					    =							BatchSampler(range(21 )			,     batch_size=4			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   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.
       UpperCamelCase_				:   Tuple					    =							BatchSampler(range(2 )			,     batch_size=4			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   List[str]					    =							[[[0, 1]], [[0, 1]]]
       self.check_batch_sampler_shards(__lowerCAmelCase			,     __lowerCAmelCase			,     split_batches=__lowerCAmelCase )

       UpperCamelCase_				:   Optional[int]					    =							BatchSampler(range(2 )			,     batch_size=4			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   str					    =							[[], []]
       self.check_batch_sampler_shards(__lowerCAmelCase			,     __lowerCAmelCase			,     split_batches=__lowerCAmelCase )
 def   _UpperCAmelCase       (							self ):
       # Check the shards when the dataset is a round multiple of total batch size.
       UpperCamelCase_				:   Union[str, Any]					    =							BatchSampler(range(24 )			,     batch_size=3			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   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 )

       UpperCamelCase_				:   int					    =							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.
       UpperCamelCase_				:   Optional[int]					    =							BatchSampler(range(21 )			,     batch_size=3			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   Optional[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 )

       UpperCamelCase_				:   str					    =							BatchSampler(range(21 )			,     batch_size=3			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   Dict					    =							[
           [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
           [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
       ]
       self.check_batch_sampler_shards(__lowerCAmelCase			,     __lowerCAmelCase			,     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.
       UpperCamelCase_				:   Any					    =							BatchSampler(range(22 )			,     batch_size=3			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   List[str]					    =							[
           [[0, 1, 2], [6, 7, 8], [12, 13, 14], [18, 19, 20]],
           [[3, 4, 5], [9, 10, 11], [15, 16, 17], [21]],
       ]
       self.check_batch_sampler_shards(__lowerCAmelCase			,     __lowerCAmelCase			,     even_batches=__lowerCAmelCase )

       UpperCamelCase_				:   int					    =							BatchSampler(range(22 )			,     batch_size=3			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   Dict					    =							[
           [[0, 1, 2], [6, 7, 8], [12, 13, 14]],
           [[3, 4, 5], [9, 10, 11], [15, 16, 17]],
       ]
       self.check_batch_sampler_shards(__lowerCAmelCase			,     __lowerCAmelCase			,     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.
       UpperCamelCase_				:   str					    =							BatchSampler(range(20 )			,     batch_size=3			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   Optional[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(__lowerCAmelCase			,     __lowerCAmelCase			,     even_batches=__lowerCAmelCase )

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

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

       UpperCamelCase_				:   Optional[int]					    =							BatchSampler(range(2 )			,     batch_size=3			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   List[Any]					    =							[[], []]
       self.check_batch_sampler_shards(__lowerCAmelCase			,     __lowerCAmelCase			,     even_batches=__lowerCAmelCase )
 def   _UpperCAmelCase       (							self ):
       # Check the shards when the dataset is a round multiple of batch size.
       UpperCamelCase_				:   Dict					    =							BatchSampler(range(24 )			,     batch_size=4			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   List[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(__lowerCAmelCase			,     __lowerCAmelCase			,     split_batches=__lowerCAmelCase			,     even_batches=__lowerCAmelCase )

       UpperCamelCase_				:   List[str]					    =							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.
       UpperCamelCase_				:   Union[str, Any]					    =							BatchSampler(range(22 )			,     batch_size=4			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   Union[str, 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 )

       UpperCamelCase_				:   List[str]					    =							BatchSampler(range(22 )			,     batch_size=4			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   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 not a round multiple of batch size or num_processes.
       UpperCamelCase_				:   Optional[Any]					    =							BatchSampler(range(21 )			,     batch_size=4			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   Optional[Any]					    =							[
           [[0, 1], [4, 5], [8, 9], [12, 13], [16, 17], [20]],
           [[2, 3], [6, 7], [10, 11], [14, 15], [18, 19]],
       ]
       self.check_batch_sampler_shards(__lowerCAmelCase			,     __lowerCAmelCase			,     split_batches=__lowerCAmelCase			,     even_batches=__lowerCAmelCase )

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

       UpperCamelCase_				:   int					    =							BatchSampler(range(2 )			,     batch_size=4			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   int					    =							[[], []]
       self.check_batch_sampler_shards(__lowerCAmelCase			,     __lowerCAmelCase			,     split_batches=__lowerCAmelCase			,     even_batches=__lowerCAmelCase )
 def   _UpperCAmelCase       (							self ):
       UpperCamelCase_				:   Optional[int]					    =							[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 10, 11], [12, 13]]
       UpperCamelCase_				:   List[Any]					    =							[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   _UpperCAmelCase       (							self			,     __lowerCAmelCase			,     __lowerCAmelCase			,     __lowerCAmelCase			,     __lowerCAmelCase=False			,     __lowerCAmelCase=2			,     __lowerCAmelCase=False ):
       random.seed(__lowerCAmelCase )
       UpperCamelCase_				:   Tuple					    =							list(__lowerCAmelCase )

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

       UpperCamelCase_				:   Optional[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
       UpperCamelCase_				:   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 )

       UpperCamelCase_				:   List[str]					    =							[]
       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   _UpperCAmelCase       (							self ):
       UpperCamelCase_				:   Union[str, Any]					    =							42
       UpperCamelCase_				:   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
       UpperCamelCase_				:   Optional[int]					    =							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   _UpperCAmelCase       (							self ):
       UpperCamelCase_				:   Union[str, Any]					    =							BatchSampler(range(16 )			,     batch_size=4			,     drop_last=__lowerCAmelCase )
       UpperCamelCase_				:   List[Any]					    =							SkipBatchSampler(__lowerCAmelCase			,     2 )
       self.assertListEqual(list(__lowerCAmelCase )			,     [[8, 9, 10, 11], [12, 13, 14, 15]] )
 def   _UpperCAmelCase       (							self ):
       UpperCamelCase_				:   List[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   _UpperCAmelCase       (							self ):
       UpperCamelCase_				:   Union[str, Any]					    =							DataLoader(list(range(16 ) )			,     batch_size=4 )
       UpperCamelCase_				:   str					    =							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   _UpperCAmelCase       (							self ):
       UpperCamelCase_				:   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   _UpperCAmelCase       (							self ):
       Accelerator()
       UpperCamelCase_				:   Optional[int]					    =							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 inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import DetrConfig, MaskFormerConfig, SwinConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import MaskFormerForInstanceSegmentation, MaskFormerModel

      if is_vision_available():
            from transformers import MaskFormerImageProcessor

if is_vision_available():
      from PIL import Image




class     A    :






 """simple docstring"""





 def __init__(							self			,     __lowerCAmelCase			,     __lowerCAmelCase=2			,     __lowerCAmelCase=True			,     __lowerCAmelCase=False			,     __lowerCAmelCase=10			,     __lowerCAmelCase=3			,     __lowerCAmelCase=32 * 4			,     __lowerCAmelCase=32 * 6			,     __lowerCAmelCase=4			,     __lowerCAmelCase=32			,     ):
       UpperCamelCase_				:   List[str]					    =							parent
       UpperCamelCase_				:   List[Any]					    =							batch_size
       UpperCamelCase_				:   int					    =							is_training
       UpperCamelCase_				:   Union[str, Any]					    =							use_auxiliary_loss
       UpperCamelCase_				:   Optional[Any]					    =							num_queries
       UpperCamelCase_				:   List[str]					    =							num_channels
       UpperCamelCase_				:   Optional[int]					    =							min_size
       UpperCamelCase_				:   Union[str, Any]					    =							max_size
       UpperCamelCase_				:   Optional[int]					    =							num_labels
       UpperCamelCase_				:   List[Any]					    =							mask_feature_size
 def   _UpperCAmelCase       (							self ):
       UpperCamelCase_				:   str					    =							floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size] ).to(
           __lowerCAmelCase )

       UpperCamelCase_				:   Union[str, Any]					    =							torch.ones([self.batch_size, self.min_size, self.max_size]			,     device=__lowerCAmelCase )

       UpperCamelCase_				:   Dict					    =							(
           torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]			,     device=__lowerCAmelCase ) > 0.5
       ).float()
       UpperCamelCase_				:   List[Any]					    =							(torch.rand((self.batch_size, self.num_labels)			,     device=__lowerCAmelCase ) > 0.5).long()

       UpperCamelCase_				:   Dict					    =							self.get_config()
       return config, pixel_values, pixel_mask, mask_labels, class_labels
 def   _UpperCAmelCase       (							self ):
       return MaskFormerConfig.from_backbone_and_decoder_configs(
           backbone_config=SwinConfig(
               depths=[1, 1, 1, 1]			,     )			,     decoder_config=DetrConfig(
               decoder_ffn_dim=1_28			,     num_queries=self.num_queries			,     decoder_attention_heads=2			,     d_model=self.mask_feature_size			,     )			,     mask_feature_size=self.mask_feature_size			,     fpn_feature_size=self.mask_feature_size			,     num_channels=self.num_channels			,     num_labels=self.num_labels			,     )
 def   _UpperCAmelCase       (							self ):
       UpperCamelCase_    ,			UpperCamelCase_    ,			UpperCamelCase_    ,			UpperCamelCase_    ,			UpperCamelCase_				:   Optional[Any]					    =							self.prepare_config_and_inputs()
       UpperCamelCase_				:   int					    =							{"""pixel_values""": pixel_values, """pixel_mask""": pixel_mask}
       return config, inputs_dict
 def   _UpperCAmelCase       (							self			,     __lowerCAmelCase			,     __lowerCAmelCase ):
       UpperCamelCase_				:   List[str]					    =							output.encoder_hidden_states
       UpperCamelCase_				:   Any					    =							output.pixel_decoder_hidden_states
       UpperCamelCase_				:   Union[str, Any]					    =							output.transformer_decoder_hidden_states

       self.parent.assertTrue(len(__lowerCAmelCase )			,     len(config.backbone_config.depths ) )
       self.parent.assertTrue(len(__lowerCAmelCase )			,     len(config.backbone_config.depths ) )
       self.parent.assertTrue(len(__lowerCAmelCase )			,     config.decoder_config.decoder_layers )
 def   _UpperCAmelCase       (							self			,     __lowerCAmelCase			,     __lowerCAmelCase			,     __lowerCAmelCase			,     __lowerCAmelCase=False ):
       with torch.no_grad():
             UpperCamelCase_				:   Tuple					    =							MaskFormerModel(config=__lowerCAmelCase )
             model.to(__lowerCAmelCase )
             model.eval()

             UpperCamelCase_				:   str					    =							model(pixel_values=__lowerCAmelCase			,     pixel_mask=__lowerCAmelCase )
             UpperCamelCase_				:   Union[str, Any]					    =							model(__lowerCAmelCase			,     output_hidden_states=__lowerCAmelCase )
       # the correct shape of output.transformer_decoder_hidden_states ensure the correcteness of the
       # encoder and pixel decoder
       self.parent.assertEqual(
           output.transformer_decoder_last_hidden_state.shape			,     (self.batch_size, self.num_queries, self.mask_feature_size)			,     )
       # let's ensure the other two hidden state exists
       self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None )
       self.parent.assertTrue(output.encoder_last_hidden_state is not None )

       if output_hidden_states:
             self.check_output_hidden_state(__lowerCAmelCase			,     __lowerCAmelCase )





 def   _UpperCAmelCase       (							self			,     __lowerCAmelCase			,     __lowerCAmelCase			,     __lowerCAmelCase			,     __lowerCAmelCase			,     __lowerCAmelCase ):
       UpperCamelCase_				:   Optional[int]					    =							MaskFormerForInstanceSegmentation(config=__lowerCAmelCase )
       model.to(__lowerCAmelCase )
       model.eval()

       def comm_check_on_output(__lowerCAmelCase ):
             # let's still check that all the required stuff is there
             self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None )
             self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None )
             self.parent.assertTrue(result.encoder_last_hidden_state is not None )
             # okay, now we need to check the logits shape
             # due to the encoder compression, masks have a //4 spatial size
             self.parent.assertEqual(
                 result.masks_queries_logits.shape			,     (self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)			,     )
             # + 1 for null class
             self.parent.assertEqual(
                 result.class_queries_logits.shape			,     (self.batch_size, self.num_queries, self.num_labels + 1) )

       with torch.no_grad():
             UpperCamelCase_				:   List[Any]					    =							model(pixel_values=__lowerCAmelCase			,     pixel_mask=__lowerCAmelCase )
             UpperCamelCase_				:   List[str]					    =							model(__lowerCAmelCase )

             comm_check_on_output(__lowerCAmelCase )

             UpperCamelCase_				:   Dict					    =							model(
                 pixel_values=__lowerCAmelCase			,     pixel_mask=__lowerCAmelCase			,     mask_labels=__lowerCAmelCase			,     class_labels=__lowerCAmelCase )

       comm_check_on_output(__lowerCAmelCase )

       self.parent.assertTrue(result.loss is not None )
       self.parent.assertEqual(result.loss.shape			,     torch.Size([1] ) )






@require_torch
class     A    (      SCREAMING_SNAKE_CASE__,   SCREAMING_SNAKE_CASE__,   unittest.TestCase						):






 """simple docstring"""





 __a			:						Tuple      =						(MaskFormerModel, MaskFormerForInstanceSegmentation) if is_torch_available() else ()
 __a			:						List[str]      =						(
     {'''feature-extraction''': MaskFormerModel, '''image-segmentation''': MaskFormerForInstanceSegmentation}
     if is_torch_available()
     else {}
 )

 __a			:						List[Any]      =						False
 __a			:						int      =						False
 __a			:						List[str]      =						False
 __a			:						Any      =						False
 def   _UpperCAmelCase       (							self ):
       UpperCamelCase_				:   Optional[int]					    =							MaskFormerModelTester(self )
       UpperCamelCase_				:   Optional[int]					    =							ConfigTester(self			,     config_class=__lowerCAmelCase			,     has_text_modality=__lowerCAmelCase )
 def   _UpperCAmelCase       (							self ):
       self.config_tester.run_common_tests()
 def   _UpperCAmelCase       (							self ):
       UpperCamelCase_    ,			UpperCamelCase_				:   List[str]					    =							self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.create_and_check_maskformer_model(__lowerCAmelCase			,     **__lowerCAmelCase			,     output_hidden_states=__lowerCAmelCase )
 def   _UpperCAmelCase       (							self ):
       UpperCamelCase_				:   int					    =							self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_maskformer_instance_segmentation_head_model(*__lowerCAmelCase )
 @unittest.skip(reason="""MaskFormer does not use inputs_embeds""" )
 def   _UpperCAmelCase       (							self ):
       pass
 @unittest.skip(reason="""MaskFormer does not have a get_input_embeddings method""" )
 def   _UpperCAmelCase       (							self ):
       pass
 @unittest.skip(reason="""MaskFormer is not a generative model""" )
 def   _UpperCAmelCase       (							self ):
       pass
 @unittest.skip(reason="""MaskFormer does not use token embeddings""" )
 def   _UpperCAmelCase       (							self ):
       pass
 @require_torch_multi_gpu
 @unittest.skip(
     reason="""MaskFormer has some layers using `add_module` which doesn't work well with `nn.DataParallel`""" )
 def   _UpperCAmelCase       (							self ):
       pass
 @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests.""" )
 def   _UpperCAmelCase       (							self ):
       pass
 def   _UpperCAmelCase       (							self ):
       UpperCamelCase_    ,			UpperCamelCase_				:   Optional[int]					    =							self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             UpperCamelCase_				:   List[Any]					    =							model_class(__lowerCAmelCase )
             UpperCamelCase_				:   Union[str, Any]					    =							inspect.signature(model.forward )
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             UpperCamelCase_				:   Optional[Any]					    =							[*signature.parameters.keys()]

             UpperCamelCase_				:   Any					    =							["""pixel_values"""]
             self.assertListEqual(arg_names[:1]			,     __lowerCAmelCase )
 @slow
 def   _UpperCAmelCase       (							self ):
       for model_name in ["facebook/maskformer-swin-small-coco"]:
             UpperCamelCase_				:   List[str]					    =							MaskFormerModel.from_pretrained(__lowerCAmelCase )
             self.assertIsNotNone(__lowerCAmelCase )
 def   _UpperCAmelCase       (							self ):
       UpperCamelCase_				:   Tuple					    =							(self.model_tester.min_size,) * 2
       UpperCamelCase_				:   Optional[int]					    =							{
           """pixel_values""": torch.randn((2, 3, *size)			,     device=__lowerCAmelCase ),
           """mask_labels""": torch.randn((2, 10, *size)			,     device=__lowerCAmelCase ),
           """class_labels""": torch.zeros(2			,     10			,     device=__lowerCAmelCase ).long(),
       }

       UpperCamelCase_				:   Union[str, Any]					    =							MaskFormerForInstanceSegmentation(MaskFormerConfig() ).to(__lowerCAmelCase )
       UpperCamelCase_				:   List[str]					    =							model(**__lowerCAmelCase )
       self.assertTrue(outputs.loss is not None )
 def   _UpperCAmelCase       (							self ):
       UpperCamelCase_    ,			UpperCamelCase_				:   Dict					    =							self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.create_and_check_maskformer_model(__lowerCAmelCase			,     **__lowerCAmelCase			,     output_hidden_states=__lowerCAmelCase )
 def   _UpperCAmelCase       (							self ):
       UpperCamelCase_    ,			UpperCamelCase_				:   Tuple					    =							self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             UpperCamelCase_				:   Optional[int]					    =							model_class(__lowerCAmelCase ).to(__lowerCAmelCase )
             UpperCamelCase_				:   Optional[int]					    =							model(**__lowerCAmelCase			,     output_attentions=__lowerCAmelCase )
             self.assertTrue(outputs.attentions is not None )
 def   _UpperCAmelCase       (							self ):
       if not self.model_tester.is_training:
             return
       # only MaskFormerForInstanceSegmentation has the loss
       UpperCamelCase_				:   int					    =							self.all_model_classes[1]
       UpperCamelCase_    ,			UpperCamelCase_    ,			UpperCamelCase_    ,			UpperCamelCase_    ,			UpperCamelCase_				:   Tuple					    =							self.model_tester.prepare_config_and_inputs()

       UpperCamelCase_				:   Dict					    =							model_class(__lowerCAmelCase )
       model.to(__lowerCAmelCase )
       model.train()

       UpperCamelCase_				:   Union[str, Any]					    =							model(__lowerCAmelCase			,     mask_labels=__lowerCAmelCase			,     class_labels=__lowerCAmelCase ).loss
       loss.backward()





 def   _UpperCAmelCase       (							self ):
       # only MaskFormerForInstanceSegmentation has the loss
       UpperCamelCase_				:   List[Any]					    =							self.all_model_classes[1]
       UpperCamelCase_    ,			UpperCamelCase_    ,			UpperCamelCase_    ,			UpperCamelCase_    ,			UpperCamelCase_				:   Dict					    =							self.model_tester.prepare_config_and_inputs()
       UpperCamelCase_				:   Optional[int]					    =							True
       UpperCamelCase_				:   List[str]					    =							True

       UpperCamelCase_				:   Any					    =							model_class(__lowerCAmelCase )
       model.to(__lowerCAmelCase )
       model.train()

       UpperCamelCase_				:   Dict					    =							model(__lowerCAmelCase			,     mask_labels=__lowerCAmelCase			,     class_labels=__lowerCAmelCase )

       UpperCamelCase_				:   Optional[int]					    =							outputs.encoder_hidden_states[0]
       encoder_hidden_states.retain_grad()

       UpperCamelCase_				:   Optional[int]					    =							outputs.pixel_decoder_hidden_states[0]
       pixel_decoder_hidden_states.retain_grad()
       # we requires_grad=True in inputs_embeds (line 2152), the original implementation don't
       UpperCamelCase_				:   List[str]					    =							outputs.transformer_decoder_hidden_states[0]
       transformer_decoder_hidden_states.retain_grad()

       UpperCamelCase_				:   List[Any]					    =							outputs.attentions[0]
       attentions.retain_grad()

       outputs.loss.backward(retain_graph=__lowerCAmelCase )

       self.assertIsNotNone(encoder_hidden_states.grad )
       self.assertIsNotNone(pixel_decoder_hidden_states.grad )
       self.assertIsNotNone(transformer_decoder_hidden_states.grad )
       self.assertIsNotNone(attentions.grad )






UpperCamelCase							=1E-4



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





      """simple docstring"""

      UpperCamelCase_				:   int					    =							Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""	)
      return image




@require_vision
@slow
class     A    (      unittest.TestCase						):






 """simple docstring"""





 @cached_property
 def   _UpperCAmelCase       (							self ):
       return (
           MaskFormerImageProcessor.from_pretrained("""facebook/maskformer-swin-small-coco""" )
           if is_vision_available()
           else None
       )
 def   _UpperCAmelCase       (							self ):
       UpperCamelCase_				:   List[Any]					    =							MaskFormerModel.from_pretrained("""facebook/maskformer-swin-small-coco""" ).to(__lowerCAmelCase )
       UpperCamelCase_				:   Optional[Any]					    =							self.default_image_processor
       UpperCamelCase_				:   Optional[int]					    =							prepare_img()
       UpperCamelCase_				:   Optional[int]					    =							image_processor(__lowerCAmelCase			,     return_tensors="""pt""" ).to(__lowerCAmelCase )
       UpperCamelCase_				:   Union[str, Any]					    =							inputs["""pixel_values"""].shape
       # check size is divisible by 32
       self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0 )
       # check size
       self.assertEqual(__lowerCAmelCase			,     (1, 3, 8_00, 10_88) )

       with torch.no_grad():
             UpperCamelCase_				:   Tuple					    =							model(**__lowerCAmelCase )

       UpperCamelCase_				:   str					    =							torch.tensor(
           [[-0.04_82, 0.92_28, 0.49_51], [-0.25_47, 0.80_17, 0.85_27], [-0.00_69, 0.33_85, -0.00_89]] ).to(__lowerCAmelCase )
       self.assertTrue(
           torch.allclose(
               outputs.encoder_last_hidden_state[0, 0, :3, :3]			,     __lowerCAmelCase			,     atol=__lowerCAmelCase ) )

       UpperCamelCase_				:   Any					    =							torch.tensor(
           [[-0.84_22, -0.84_34, -0.97_18], [-1.01_44, -0.55_65, -0.41_95], [-1.00_38, -0.44_84, -0.19_61]] ).to(__lowerCAmelCase )
       self.assertTrue(
           torch.allclose(
               outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]			,     __lowerCAmelCase			,     atol=__lowerCAmelCase ) )

       UpperCamelCase_				:   List[Any]					    =							torch.tensor(
           [[0.28_52, -0.01_59, 0.97_35], [0.62_54, 0.18_58, 0.85_29], [-0.06_80, -0.41_16, 1.84_13]] ).to(__lowerCAmelCase )
       self.assertTrue(
           torch.allclose(
               outputs.transformer_decoder_last_hidden_state[0, :3, :3]			,     __lowerCAmelCase			,     atol=__lowerCAmelCase ) )
 def   _UpperCAmelCase       (							self ):
       UpperCamelCase_				:   Optional[int]					    =							(
           MaskFormerForInstanceSegmentation.from_pretrained("""facebook/maskformer-swin-small-coco""" )
           .to(__lowerCAmelCase )
           .eval()
       )
       UpperCamelCase_				:   int					    =							self.default_image_processor
       UpperCamelCase_				:   Tuple					    =							prepare_img()
       UpperCamelCase_				:   Optional[Any]					    =							image_processor(__lowerCAmelCase			,     return_tensors="""pt""" ).to(__lowerCAmelCase )
       UpperCamelCase_				:   List[Any]					    =							inputs["""pixel_values"""].shape
       # check size is divisible by 32
       self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0 )
       # check size
       self.assertEqual(__lowerCAmelCase			,     (1, 3, 8_00, 10_88) )

       with torch.no_grad():
             UpperCamelCase_				:   List[str]					    =							model(**__lowerCAmelCase )
       # masks_queries_logits
       UpperCamelCase_				:   Any					    =							outputs.masks_queries_logits
       self.assertEqual(
           masks_queries_logits.shape			,     (1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)			,     )
       UpperCamelCase_				:   Union[str, Any]					    =							[
           [-1.3_73_71_24, -1.7_72_49_37, -1.9_36_42_33],
           [-1.5_97_72_81, -1.9_86_79_39, -2.1_52_36_95],
           [-1.5_79_53_98, -1.9_26_98_32, -2.09_39_42],
       ]
       UpperCamelCase_				:   Optional[Any]					    =							torch.tensor(__lowerCAmelCase ).to(__lowerCAmelCase )
       self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]			,     __lowerCAmelCase			,     atol=__lowerCAmelCase ) )
       # class_queries_logits
       UpperCamelCase_				:   List[str]					    =							outputs.class_queries_logits
       self.assertEqual(
           class_queries_logits.shape			,     (1, model.config.decoder_config.num_queries, model.config.num_labels + 1) )
       UpperCamelCase_				:   Optional[Any]					    =							torch.tensor(
           [
               [1.6512E00, -5.2572E00, -3.3519E00],
               [3.6169E-02, -5.9025E00, -2.9313E00],
               [1.0766E-04, -7.7630E00, -5.1263E00],
           ] ).to(__lowerCAmelCase )
       self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]			,     __lowerCAmelCase			,     atol=__lowerCAmelCase ) )
 def   _UpperCAmelCase       (							self ):
       UpperCamelCase_				:   Optional[int]					    =							(
           MaskFormerForInstanceSegmentation.from_pretrained("""facebook/maskformer-resnet101-coco-stuff""" )
           .to(__lowerCAmelCase )
           .eval()
       )
       UpperCamelCase_				:   Optional[int]					    =							self.default_image_processor
       UpperCamelCase_				:   Optional[Any]					    =							prepare_img()
       UpperCamelCase_				:   str					    =							image_processor(__lowerCAmelCase			,     return_tensors="""pt""" ).to(__lowerCAmelCase )
       UpperCamelCase_				:   Union[str, Any]					    =							inputs["""pixel_values"""].shape
       # check size is divisible by 32
       self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0 )
       # check size
       self.assertEqual(__lowerCAmelCase			,     (1, 3, 8_00, 10_88) )

       with torch.no_grad():
             UpperCamelCase_				:   List[Any]					    =							model(**__lowerCAmelCase )
       # masks_queries_logits
       UpperCamelCase_				:   Union[str, Any]					    =							outputs.masks_queries_logits
       self.assertEqual(
           masks_queries_logits.shape			,     (1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)			,     )
       UpperCamelCase_				:   List[str]					    =							[[-0.90_46, -2.63_66, -4.60_62], [-3.41_79, -5.78_90, -8.80_57], [-4.91_79, -7.65_60, -10.77_11]]
       UpperCamelCase_				:   Any					    =							torch.tensor(__lowerCAmelCase ).to(__lowerCAmelCase )
       self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]			,     __lowerCAmelCase			,     atol=__lowerCAmelCase ) )
       # class_queries_logits
       UpperCamelCase_				:   List[Any]					    =							outputs.class_queries_logits
       self.assertEqual(
           class_queries_logits.shape			,     (1, model.config.decoder_config.num_queries, model.config.num_labels + 1) )
       UpperCamelCase_				:   Any					    =							torch.tensor(
           [[4.71_88, -3.25_85, -2.88_57], [6.68_71, -2.91_81, -1.24_87], [7.24_49, -2.27_64, -2.18_74]] ).to(__lowerCAmelCase )
       self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]			,     __lowerCAmelCase			,     atol=__lowerCAmelCase ) )





 def   _UpperCAmelCase       (							self ):
       UpperCamelCase_				:   List[Any]					    =							(
           MaskFormerForInstanceSegmentation.from_pretrained("""facebook/maskformer-swin-small-coco""" )
           .to(__lowerCAmelCase )
           .eval()
       )
       UpperCamelCase_				:   Optional[int]					    =							self.default_image_processor

       UpperCamelCase_				:   Tuple					    =							image_processor(
           [np.zeros((3, 8_00, 13_33) ), np.zeros((3, 8_00, 13_33) )]			,     segmentation_maps=[np.zeros((3_84, 3_84) ).astype(np.floataa ), np.zeros((3_84, 3_84) ).astype(np.floataa )]			,     return_tensors="""pt"""			,     )

       UpperCamelCase_				:   List[Any]					    =							inputs["""pixel_values"""].to(__lowerCAmelCase )
       UpperCamelCase_				:   List[Any]					    =							[el.to(__lowerCAmelCase ) for el in inputs["""mask_labels"""]]
       UpperCamelCase_				:   Tuple					    =							[el.to(__lowerCAmelCase ) for el in inputs["""class_labels"""]]

       with torch.no_grad():
             UpperCamelCase_				:   Union[str, Any]					    =							model(**__lowerCAmelCase )

       self.assertTrue(outputs.loss is not None )






