






"""simple docstring"""







import argparse
import os
from functools import reduce

import fairseq
import torch
from datasets import load_dataset

from transformers import WavaVecaProcessor, logging
from transformers.models.dataavec.configuration_dataavec_audio import DataaVecAudioConfig

# Copied from https://github.com/pytorch/fairseq/blob/main/examples/data2vec/models/data2vec_audio.py
from transformers.models.dataavec.dataavec_audio import DataaVecAudioModel as Dummy  # noqa: F401
from transformers.models.dataavec.modeling_dataavec_audio import DataaVecAudioForCTC, DataaVecAudioModel


logging.set_verbosity_info()
lowerCAmelCase_		:			Optional[Any]				   =					logging.get_logger(__name__)

lowerCAmelCase_		:			List[str]				   =					{
    '''post_extract_proj''': '''feature_projection.projection''',
    '''models.0.layer_norm''': '''feature_projection.layer_norm''',
    '''self_attn.k_proj''': '''encoder.layers.*.attention.k_proj''',
    '''self_attn.v_proj''': '''encoder.layers.*.attention.v_proj''',
    '''self_attn.q_proj''': '''encoder.layers.*.attention.q_proj''',
    '''self_attn.out_proj''': '''encoder.layers.*.attention.out_proj''',
    '''self_attn_layer_norm''': '''encoder.layers.*.layer_norm''',
    '''fc1''': '''encoder.layers.*.feed_forward.intermediate_dense''',
    '''fc2''': '''encoder.layers.*.feed_forward.output_dense''',
    '''final_layer_norm''': '''encoder.layers.*.final_layer_norm''',
    '''encoder.layer_norm''': '''encoder.layer_norm''',
    '''w2v_model.layer_norm''': '''feature_projection.layer_norm''',
    '''w2v_encoder.proj''': '''lm_head''',
    '''mask_emb''': '''masked_spec_embed''',
}
lowerCAmelCase_		:			Union[str, Any]				   =					[
    '''lm_head''',
]






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       for attribute in key.split("""."""						):
              UpperCAmelCase					=	getattr(lowerCAmelCase		,					lowerCAmelCase						)

       if weight_type is not None:
              UpperCAmelCase					=	getattr(lowerCAmelCase		,					lowerCAmelCase						).shape
       else:
              UpperCAmelCase					=	hf_pointer.shape

       if hf_shape != value.shape:
              raise ValueError(
                  F'''Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be'''
                  F''' {value.shape} for {full_name}'''						)

       if weight_type == "weight":
              UpperCAmelCase					=	value
       elif weight_type == "weight_g":
              UpperCAmelCase					=	value
       elif weight_type == "weight_v":
              UpperCAmelCase					=	value
       elif weight_type == "bias":
              UpperCAmelCase					=	value
       else:
              UpperCAmelCase					=	value

       logger.info(F'''{key + '.' + weight_type if weight_type is not None else ''} was initialized from {full_name}.'''						)






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       UpperCAmelCase					=	[]
       UpperCAmelCase					=	fairseq_model.state_dict()

       if not is_headless:
              UpperCAmelCase					=	hf_model.dataavec_audio.feature_extractor
              UpperCAmelCase					=	hf_model.dataavec_audio.encoder.pos_conv_embed

       else:
              UpperCAmelCase					=	hf_model.feature_extractor
              UpperCAmelCase					=	hf_model.encoder.pos_conv_embed

       for name, value in fairseq_dict.items():
              UpperCAmelCase					=	False
              if "conv_layers" in name:
                     load_conv_layer(
                         lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					)
                     UpperCAmelCase					=	True
              elif "pos_conv" in name:
                     load_pos_conv_layer(
                         lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					)
                     UpperCAmelCase					=	True
              else:
                     for key, mapped_key in MAPPING.items():
                            if not is_headless:
                                   UpperCAmelCase					=	"""data2vec_audio.""" + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
                            if key in name or key.split("""w2v_model."""						)[-1] == name.split("""."""						)[0]:
                                   UpperCAmelCase					=	True
                                   if "*" in mapped_key:
                                          UpperCAmelCase					=	name.split(lowerCAmelCase						)[0].split("""."""						)[-2]
                                          UpperCAmelCase					=	mapped_key.replace("""*"""		,					lowerCAmelCase						)
                                   if "weight_g" in name:
                                          UpperCAmelCase					=	"""weight_g"""
                                   elif "weight_v" in name:
                                          UpperCAmelCase					=	"""weight_v"""
                                   elif "bias" in name:
                                          UpperCAmelCase					=	"""bias"""
                                   elif "weight" in name:
                                          # TODO: don't match quantizer.weight_proj
                                          UpperCAmelCase					=	"""weight"""
                                   else:
                                          UpperCAmelCase					=	None
                                   set_recursively(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						)
                            continue
              if not is_used:
                     unused_weights.append(lowerCAmelCase						)

       logger.warning(F'''Unused weights: {unused_weights}'''						)






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       UpperCAmelCase					=	path.split("""."""						)
       return reduce(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						)






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       UpperCAmelCase					=	access_by_string(lowerCAmelCase		,					lowerCAmelCase						)
       UpperCAmelCase					=	hf_weight.data

       if fsq_value.shape != hf_value.shape:
              raise ValueError(F'''{full_name} has size {fsq_value.shape}, but {hf_value.shape} was found.'''						)
       UpperCAmelCase					=	fsq_value
       logger.info(F'''{full_name} was correctly initialized from {hf_weight_path}.'''						)






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       UpperCAmelCase					=	full_name.split("""conv_layers."""						)[-1]
       UpperCAmelCase					=	name.split("""."""						)
       UpperCAmelCase					=	int(items[0]						)
       UpperCAmelCase					=	int(items[1]						)

       UpperCAmelCase					=	name.split("""."""						)[-1]
       if type_id == 0:
              UpperCAmelCase					=	"""conv"""
       elif type_id == 2:
              UpperCAmelCase					=	"""layer_norm"""
       else:
              unused_weights.append(lowerCAmelCase						)
              return

       set_weights(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					F'''conv_layers.{layer_id}.{layer_type}.{weight_type}'''						)






