



import collections
import inspect
import unittest

from transformers import DonutSwinConfig
from transformers.testing_utils import require_torch, slow, torch_device
from transformers.utils import is_torch_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch
						from torch import nn

						from transformers import DonutSwinModel
						from transformers.models.donut.modeling_donut_swin import DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST






class 						_UpperCamelCase  :


			'''simple docstring'''

			def __init__(							self					:			Any							,					lowerCAmelCase__					:			Optional[Any]							,					lowerCAmelCase__					:			Dict=1_3							,					lowerCAmelCase__					:			str=3_2							,					lowerCAmelCase__					:			Union[str, Any]=2							,					lowerCAmelCase__					:			Tuple=3							,					lowerCAmelCase__					:			Optional[int]=1_6							,					lowerCAmelCase__					:			List[Any]=[1, 2, 1]							,					lowerCAmelCase__					:			List[str]=[2, 2, 4]							,					lowerCAmelCase__					:			Union[str, Any]=2							,					lowerCAmelCase__					:			Optional[Any]=2.0							,					lowerCAmelCase__					:			int=True							,					lowerCAmelCase__					:			Union[str, Any]=0.0							,					lowerCAmelCase__					:			Optional[Any]=0.0							,					lowerCAmelCase__					:			Optional[int]=0.1							,					lowerCAmelCase__					:			List[Any]="gelu"							,					lowerCAmelCase__					:			List[str]=False							,					lowerCAmelCase__					:			Union[str, Any]=True							,					lowerCAmelCase__					:			Optional[int]=0.02							,					lowerCAmelCase__					:			Optional[Any]=1E-5							,					lowerCAmelCase__					:			int=True							,					lowerCAmelCase__					:			List[str]=None							,					lowerCAmelCase__					:			List[str]=True							,					lowerCAmelCase__					:			List[Any]=1_0							,					lowerCAmelCase__					:			Optional[Any]=8							,					):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		Dict   	=		parent
					__SCREAMING_SNAKE_CASE	:		Any   	=		batch_size
					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		image_size
					__SCREAMING_SNAKE_CASE	:		List[str]   	=		patch_size
					__SCREAMING_SNAKE_CASE	:		Any   	=		num_channels
					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		embed_dim
					__SCREAMING_SNAKE_CASE	:		List[str]   	=		depths
					__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		num_heads
					__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		window_size
					__SCREAMING_SNAKE_CASE	:		int   	=		mlp_ratio
					__SCREAMING_SNAKE_CASE	:		Dict   	=		qkv_bias
					__SCREAMING_SNAKE_CASE	:		Any   	=		hidden_dropout_prob
					__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		attention_probs_dropout_prob
					__SCREAMING_SNAKE_CASE	:		Dict   	=		drop_path_rate
					__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		hidden_act
					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		use_absolute_embeddings
					__SCREAMING_SNAKE_CASE	:		List[str]   	=		patch_norm
					__SCREAMING_SNAKE_CASE	:		Tuple   	=		layer_norm_eps
					__SCREAMING_SNAKE_CASE	:		List[str]   	=		initializer_range
					__SCREAMING_SNAKE_CASE	:		Dict   	=		is_training
					__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		scope
					__SCREAMING_SNAKE_CASE	:		Tuple   	=		use_labels
					__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		type_sequence_label_size
					__SCREAMING_SNAKE_CASE	:		int   	=		encoder_stride

			def UpperCamelCase__ (							self					:			int						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		List[Any]   	=		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

					__SCREAMING_SNAKE_CASE	:		str   	=		None
					if self.use_labels:
							__SCREAMING_SNAKE_CASE	:		List[str]   	=		ids_tensor([self.batch_size]							,					self.type_sequence_label_size						)

					__SCREAMING_SNAKE_CASE	:		str   	=		self.get_config()

					return config, pixel_values, labels

			def UpperCamelCase__ (							self					:			Tuple						):





					"""simple docstring"""






					return DonutSwinConfig(
					    image_size=self.image_size							,					patch_size=self.patch_size							,					num_channels=self.num_channels							,					embed_dim=self.embed_dim							,					depths=self.depths							,					num_heads=self.num_heads							,					window_size=self.window_size							,					mlp_ratio=self.mlp_ratio							,					qkv_bias=self.qkv_bias							,					hidden_dropout_prob=self.hidden_dropout_prob							,					attention_probs_dropout_prob=self.attention_probs_dropout_prob							,					drop_path_rate=self.drop_path_rate							,					hidden_act=self.hidden_act							,					use_absolute_embeddings=self.use_absolute_embeddings							,					path_norm=self.patch_norm							,					layer_norm_eps=self.layer_norm_eps							,					initializer_range=self.initializer_range							,					encoder_stride=self.encoder_stride							,					)

			def UpperCamelCase__ (							self					:			Dict							,					lowerCAmelCase__					:			List[Any]							,					lowerCAmelCase__					:			Optional[int]							,					lowerCAmelCase__					:			List[Any]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		List[str]   	=		DonutSwinModel(config=lowerCAmelCase__						)
					model.to(lowerCAmelCase__						)
					model.eval()
					__SCREAMING_SNAKE_CASE	:		Any   	=		model(lowerCAmelCase__						)

					__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths						) - 1))
					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		int(config.embed_dim * 2 ** (len(config.depths						) - 1)						)

					self.parent.assertEqual(result.last_hidden_state.shape							,					(self.batch_size, expected_seq_len, expected_dim)						)



			def UpperCamelCase__ (							self					:			Union[str, Any]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		self.prepare_config_and_inputs()
					(
					    __SCREAMING_SNAKE_CASE						
					)	:		Dict   	=		config_and_inputs
					__SCREAMING_SNAKE_CASE	:		int   	=		{"""pixel_values""": pixel_values}
					return config, inputs_dict






@require_torch
class 						_UpperCamelCase  (							lowerCamelCase__   , lowerCamelCase__   , unittest.TestCase   ):


			'''simple docstring'''

			_A				:		int         =		(DonutSwinModel,) if is_torch_available() else ()
			_A				:		List[Any]         =		{'''feature-extraction''': DonutSwinModel} if is_torch_available() else {}
			_A				:		str         =		True

			_A				:		Tuple         =		False
			_A				:		List[Any]         =		False
			_A				:		int         =		False

			def UpperCamelCase__ (							self					:			Optional[Any]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		int   	=		DonutSwinModelTester(self						)
					__SCREAMING_SNAKE_CASE	:		Any   	=		ConfigTester(self							,					config_class=lowerCAmelCase__							,					embed_dim=3_7						)

			def UpperCamelCase__ (							self					:			Dict						):





					"""simple docstring"""






					self.create_and_test_config_common_properties()
					self.config_tester.create_and_test_config_to_json_string()
					self.config_tester.create_and_test_config_to_json_file()
					self.config_tester.create_and_test_config_from_and_save_pretrained()
					self.config_tester.create_and_test_config_with_num_labels()
					self.config_tester.check_config_can_be_init_without_params()
					self.config_tester.check_config_arguments_init()

			def UpperCamelCase__ (							self					:			Optional[Any]						):





					"""simple docstring"""






					return

			def UpperCamelCase__ (							self					:			Tuple						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*lowerCAmelCase__						)

			def UpperCamelCase__ (							self					:			int						):





					"""simple docstring"""






					pass

			def UpperCamelCase__ (							self					:			str						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		Dict   	=		self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							__SCREAMING_SNAKE_CASE	:		str   	=		model_class(lowerCAmelCase__						)
							self.assertIsInstance(model.get_input_embeddings()							,					(nn.Module)						)
							__SCREAMING_SNAKE_CASE	:		int   	=		model.get_output_embeddings()
							self.assertTrue(x is None or isinstance(lowerCAmelCase__							,					nn.Linear						)						)

			def UpperCamelCase__ (							self					:			Optional[int]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		model_class(lowerCAmelCase__						)
							__SCREAMING_SNAKE_CASE	:		List[str]   	=		inspect.signature(model.forward						)
							# signature.parameters is an OrderedDict => so arg_names order is deterministic
							__SCREAMING_SNAKE_CASE	:		Tuple   	=		[*signature.parameters.keys()]

							__SCREAMING_SNAKE_CASE	:		List[Any]   	=		["""pixel_values"""]
							self.assertListEqual(arg_names[:1]							,					lowerCAmelCase__						)

			def UpperCamelCase__ (							self					:			Union[str, Any]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		self.model_tester.prepare_config_and_inputs_for_common()
					__SCREAMING_SNAKE_CASE	:		int   	=		True

					for model_class in self.all_model_classes:
							__SCREAMING_SNAKE_CASE	:		Tuple   	=		True
							__SCREAMING_SNAKE_CASE	:		Dict   	=		False
							__SCREAMING_SNAKE_CASE	:		int   	=		True
							__SCREAMING_SNAKE_CASE	:		List[Any]   	=		model_class(lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							with torch.no_grad():
									__SCREAMING_SNAKE_CASE	:		Any   	=		model(**self._prepare_for_class(lowerCAmelCase__							,					lowerCAmelCase__						)						)
							__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		outputs.attentions
							__SCREAMING_SNAKE_CASE	:		List[str]   	=		len(self.model_tester.depths						)
							self.assertEqual(len(lowerCAmelCase__						)							,					lowerCAmelCase__						)

							# check that output_attentions also work using config
							del inputs_dict["output_attentions"]
							__SCREAMING_SNAKE_CASE	:		int   	=		True
							__SCREAMING_SNAKE_CASE	:		Any   	=		config.window_size**2
							__SCREAMING_SNAKE_CASE	:		int   	=		model_class(lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							with torch.no_grad():
									__SCREAMING_SNAKE_CASE	:		str   	=		model(**self._prepare_for_class(lowerCAmelCase__							,					lowerCAmelCase__						)						)
							__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		outputs.attentions
							self.assertEqual(len(lowerCAmelCase__						)							,					lowerCAmelCase__						)

							self.assertListEqual(
							    list(attentions[0].shape[-3:]						)							,					[self.model_tester.num_heads[0], window_size_squared, window_size_squared]							,					)
							__SCREAMING_SNAKE_CASE	:		Tuple   	=		len(lowerCAmelCase__						)

							# Check attention is always last and order is fine
							__SCREAMING_SNAKE_CASE	:		Dict   	=		True
							__SCREAMING_SNAKE_CASE	:		Tuple   	=		True
							__SCREAMING_SNAKE_CASE	:		Dict   	=		model_class(lowerCAmelCase__						)
							model.to(lowerCAmelCase__						)
							model.eval()
							with torch.no_grad():
									__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		model(**self._prepare_for_class(lowerCAmelCase__							,					lowerCAmelCase__						)						)

							if hasattr(self.model_tester							,					"""num_hidden_states_types"""						):
									__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		self.model_tester.num_hidden_states_types
							else:
									# also another +1 for reshaped_hidden_states
									__SCREAMING_SNAKE_CASE	:		int   	=		2
							self.assertEqual(out_len + added_hidden_states							,					len(lowerCAmelCase__						)						)

							__SCREAMING_SNAKE_CASE	:		List[str]   	=		outputs.attentions

							self.assertEqual(len(lowerCAmelCase__						)							,					lowerCAmelCase__						)

							self.assertListEqual(
							    list(self_attentions[0].shape[-3:]						)							,					[self.model_tester.num_heads[0], window_size_squared, window_size_squared]							,					)

			def UpperCamelCase__ (							self					:			Optional[int]							,					lowerCAmelCase__					:			Dict							,					lowerCAmelCase__					:			List[Any]							,					lowerCAmelCase__					:			List[str]							,					lowerCAmelCase__					:			str						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		model_class(lowerCAmelCase__						)
					model.to(lowerCAmelCase__						)
					model.eval()

					with torch.no_grad():
							__SCREAMING_SNAKE_CASE	:		List[str]   	=		model(**self._prepare_for_class(lowerCAmelCase__							,					lowerCAmelCase__						)						)

					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		outputs.hidden_states

					__SCREAMING_SNAKE_CASE	:		List[str]   	=		getattr(
					    self.model_tester							,					"""expected_num_hidden_layers"""							,					len(self.model_tester.depths						) + 1						)
					self.assertEqual(len(lowerCAmelCase__						)							,					lowerCAmelCase__						)

					# DonutSwin has a different seq_length
					__SCREAMING_SNAKE_CASE	:		List[str]   	=		(
					    config.patch_size
					    if isinstance(config.patch_size							,					collections.abc.Iterable						)
					    else (config.patch_size, config.patch_size)
					)

					__SCREAMING_SNAKE_CASE	:		List[str]   	=		(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

					self.assertListEqual(
					    list(hidden_states[0].shape[-2:]						)							,					[num_patches, self.model_tester.embed_dim]							,					)

					__SCREAMING_SNAKE_CASE	:		List[str]   	=		outputs.reshaped_hidden_states
					self.assertEqual(len(lowerCAmelCase__						)							,					lowerCAmelCase__						)

					__SCREAMING_SNAKE_CASE	:		List[str]   	=		reshaped_hidden_states[0].shape
					__SCREAMING_SNAKE_CASE	:		List[str]   	=		(
					    reshaped_hidden_states[0].view(lowerCAmelCase__							,					lowerCAmelCase__							,					height * width						).permute(0							,					2							,					1						)
					)
					self.assertListEqual(
					    list(reshaped_hidden_states.shape[-2:]						)							,					[num_patches, self.model_tester.embed_dim]							,					)

			def UpperCamelCase__ (							self					:			Optional[Any]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		self.model_tester.prepare_config_and_inputs_for_common()

					__SCREAMING_SNAKE_CASE	:		List[str]   	=		(
					    self.model_tester.image_size
					    if isinstance(self.model_tester.image_size							,					collections.abc.Iterable						)
					    else (self.model_tester.image_size, self.model_tester.image_size)
					)

					for model_class in self.all_model_classes:
							__SCREAMING_SNAKE_CASE	:		Tuple   	=		True
							self.check_hidden_states_output(lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__						)

							# check that output_hidden_states also work using config
							del inputs_dict["output_hidden_states"]
							__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		True

							self.check_hidden_states_output(lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__						)

			def UpperCamelCase__ (							self					:			int						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		str   	=		self.model_tester.prepare_config_and_inputs_for_common()
					__SCREAMING_SNAKE_CASE	:		int   	=		3

					__SCREAMING_SNAKE_CASE	:		int   	=		(
					    self.model_tester.image_size
					    if isinstance(self.model_tester.image_size							,					collections.abc.Iterable						)
					    else (self.model_tester.image_size, self.model_tester.image_size)
					)
					__SCREAMING_SNAKE_CASE	:		str   	=		(
					    config.patch_size
					    if isinstance(config.patch_size							,					collections.abc.Iterable						)
					    else (config.patch_size, config.patch_size)
					)

					__SCREAMING_SNAKE_CASE	:		int   	=		image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
					__SCREAMING_SNAKE_CASE	:		str   	=		image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

					for model_class in self.all_model_classes:
							__SCREAMING_SNAKE_CASE	:		int   	=		True
							self.check_hidden_states_output(lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__							,					(padded_height, padded_width)						)

							# check that output_hidden_states also work using config
							del inputs_dict["output_hidden_states"]
							__SCREAMING_SNAKE_CASE	:		Dict   	=		True
							self.check_hidden_states_output(lowerCAmelCase__							,					lowerCAmelCase__							,					lowerCAmelCase__							,					(padded_height, padded_width)						)

			@slow
			def UpperCamelCase__ (							self					:			Union[str, Any]						):





					"""simple docstring"""






					for model_name in DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							__SCREAMING_SNAKE_CASE	:		str   	=		DonutSwinModel.from_pretrained(lowerCAmelCase__						)
							self.assertIsNotNone(lowerCAmelCase__						)



			def UpperCamelCase__ (							self					:			List[Any]						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		Union[str, Any]   	=		self.model_tester.prepare_config_and_inputs_for_common()

					__SCREAMING_SNAKE_CASE	:		List[Any]   	=		_config_zero_init(lowerCAmelCase__						)
					for model_class in self.all_model_classes:
							__SCREAMING_SNAKE_CASE	:		List[str]   	=		model_class(config=lowerCAmelCase__						)
							for name, param in model.named_parameters():
									if "embeddings" not in name and param.requires_grad:
											self.assertIn(
											    ((param.data.mean() * 1E9).round() / 1E9).item()							,					[0.0, 1.0]							,					msg=F"Parameter {name} of model {model_class} seems not properly initialized"							,					)




'''simple docstring'''




import torch
import torch.nn.functional as F

from diffusers import VQDiffusionScheduler

from .test_schedulers import SchedulerCommonTest






class 						_UpperCamelCase  (							lowerCamelCase__   ):


			'''simple docstring'''

			_A				:		int         =		(VQDiffusionScheduler,)

			def UpperCamelCase__ (							self					:			List[str]							,					**lowerCAmelCase__					:			Dict						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		int   	=		{
					    """num_vec_classes""": 4_0_9_7,
					    """num_train_timesteps""": 1_0_0,
					}

					config.update(**lowerCAmelCase__						)
					return config

			def UpperCamelCase__ (							self					:			Dict							,					lowerCAmelCase__					:			int						):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE	:		Dict   	=		4
					__SCREAMING_SNAKE_CASE	:		Tuple   	=		8
					__SCREAMING_SNAKE_CASE	:		int   	=		8

					__SCREAMING_SNAKE_CASE	:		Dict   	=		torch.randint(0							,					lowerCAmelCase__							,					(batch_size, height * width)						)

					return sample

			@property
			def UpperCamelCase__ (							self					:			str						):





					"""simple docstring"""






					assert False

			def UpperCamelCase__ (							self					:			Dict							,					lowerCAmelCase__					:			str						):





					"""simple docstring"""






					def model(lowerCAmelCase__					:			Optional[Any]							,					lowerCAmelCase__					:			Any							,					*lowerCAmelCase__					:			Optional[int]						):
							__SCREAMING_SNAKE_CASE	,						__SCREAMING_SNAKE_CASE	:		Optional[int]   	=		sample.shape
							__SCREAMING_SNAKE_CASE	:		Optional[Any]   	=		torch.rand((batch_size, num_vec_classes - 1, num_latent_pixels)						)
							__SCREAMING_SNAKE_CASE	:		str   	=		F.log_softmax(logits.double()							,					dim=1						).float()
							return return_value

					return model

			def UpperCamelCase__ (							self					:			Optional[int]						):





					"""simple docstring"""






					for timesteps in [2, 5, 1_0_0, 1_0_0_0]:
							self.check_over_configs(num_train_timesteps=lowerCAmelCase__						)

			def UpperCamelCase__ (							self					:			int						):





					"""simple docstring"""






					for num_vec_classes in [5, 1_0_0, 1_0_0_0, 4_0_0_0]:
							self.check_over_configs(num_vec_classes=lowerCAmelCase__						)

			def UpperCamelCase__ (							self					:			Any						):





					"""simple docstring"""






					for t in [0, 5_0, 9_9]:
							self.check_over_forward(time_step=lowerCAmelCase__						)



			def UpperCamelCase__ (							self					:			List[Any]						):





					"""simple docstring"""






					pass