
"""simple docstring"""




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


lowerCAmelCase__							=				logging.get_logger(__name__)

lowerCAmelCase__							=				{
    'facebook/vit-mae-base': 'https://huggingface.co/facebook/vit-mae-base/resolve/main/config.json',
    # See all ViT MAE models at https://huggingface.co/models?filter=vit-mae
}





class  snake_case     (							__lowercase   ):
							UpperCAmelCase__						     =      '''vit_mae'''

							def __init__(self       ,       SCREAMING_SNAKE_CASE_=7_68       ,       SCREAMING_SNAKE_CASE_=12       ,       SCREAMING_SNAKE_CASE_=12       ,       SCREAMING_SNAKE_CASE_=30_72       ,       SCREAMING_SNAKE_CASE_="gelu"       ,       SCREAMING_SNAKE_CASE_=0.0       ,       SCREAMING_SNAKE_CASE_=0.0       ,       SCREAMING_SNAKE_CASE_=0.02       ,       SCREAMING_SNAKE_CASE_=1e-1_2       ,       SCREAMING_SNAKE_CASE_=2_24       ,       SCREAMING_SNAKE_CASE_=16       ,       SCREAMING_SNAKE_CASE_=3       ,       SCREAMING_SNAKE_CASE_=True       ,       SCREAMING_SNAKE_CASE_=16       ,       SCREAMING_SNAKE_CASE_=5_12       ,       SCREAMING_SNAKE_CASE_=8       ,       SCREAMING_SNAKE_CASE_=20_48       ,       SCREAMING_SNAKE_CASE_=0.75       ,       SCREAMING_SNAKE_CASE_=False       ,       **SCREAMING_SNAKE_CASE_       ,       ):

									"""simple docstring"""







									super().__init__(**SCREAMING_SNAKE_CASE_    )

									SCREAMING_SNAKE_CASE_						 = hidden_size
									SCREAMING_SNAKE_CASE_						 = num_hidden_layers
									SCREAMING_SNAKE_CASE_						 = num_attention_heads
									SCREAMING_SNAKE_CASE_						 = intermediate_size
									SCREAMING_SNAKE_CASE_						 = hidden_act
									SCREAMING_SNAKE_CASE_						 = hidden_dropout_prob
									SCREAMING_SNAKE_CASE_						 = attention_probs_dropout_prob
									SCREAMING_SNAKE_CASE_						 = initializer_range
									SCREAMING_SNAKE_CASE_						 = layer_norm_eps
									SCREAMING_SNAKE_CASE_						 = image_size
									SCREAMING_SNAKE_CASE_						 = patch_size
									SCREAMING_SNAKE_CASE_						 = num_channels
									SCREAMING_SNAKE_CASE_						 = qkv_bias
									SCREAMING_SNAKE_CASE_						 = decoder_num_attention_heads
									SCREAMING_SNAKE_CASE_						 = decoder_hidden_size
									SCREAMING_SNAKE_CASE_						 = decoder_num_hidden_layers
									SCREAMING_SNAKE_CASE_						 = decoder_intermediate_size
									SCREAMING_SNAKE_CASE_						 = mask_ratio
									SCREAMING_SNAKE_CASE_						 = norm_pix_loss

"""simple docstring"""




import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch



def        _lowerCamelCase ( __a    ):  # picklable for multiprocessing
		return x.sum()



def        _lowerCamelCase ( __a    ):  # picklable for multiprocessing
		return i + 1





@dataclass
class  snake_case     :
							UpperCAmelCase__						     =      42
							UpperCAmelCase__						     =      42







