





"""simple docstring"""



import os
import tempfile
import unittest
from pathlib import Path

from transformers import AutoConfig, is_torch_available
from transformers.testing_utils import require_torch, torch_device


if is_torch_available():
       from transformers import PyTorchBenchmark, PyTorchBenchmarkArguments






@require_torch
class 		__snake_case					(		unittest.TestCase     ):





 """simple docstring"""


 def 					UpperCamelCase__(    self					,						__lowerCamelCase  ):



     '''simple docstring'''




     for model_result in results.values():
         for batch_size, sequence_length in zip(model_result['''bs''']					,						model_result['''ss''']  ):
             __A    :						int						 =     model_result['''result'''][batch_size][sequence_length]
             self.assertIsNotNone(__lowerCamelCase  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						List[Any]						 =     '''sshleifer/tiny-gpt2'''
     __A    :						Dict						 =     PyTorchBenchmarkArguments(
         models=[MODEL_ID]					,						training=__lowerCamelCase					,						inference=__lowerCamelCase					,						sequence_lengths=[8]					,						batch_sizes=[1]					,						multi_process=__lowerCamelCase					,						)
     __A    :						Any						 =     PyTorchBenchmark(__lowerCamelCase  )
     __A    :						int						 =     benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						List[str]						 =     '''sgugger/tiny-distilbert-classification'''
     __A    :						Any						 =     PyTorchBenchmarkArguments(
         models=[MODEL_ID]					,						training=__lowerCamelCase					,						inference=__lowerCamelCase					,						sequence_lengths=[8]					,						batch_sizes=[1]					,						multi_process=__lowerCamelCase					,						only_pretrain_model=__lowerCamelCase					,						)
     __A    :						Union[str, Any]						 =     PyTorchBenchmark(__lowerCamelCase  )
     __A    :						Optional[int]						 =     benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						str						 =     '''sshleifer/tiny-gpt2'''
     __A    :						Any						 =     PyTorchBenchmarkArguments(
         models=[MODEL_ID]					,						training=__lowerCamelCase					,						inference=__lowerCamelCase					,						torchscript=__lowerCamelCase					,						sequence_lengths=[8]					,						batch_sizes=[1]					,						multi_process=__lowerCamelCase					,						)
     __A    :						Tuple						 =     PyTorchBenchmark(__lowerCamelCase  )
     __A    :						List[str]						 =     benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )

 @unittest.skipIf(torch_device == '''cpu'''					,						'''Cant do half precision'''  )
 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Dict						 =     '''sshleifer/tiny-gpt2'''
     __A    :						Union[str, Any]						 =     PyTorchBenchmarkArguments(
         models=[MODEL_ID]					,						training=__lowerCamelCase					,						inference=__lowerCamelCase					,						fpaa=__lowerCamelCase					,						sequence_lengths=[8]					,						batch_sizes=[1]					,						multi_process=__lowerCamelCase					,						)
     __A    :						Union[str, Any]						 =     PyTorchBenchmark(__lowerCamelCase  )
     __A    :						str						 =     benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Optional[int]						 =     '''sshleifer/tiny-gpt2'''
     __A    :						int						 =     AutoConfig.from_pretrained(__lowerCamelCase  )
     # set architectures equal to `None`
     __A    :						str						 =     None
     __A    :						Any						 =     PyTorchBenchmarkArguments(
         models=[MODEL_ID]					,						training=__lowerCamelCase					,						inference=__lowerCamelCase					,						sequence_lengths=[8]					,						batch_sizes=[1]					,						multi_process=__lowerCamelCase					,						)
     __A    :						Tuple						 =     PyTorchBenchmark(__lowerCamelCase					,						configs=[config]  )
     __A    :						Optional[Any]						 =     benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						str						 =     '''sshleifer/tiny-gpt2'''
     __A    :						Any						 =     PyTorchBenchmarkArguments(
         models=[MODEL_ID]					,						training=__lowerCamelCase					,						inference=__lowerCamelCase					,						sequence_lengths=[8]					,						batch_sizes=[1]					,						multi_process=__lowerCamelCase					,						)
     __A    :						Optional[int]						 =     PyTorchBenchmark(__lowerCamelCase  )
     __A    :						Union[str, Any]						 =     benchmark.run()
     self.check_results_dict_not_empty(results.time_train_result  )
     self.check_results_dict_not_empty(results.memory_train_result  )

 @unittest.skipIf(torch_device == '''cpu'''					,						'''Can\'t do half precision'''  )
 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						int						 =     '''sshleifer/tiny-gpt2'''
     __A    :						Any						 =     PyTorchBenchmarkArguments(
         models=[MODEL_ID]					,						training=__lowerCamelCase					,						inference=__lowerCamelCase					,						sequence_lengths=[8]					,						batch_sizes=[1]					,						fpaa=__lowerCamelCase					,						multi_process=__lowerCamelCase					,						)
     __A    :						Optional[int]						 =     PyTorchBenchmark(__lowerCamelCase  )
     __A    :						Dict						 =     benchmark.run()
     self.check_results_dict_not_empty(results.time_train_result  )
     self.check_results_dict_not_empty(results.memory_train_result  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Any						 =     '''sshleifer/tiny-gpt2'''
     __A    :						str						 =     AutoConfig.from_pretrained(__lowerCamelCase  )
     __A    :						Tuple						 =     PyTorchBenchmarkArguments(
         models=[MODEL_ID]					,						training=__lowerCamelCase					,						inference=__lowerCamelCase					,						sequence_lengths=[8]					,						batch_sizes=[1]					,						multi_process=__lowerCamelCase					,						)
     __A    :						Optional[int]						 =     PyTorchBenchmark(__lowerCamelCase					,						configs=[config]  )
     __A    :						int						 =     benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						List[str]						 =     '''sshleifer/tinier_bart'''
     __A    :						Optional[Any]						 =     AutoConfig.from_pretrained(__lowerCamelCase  )
     __A    :						str						 =     PyTorchBenchmarkArguments(
         models=[MODEL_ID]					,						training=__lowerCamelCase					,						inference=__lowerCamelCase					,						sequence_lengths=[8]					,						batch_sizes=[1]					,						multi_process=__lowerCamelCase					,						)
     __A    :						List[Any]						 =     PyTorchBenchmark(__lowerCamelCase					,						configs=[config]  )
     __A    :						Tuple						 =     benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result  )
     self.check_results_dict_not_empty(results.memory_inference_result  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						List[str]						 =     '''sshleifer/tiny-gpt2'''
     __A    :						str						 =     AutoConfig.from_pretrained(__lowerCamelCase  )
     __A    :						Dict						 =     PyTorchBenchmarkArguments(
         models=[MODEL_ID]					,						training=__lowerCamelCase					,						inference=__lowerCamelCase					,						sequence_lengths=[8]					,						batch_sizes=[1]					,						multi_process=__lowerCamelCase					,						)
     __A    :						List[Any]						 =     PyTorchBenchmark(__lowerCamelCase					,						configs=[config]  )
     __A    :						str						 =     benchmark.run()
     self.check_results_dict_not_empty(results.time_train_result  )
     self.check_results_dict_not_empty(results.memory_train_result  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						str						 =     '''sshleifer/tinier_bart'''
     __A    :						Optional[Any]						 =     AutoConfig.from_pretrained(__lowerCamelCase  )
     __A    :						Optional[int]						 =     PyTorchBenchmarkArguments(
         models=[MODEL_ID]					,						training=__lowerCamelCase					,						inference=__lowerCamelCase					,						sequence_lengths=[8]					,						batch_sizes=[1]					,						multi_process=__lowerCamelCase					,						)
     __A    :						Optional[Any]						 =     PyTorchBenchmark(__lowerCamelCase					,						configs=[config]  )
     __A    :						List[str]						 =     benchmark.run()
     self.check_results_dict_not_empty(results.time_train_result  )
     self.check_results_dict_not_empty(results.memory_train_result  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						List[Any]						 =     '''sshleifer/tiny-gpt2'''
     with tempfile.TemporaryDirectory() as tmp_dir:
         __A    :						List[str]						 =     PyTorchBenchmarkArguments(
             models=[MODEL_ID]					,						training=__lowerCamelCase					,						inference=__lowerCamelCase					,						save_to_csv=__lowerCamelCase					,						sequence_lengths=[8]					,						batch_sizes=[1]					,						inference_time_csv_file=os.path.join(__lowerCamelCase					,						'''inf_time.csv'''  )					,						train_memory_csv_file=os.path.join(__lowerCamelCase					,						'''train_mem.csv'''  )					,						inference_memory_csv_file=os.path.join(__lowerCamelCase					,						'''inf_mem.csv'''  )					,						train_time_csv_file=os.path.join(__lowerCamelCase					,						'''train_time.csv'''  )					,						env_info_csv_file=os.path.join(__lowerCamelCase					,						'''env.csv'''  )					,						multi_process=__lowerCamelCase					,						)
         __A    :						int						 =     PyTorchBenchmark(__lowerCamelCase  )
         benchmark.run()
         self.assertTrue(Path(os.path.join(__lowerCamelCase					,						'''inf_time.csv'''  )  ).exists()  )
         self.assertTrue(Path(os.path.join(__lowerCamelCase					,						'''train_time.csv'''  )  ).exists()  )
         self.assertTrue(Path(os.path.join(__lowerCamelCase					,						'''inf_mem.csv'''  )  ).exists()  )
         self.assertTrue(Path(os.path.join(__lowerCamelCase					,						'''train_mem.csv'''  )  ).exists()  )
         self.assertTrue(Path(os.path.join(__lowerCamelCase					,						'''env.csv'''  )  ).exists()  )







 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Dict						 =     '''sshleifer/tiny-gpt2'''

     def _check_summary_is_not_empty(__lowerCamelCase  ):
         self.assertTrue(hasattr(__lowerCamelCase					,						'''sequential'''  )  )
         self.assertTrue(hasattr(__lowerCamelCase					,						'''cumulative'''  )  )
         self.assertTrue(hasattr(__lowerCamelCase					,						'''current'''  )  )
         self.assertTrue(hasattr(__lowerCamelCase					,						'''total'''  )  )

     with tempfile.TemporaryDirectory() as tmp_dir:
         __A    :						int						 =     PyTorchBenchmarkArguments(
             models=[MODEL_ID]					,						training=__lowerCamelCase					,						inference=__lowerCamelCase					,						sequence_lengths=[8]					,						batch_sizes=[1]					,						log_filename=os.path.join(__lowerCamelCase					,						'''log.txt'''  )					,						log_print=__lowerCamelCase					,						trace_memory_line_by_line=__lowerCamelCase					,						multi_process=__lowerCamelCase					,						)
         __A    :						str						 =     PyTorchBenchmark(__lowerCamelCase  )
         __A    :						Dict						 =     benchmark.run()
         _check_summary_is_not_empty(result.inference_summary  )
         _check_summary_is_not_empty(result.train_summary  )
         self.assertTrue(Path(os.path.join(__lowerCamelCase					,						'''log.txt'''  )  ).exists()  )











"""simple docstring"""



from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


a_ =			logging.get_logger(__name__)

a_ =			{
    """shi-labs/nat-mini-in1k-224""": """https://huggingface.co/shi-labs/nat-mini-in1k-224/resolve/main/config.json""",
    # See all Nat models at https://huggingface.co/models?filter=nat
}






class 		__snake_case					(		SCREAMING_SNAKE_CASE__							,							SCREAMING_SNAKE_CASE__     ):





 """simple docstring"""


 _lowerCamelCase        =     """nat"""

 _lowerCamelCase        =     {
     """num_attention_heads""": """num_heads""",
     """num_hidden_layers""": """num_layers""",
 }

 def __init__(    self					,						__lowerCamelCase=4					,						__lowerCamelCase=3					,						__lowerCamelCase=64					,						__lowerCamelCase=[3, 4, 6, 5]					,						__lowerCamelCase=[2, 4, 8, 16]					,						__lowerCamelCase=7					,						__lowerCamelCase=3.0					,						__lowerCamelCase=True					,						__lowerCamelCase=0.0					,						__lowerCamelCase=0.0					,						__lowerCamelCase=0.1					,						__lowerCamelCase="gelu"					,						__lowerCamelCase=0.0_2					,						__lowerCamelCase=1e-5					,						__lowerCamelCase=0.0					,						__lowerCamelCase=None					,						__lowerCamelCase=None					,						**__lowerCamelCase					,						):



     '''simple docstring'''




     super().__init__(**__lowerCamelCase  )

     __A    :						Union[str, Any]						 =     patch_size
     __A    :						Optional[Any]						 =     num_channels
     __A    :						Tuple						 =     embed_dim
     __A    :						Dict						 =     depths
     __A    :						str						 =     len(__lowerCamelCase  )
     __A    :						Optional[Any]						 =     num_heads
     __A    :						str						 =     kernel_size
     __A    :						Any						 =     mlp_ratio
     __A    :						Optional[int]						 =     qkv_bias
     __A    :						str						 =     hidden_dropout_prob
     __A    :						Any						 =     attention_probs_dropout_prob
     __A    :						int						 =     drop_path_rate
     __A    :						int						 =     hidden_act
     __A    :						Any						 =     layer_norm_eps
     __A    :						Tuple						 =     initializer_range
     # we set the hidden_size attribute in order to make Nat work with VisionEncoderDecoderModel
     # this indicates the channel dimension after the last stage of the model
     __A    :						int						 =     int(embed_dim * 2 ** (len(__lowerCamelCase  ) - 1)  )
     __A    :						Union[str, Any]						 =     layer_scale_init_value
     __A    :						List[str]						 =     ['''stem'''] + [F"""stage{idx}""" for idx in range(1					,						len(__lowerCamelCase  ) + 1  )]
     __A ,       __A    :						Any						 =     get_aligned_output_features_output_indices(
         out_features=__lowerCamelCase					,						out_indices=__lowerCamelCase					,						stage_names=self.stage_names  )