def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase						):







       '''simple docstring'''
       UpperCAmelCase					=	full_name.split("""pos_conv."""						)[-1]
       UpperCAmelCase					=	name.split("""."""						)
       UpperCAmelCase					=	int(items[0]						)
       UpperCAmelCase					=	int(items[1]						)

       UpperCAmelCase					=	name.split("""."""						)[-1]
       if type_id != 0:
              unused_weights.append(lowerCAmelCase						)
              return
       else:
              UpperCAmelCase					=	"""conv"""

       set_weights(lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase		,					F'''layers.{layer_id}.{layer_type}.{weight_type}'''						)






@torch.no_grad()
def 					_lowerCAmelCase				(					lowerCAmelCase		,					lowerCAmelCase		,					lowerCAmelCase=None		,					lowerCAmelCase=None		,					lowerCAmelCase=True						):







       '''simple docstring'''
       if config_path is not None:
              UpperCAmelCase					=	DataaVecAudioConfig.from_pretrained(lowerCAmelCase						)
       else:
              UpperCAmelCase					=	DataaVecAudioConfig()

       if not is_finetuned:
              # Modify final_proj layer name
              UpperCAmelCase					=	DataaVecAudioModel(lowerCAmelCase						)
              UpperCAmelCase					=	os.path.dirname(lowerCAmelCase						)

              UpperCAmelCase					=	torch.load(lowerCAmelCase						)
              UpperCAmelCase					=	state_dict["""model"""].pop("""final_proj.0.weight"""						)
              UpperCAmelCase					=	state_dict["""model"""].pop("""final_proj.0.bias"""						)
              UpperCAmelCase					=	os.path.join(lowerCAmelCase		,					"""converted.pt"""						)
              torch.save(lowerCAmelCase		,					lowerCAmelCase						)
       else:
              UpperCAmelCase					=	DataaVecAudioForCTC(lowerCAmelCase						)
              UpperCAmelCase					=	checkpoint_path

       def load_dataavec(lowerCAmelCase						):
              UpperCAmelCase			,				UpperCAmelCase			,				UpperCAmelCase					=	fairseq.checkpoint_utils.load_model_ensemble_and_task([path]						)
              return model[0].eval()

       UpperCAmelCase					=	load_dataavec(lowerCAmelCase						)

       recursively_load_weights(lowerCAmelCase		,					lowerCAmelCase		,					not is_finetuned						)

       UpperCAmelCase					=	WavaVecaProcessor.from_pretrained("""facebook/wav2vec2-large-lv60"""						)

       UpperCAmelCase					=	load_dataset("""patrickvonplaten/librispeech_asr_dummy"""		,					"""clean"""		,					split="""validation"""						)
       UpperCAmelCase					=	[x["""array"""] for x in ds[:4]["""audio"""]]

       UpperCAmelCase					=	processor(lowerCAmelCase		,					return_tensors="""pt"""		,					padding=lowerCAmelCase						)

       UpperCAmelCase					=	inputs.input_values
       UpperCAmelCase					=	inputs.attention_mask
       #    input_values = inputs.input_values[:, :-1]
       #    attention_mask = inputs.attention_mask[:, :-1]

       hf_wavavec.eval()
       model.eval()
       if is_finetuned:
              UpperCAmelCase					=	model(source=lowerCAmelCase		,					padding_mask=(1 - attention_mask)		,					mask=lowerCAmelCase		,					features_only=lowerCAmelCase						)[
                  """encoder_out"""
              ].transpose(0		,					1						)
              UpperCAmelCase					=	hf_wavavec(lowerCAmelCase		,					attention_mask=lowerCAmelCase						)["""logits"""]

              UpperCAmelCase					=	torch.argmax(lowerCAmelCase		,					dim=-1						)
              UpperCAmelCase					=	processor.batch_decode(lowerCAmelCase						)

              print(F'''Expected Output: {ds[:4]['text']}, Pred: {output_string}'''						)
       else:
              UpperCAmelCase					=	model(source=lowerCAmelCase		,					padding_mask=(1 - attention_mask)		,					mask=lowerCAmelCase		,					features_only=lowerCAmelCase						)[
                  """layer_results"""
              ][-1][0].transpose(0		,					1						)
              UpperCAmelCase					=	hf_wavavec(lowerCAmelCase		,					attention_mask=lowerCAmelCase						)["""last_hidden_state"""]

       print(our_output.shape		,					their_output.shape						)
       UpperCAmelCase					=	torch.max(torch.abs(our_output - their_output						)						).item()
       print(F'''max_absolute_diff = {max_absolute_diff}'''						)  # ~ 1e-7
       UpperCAmelCase					=	torch.allclose(lowerCAmelCase		,					lowerCAmelCase		,					atol=1e-3						)
       print("""Do both models output the same tensors?"""		,					"""🔥""" if success else """💩"""						)
       if not success:
              raise Exception("""Something went wRoNg"""						)

       hf_wavavec.save_pretrained(lowerCAmelCase						)

       if is_finetuned:
              processor.save_pretrained(lowerCAmelCase						)
       else:
              processor.feature_extractor.save_pretrained(lowerCAmelCase						)