class  snake_case     (							__lowercase   ):

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = {}
									SCREAMING_SNAKE_CASE_						 = []
									SCREAMING_SNAKE_CASE_						 = 1
									SCREAMING_SNAKE_CASE_						 = [1, 2]
									SCREAMING_SNAKE_CASE_						 = {'''a''': 1, '''b''': 2}
									SCREAMING_SNAKE_CASE_						 = {'''a''': [1, 2], '''b''': [3, 4]}
									SCREAMING_SNAKE_CASE_						 = {'''a''': {'''1''': 1}, '''b''': 2}
									SCREAMING_SNAKE_CASE_						 = {'''a''': 1, '''b''': 2, '''c''': 3, '''d''': 4}
									SCREAMING_SNAKE_CASE_						 = {}
									SCREAMING_SNAKE_CASE_						 = []
									SCREAMING_SNAKE_CASE_						 = 2
									SCREAMING_SNAKE_CASE_						 = [2, 3]
									SCREAMING_SNAKE_CASE_						 = {'''a''': 2, '''b''': 3}
									SCREAMING_SNAKE_CASE_						 = {'''a''': [2, 3], '''b''': [4, 5]}
									SCREAMING_SNAKE_CASE_						 = {'''a''': {'''1''': 2}, '''b''': 3}
									SCREAMING_SNAKE_CASE_						 = {'''a''': 2, '''b''': 3, '''c''': 4, '''d''': 5}
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )

									SCREAMING_SNAKE_CASE_						 = 2
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )

									SCREAMING_SNAKE_CASE_						 = {'''a''': np.eye(2    ), '''b''': np.zeros(3    ), '''c''': np.ones(2    )}
									SCREAMING_SNAKE_CASE_						 = {'''a''': 2, '''b''': 0, '''c''': 2}
									SCREAMING_SNAKE_CASE_						 = {
									    '''a''': np.eye(2    ).astype(SCREAMING_SNAKE_CASE_    ),
									    '''b''': np.zeros(3    ).astype(SCREAMING_SNAKE_CASE_    ),
									    '''c''': np.ones(2    ).astype(SCREAMING_SNAKE_CASE_    ),
									}
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       map_numpy=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(
									    {k: v.tolist() for k, v in map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       map_numpy=SCREAMING_SNAKE_CASE_    ).items()}       ,       {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}       ,       )
									self.assertEqual(map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       map_numpy=SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )
									self.assertEqual(
									    {k: v.tolist() for k, v in map_nested(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       map_numpy=SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    ).items()}       ,       {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}       ,       )
									with self.assertRaises(SCREAMING_SNAKE_CASE_    ):  # can't pickle a local lambda
											map_nested(lambda SCREAMING_SNAKE_CASE_    : x + 1       ,       SCREAMING_SNAKE_CASE_       ,       num_proc=SCREAMING_SNAKE_CASE_    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = {'''a''': 1, '''b''': 2}
									SCREAMING_SNAKE_CASE_						 = {'''a''': 3, '''b''': 4}
									SCREAMING_SNAKE_CASE_						 = {'''a''': 5, '''b''': 6}
									SCREAMING_SNAKE_CASE_						 = sorted([('''a''', (1, 3, 5)), ('''b''', (2, 4, 6))]    )
									self.assertEqual(sorted(zip_dict(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )    )       ,       SCREAMING_SNAKE_CASE_    )







							def _lowercase					(self    ):

									"""simple docstring"""







									class  snake_case     :
																UpperCAmelCase__						     =      '''bar'''







									SCREAMING_SNAKE_CASE_						 = Foo()
									self.assertEqual(foo.my_attr       ,       '''bar'''    )
									with temporary_assignment(SCREAMING_SNAKE_CASE_       ,       '''my_attr'''       ,       '''BAR'''    ):
											self.assertEqual(foo.my_attr       ,       '''BAR'''    )
									self.assertEqual(foo.my_attr       ,       '''bar'''    )







@pytest.mark.parametrize(
    '''iterable_length, num_proc, expected_num_proc''',							[
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ],							)
def        _lowerCamelCase ( __a,							__a,							__a    ):
		with patch('''datasets.utils.py_utils._single_map_nested'''    ) as mock_single_map_nested, patch(
		    '''datasets.parallel.parallel.Pool'''    ) as mock_multiprocessing_pool:
				SCREAMING_SNAKE_CASE_						 = {F'{i}': i for i in range(__a    )}
				SCREAMING_SNAKE_CASE_						 = map_nested(lambda __a    : x + 10,							__a,							num_proc=__a,							parallel_min_length=16    )
				if expected_num_proc == 1:
						assert mock_single_map_nested.called
						assert not mock_multiprocessing_pool.called
				else:
						assert not mock_single_map_nested.called
						assert mock_multiprocessing_pool.called
						assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc





class  snake_case     (							__lowercase   ):

							@require_tf
							def _lowercase					(self    ):

									"""simple docstring"""







									import tensorflow as tf
									from tensorflow.keras import layers

									SCREAMING_SNAKE_CASE_						 = layers.Dense(2    )

									def gen_random_output():
											SCREAMING_SNAKE_CASE_						 = tf.random.uniform((1, 3)    )
											return model(SCREAMING_SNAKE_CASE_    ).numpy()

									with temp_seed(42       ,       set_tensorflow=SCREAMING_SNAKE_CASE_    ):
											SCREAMING_SNAKE_CASE_						 = gen_random_output()
									with temp_seed(42       ,       set_tensorflow=SCREAMING_SNAKE_CASE_    ):
											SCREAMING_SNAKE_CASE_						 = gen_random_output()
									SCREAMING_SNAKE_CASE_						 = gen_random_output()

									np.testing.assert_equal(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )
									self.assertGreater(np.abs(outa - outa    ).sum()       ,       0    )

							@require_torch
							def _lowercase					(self    ):

									"""simple docstring"""







									import torch

									def gen_random_output():
											SCREAMING_SNAKE_CASE_						 = torch.nn.Linear(3       ,       2    )
											SCREAMING_SNAKE_CASE_						 = torch.rand(1       ,       3    )
											return model(SCREAMING_SNAKE_CASE_    ).detach().numpy()

									with temp_seed(42       ,       set_pytorch=SCREAMING_SNAKE_CASE_    ):
											SCREAMING_SNAKE_CASE_						 = gen_random_output()
									with temp_seed(42       ,       set_pytorch=SCREAMING_SNAKE_CASE_    ):
											SCREAMING_SNAKE_CASE_						 = gen_random_output()
									SCREAMING_SNAKE_CASE_						 = gen_random_output()

									np.testing.assert_equal(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )
									self.assertGreater(np.abs(outa - outa    ).sum()       ,       0    )







							def _lowercase					(self    ):

									"""simple docstring"""







									def gen_random_output():
											return np.random.rand(1       ,       3    )

									with temp_seed(42    ):
											SCREAMING_SNAKE_CASE_						 = gen_random_output()
									with temp_seed(42    ):
											SCREAMING_SNAKE_CASE_						 = gen_random_output()
									SCREAMING_SNAKE_CASE_						 = gen_random_output()

									np.testing.assert_equal(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )
									self.assertGreater(np.abs(outa - outa    ).sum()       ,       0    )







@pytest.mark.parametrize('''input_data''',							[{}]    )
def        _lowerCamelCase ( __a    ):
		SCREAMING_SNAKE_CASE_						 = NestedDataStructure(__a    ).data
		assert output_data == input_data



@pytest.mark.parametrize(
    '''data, expected_output''',							[
        ({}, []),
        ([], []),
        ('''foo''', ['''foo''']),
        (['''foo''', '''bar'''], ['''foo''', '''bar''']),
        ([['''foo''', '''bar''']], ['''foo''', '''bar''']),
        ([[['''foo'''], ['''bar''']]], ['''foo''', '''bar''']),
        ([[['''foo'''], '''bar''']], ['''foo''', '''bar''']),
        ({'''a''': 1, '''b''': 2}, [1, 2]),
        ({'''a''': [1, 2], '''b''': [3, 4]}, [1, 2, 3, 4]),
        ({'''a''': [[1, 2]], '''b''': [[3, 4]]}, [1, 2, 3, 4]),
        ({'''a''': [[1, 2]], '''b''': [3, 4]}, [1, 2, 3, 4]),
        ({'''a''': [[[1], [2]]], '''b''': [[[3], [4]]]}, [1, 2, 3, 4]),
        ({'''a''': [[[1], [2]]], '''b''': [[3, 4]]}, [1, 2, 3, 4]),
        ({'''a''': [[[1], [2]]], '''b''': [3, 4]}, [1, 2, 3, 4]),
        ({'''a''': [[[1], [2]]], '''b''': [3, [4]]}, [1, 2, 3, 4]),
        ({'''a''': {'''1''': 1}, '''b''': 2}, [1, 2]),
        ({'''a''': {'''1''': [1]}, '''b''': 2}, [1, 2]),
        ({'''a''': {'''1''': [1]}, '''b''': [2]}, [1, 2]),
    ],							)
def        _lowerCamelCase ( __a,							__a    ):
		SCREAMING_SNAKE_CASE_						 = NestedDataStructure(__a    ).flatten()
		assert output == expected_output



def        _lowerCamelCase ( ):
		SCREAMING_SNAKE_CASE_						 = A(x=1,							y='''foobar'''    )
		SCREAMING_SNAKE_CASE_						 = {'''x''': 1, '''y''': '''foobar'''}
		assert asdict(__a    ) == expected_output

		SCREAMING_SNAKE_CASE_						 = {'''a''': {'''b''': A(x=10,							y='''foo'''    )}, '''c''': [A(x=20,							y='''bar'''    )]}
		SCREAMING_SNAKE_CASE_						 = {'''a''': {'''b''': {'''x''': 10, '''y''': '''foo'''}}, '''c''': [{'''x''': 20, '''y''': '''bar'''}]}
		assert asdict(__a    ) == expected_output

		with pytest.raises(__a    ):
				asdict([1, A(x=10,							y='''foo'''    )]    )



def        _lowerCamelCase ( __a    ):
		return text.split()



def        _lowerCamelCase ( __a    ):
		yield (time.time(), content)
		time.sleep(2    )
		yield (time.time(), content)



def        _lowerCamelCase ( ):
		with Pool(2    ) as pool:
				SCREAMING_SNAKE_CASE_						 = list(iflatmap_unordered(__a,							_split_text,							kwargs_iterable=[{'''text''': '''hello there'''}] * 10    )    )
				assert out.count('''hello'''    ) == 10
				assert out.count('''there'''    ) == 10
				assert len(__a    ) == 20

		# check multiprocess from pathos (uses dill for pickling)
		with multiprocess.Pool(2    ) as pool:
				SCREAMING_SNAKE_CASE_						 = list(iflatmap_unordered(__a,							_split_text,							kwargs_iterable=[{'''text''': '''hello there'''}] * 10    )    )
				assert out.count('''hello'''    ) == 10
				assert out.count('''there'''    ) == 10
				assert len(__a    ) == 20

		# check that we get items as fast as possible
		with Pool(2    ) as pool:
				SCREAMING_SNAKE_CASE_						 = []
				for yield_time, content in iflatmap_unordered(
				    __a,							_aseconds_generator_of_aitems_with_timing,							kwargs_iterable=[{'''content''': '''a'''}, {'''content''': '''b'''}]    ):
						assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
						out.append(__a    )
				assert out.count('''a'''    ) == 2
				assert out.count('''b'''    ) == 2
				assert len(__a    ) == 4