if __name__ == "__main__":
 lowerCAmelCase_		:			int				   =					argparse.ArgumentParser()
 parser.add_argument('''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model.''')
 parser.add_argument('''--checkpoint_path''', default=None, type=str, help='''Path to fairseq checkpoint''')
 parser.add_argument('''--dict_path''', default=None, type=str, help='''Path to dict of fine-tuned model''')
 parser.add_argument('''--config_path''', default=None, type=str, help='''Path to hf config.json of model to convert''')
 parser.add_argument(
     '''--not_finetuned''', action='''store_true''', help='''Whether the model to convert is a fine-tuned model or not'''
 )
 lowerCAmelCase_		:			int				   =					parser.parse_args()
 convert_wavaveca_checkpoint(
     args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
 )








"""simple docstring"""







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_       :



  def __init__(					self       ,						snake_case__       ,						snake_case__=13       ,						snake_case__=32       ,						snake_case__=2       ,						snake_case__=3       ,						snake_case__=16       ,						snake_case__=[1, 2, 1]       ,						snake_case__=[2, 2, 4]       ,						snake_case__=2       ,						snake_case__=2.0       ,						snake_case__=True       ,						snake_case__=0.0       ,						snake_case__=0.0       ,						snake_case__=0.1       ,						snake_case__="gelu"       ,						snake_case__=False       ,						snake_case__=True       ,						snake_case__=0.02       ,						snake_case__=1e-5       ,						snake_case__=True       ,						snake_case__=None       ,						snake_case__=True       ,						snake_case__=10       ,						snake_case__=8       ,						)	->	Optional[int]:



         """simple docstring"""





         UpperCAmelCase					=	parent
         UpperCAmelCase					=	batch_size
         UpperCAmelCase					=	image_size
         UpperCAmelCase					=	patch_size
         UpperCAmelCase					=	num_channels
         UpperCAmelCase					=	embed_dim
         UpperCAmelCase					=	depths
         UpperCAmelCase					=	num_heads
         UpperCAmelCase					=	window_size
         UpperCAmelCase					=	mlp_ratio
         UpperCAmelCase					=	qkv_bias
         UpperCAmelCase					=	hidden_dropout_prob
         UpperCAmelCase					=	attention_probs_dropout_prob
         UpperCAmelCase					=	drop_path_rate
         UpperCAmelCase					=	hidden_act
         UpperCAmelCase					=	use_absolute_embeddings
         UpperCAmelCase					=	patch_norm
         UpperCAmelCase					=	layer_norm_eps
         UpperCAmelCase					=	initializer_range
         UpperCAmelCase					=	is_training
         UpperCAmelCase					=	scope
         UpperCAmelCase					=	use_labels
         UpperCAmelCase					=	type_sequence_label_size
         UpperCAmelCase					=	encoder_stride



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



         """simple docstring"""





         UpperCAmelCase					=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

         UpperCAmelCase					=	None
         if self.use_labels:
                UpperCAmelCase					=	ids_tensor([self.batch_size]       ,						self.type_sequence_label_size )

         UpperCAmelCase					=	self.get_config()

         return config, pixel_values, labels



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



         """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       ,						snake_case__       ,						snake_case__       ,						snake_case__ )	->	Optional[int]:



         """simple docstring"""





         UpperCAmelCase					=	DonutSwinModel(config=snake_case__ )
         model.to(snake_case__ )
         model.eval()
         UpperCAmelCase					=	model(snake_case__ )

         UpperCAmelCase					=	((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths ) - 1))
         UpperCAmelCase					=	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 )	->	Optional[Any]:



         """simple docstring"""





         UpperCAmelCase					=	self.prepare_config_and_inputs()
         (
             (
             UpperCAmelCase
         )			,				(
             UpperCAmelCase
         )			,				(
             UpperCAmelCase
         )			,				
         )					=	config_and_inputs
         UpperCAmelCase					=	{"""pixel_values""": pixel_values}
         return config, inputs_dict





@require_torch
class      UpperCamelCase_       (				a_      ,    a_      ,    unittest.TestCase     ):
  _A				:    Tuple =		(DonutSwinModel,) if is_torch_available() else ()
  _A				:    Any =		{'feature-extraction': DonutSwinModel} if is_torch_available() else {}
  _A				:    Optional[Any] =		True

  _A				:    Dict =		False
  _A				:    Any =		False
  _A				:    Tuple =		False



  def 						UpperCamelCase_ (					self )	->	Any:



         """simple docstring"""





         UpperCAmelCase					=	DonutSwinModelTester(self )
         UpperCAmelCase					=	ConfigTester(self       ,						config_class=snake_case__       ,						embed_dim=37 )



  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 )	->	List[str]:



         """simple docstring"""





         return



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



         """simple docstring"""





         UpperCAmelCase					=	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*snake_case__ )



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



         """simple docstring"""





         pass



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



         """simple docstring"""





         UpperCAmelCase			,				UpperCAmelCase					=	self.model_tester.prepare_config_and_inputs_for_common()

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



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



         """simple docstring"""





         UpperCAmelCase			,				UpperCAmelCase					=	self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
                UpperCAmelCase					=	model_class(snake_case__ )
                UpperCAmelCase					=	inspect.signature(model.forward )
                # signature.parameters is an OrderedDict => so arg_names order is deterministic
                UpperCAmelCase					=	[*signature.parameters.keys()]

                UpperCAmelCase					=	["""pixel_values"""]
                self.assertListEqual(arg_names[:1]       ,						snake_case__ )



  def 						UpperCamelCase_ (					self )	->	Tuple:



         """simple docstring"""





         UpperCAmelCase			,				UpperCAmelCase					=	self.model_tester.prepare_config_and_inputs_for_common()
         UpperCAmelCase					=	True

         for model_class in self.all_model_classes:
                UpperCAmelCase					=	True
                UpperCAmelCase					=	False
                UpperCAmelCase					=	True
                UpperCAmelCase					=	model_class(snake_case__ )
                model.to(snake_case__ )
                model.eval()
                with torch.no_grad():
                       UpperCAmelCase					=	model(**self._prepare_for_class(snake_case__       ,						snake_case__ ) )
                UpperCAmelCase					=	outputs.attentions
                UpperCAmelCase					=	len(self.model_tester.depths )
                self.assertEqual(len(snake_case__ )       ,						snake_case__ )

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

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

                # Check attention is always last and order is fine
                UpperCAmelCase					=	True
                UpperCAmelCase					=	True
                UpperCAmelCase					=	model_class(snake_case__ )
                model.to(snake_case__ )
                model.eval()
                with torch.no_grad():
                       UpperCAmelCase					=	model(**self._prepare_for_class(snake_case__       ,						snake_case__ ) )

                if hasattr(self.model_tester       ,						"""num_hidden_states_types""" ):
                       UpperCAmelCase					=	self.model_tester.num_hidden_states_types
                else:
                       # also another +1 for reshaped_hidden_states
                       UpperCAmelCase					=	2
                self.assertEqual(out_len + added_hidden_states       ,						len(snake_case__ ) )

                UpperCAmelCase					=	outputs.attentions

                self.assertEqual(len(snake_case__ )       ,						snake_case__ )

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



  def 						UpperCamelCase_ (					self       ,						snake_case__       ,						snake_case__       ,						snake_case__       ,						snake_case__ )	->	Optional[int]:



         """simple docstring"""





         UpperCAmelCase					=	model_class(snake_case__ )
         model.to(snake_case__ )
         model.eval()

         with torch.no_grad():
                UpperCAmelCase					=	model(**self._prepare_for_class(snake_case__       ,						snake_case__ ) )

         UpperCAmelCase					=	outputs.hidden_states

         UpperCAmelCase					=	getattr(
             self.model_tester       ,						"""expected_num_hidden_layers"""       ,						len(self.model_tester.depths ) + 1 )
         self.assertEqual(len(snake_case__ )       ,						snake_case__ )

         # DonutSwin has a different seq_length
         UpperCAmelCase					=	(
             config.patch_size
             if isinstance(config.patch_size       ,						collections.abc.Iterable )
             else (config.patch_size, config.patch_size)
         )

         UpperCAmelCase					=	(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]       ,						)

         UpperCAmelCase					=	outputs.reshaped_hidden_states
         self.assertEqual(len(snake_case__ )       ,						snake_case__ )

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



  def 						UpperCamelCase_ (					self )	->	str:



         """simple docstring"""





         UpperCAmelCase			,				UpperCAmelCase					=	self.model_tester.prepare_config_and_inputs_for_common()

         UpperCAmelCase					=	(
             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:
                UpperCAmelCase					=	True
                self.check_hidden_states_output(snake_case__       ,						snake_case__       ,						snake_case__       ,						snake_case__ )

                # check that output_hidden_states also work using config
                del inputs_dict["output_hidden_states"]
                UpperCAmelCase					=	True

                self.check_hidden_states_output(snake_case__       ,						snake_case__       ,						snake_case__       ,						snake_case__ )



  def 						UpperCamelCase_ (					self )	->	Dict:



         """simple docstring"""





         UpperCAmelCase			,				UpperCAmelCase					=	self.model_tester.prepare_config_and_inputs_for_common()
         UpperCAmelCase					=	3

         UpperCAmelCase					=	(
             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)
         )
         UpperCAmelCase					=	(
             config.patch_size
             if isinstance(config.patch_size       ,						collections.abc.Iterable )
             else (config.patch_size, config.patch_size)
         )

         UpperCAmelCase					=	image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
         UpperCAmelCase					=	image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

         for model_class in self.all_model_classes:
                UpperCAmelCase					=	True
                self.check_hidden_states_output(snake_case__       ,						snake_case__       ,						snake_case__       ,						(padded_height, padded_width) )

                # check that output_hidden_states also work using config
                del inputs_dict["output_hidden_states"]
                UpperCAmelCase					=	True
                self.check_hidden_states_output(snake_case__       ,						snake_case__       ,						snake_case__       ,						(padded_height, padded_width) )



  @slow
  def 						UpperCamelCase_ (					self )	->	List[str]:



         """simple docstring"""





         for model_name in DONUT_SWIN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                UpperCAmelCase					=	DonutSwinModel.from_pretrained(snake_case__ )
                self.assertIsNotNone(snake_case__ )



  def 						UpperCamelCase_ (					self )	->	Any:



         """simple docstring"""





         UpperCAmelCase			,				UpperCAmelCase					=	self.model_tester.prepare_config_and_inputs_for_common()

         UpperCAmelCase					=	_config_zero_init(snake_case__ )
         for model_class in self.all_model_classes:
                UpperCAmelCase					=	model_class(config=snake_case__ )
                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'''       ,						)
