"""simple docstring"""







# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import os

from ...utils.constants import SAGEMAKER_PARALLEL_EC2_INSTANCES, TORCH_DYNAMO_MODES
from ...utils.dataclasses import ComputeEnvironment, SageMakerDistributedType
from ...utils.imports import is_botoa_available
from .config_args import SageMakerConfig
from .config_utils import (
    DYNAMO_BACKENDS,
    _ask_field,
    _ask_options,
    _convert_dynamo_backend,
    _convert_mixed_precision,
    _convert_sagemaker_distributed_mode,
    _convert_yes_no_to_bool,
)


if is_botoa_available():
       import botoa  # noqa: F401



def 			lowercase__(       __SCREAMING_SNAKE_CASE							:     List[str]      ):
     lowercase_	:				Tuple													=	botoa.client('iam'      )

     lowercase_	:				List[Any]													=	{
         'Version': '2012-10-17',
         'Statement': [
             {'Effect': 'Allow', 'Principal': {'Service': 'sagemaker.amazonaws.com'}, 'Action': 'sts:AssumeRole'}
         ],
     }
     try:
          # create the role, associated with the chosen trust policy
          iam_client.create_role(
              RoleName=__SCREAMING_SNAKE_CASE     ,			AssumeRolePolicyDocument=json.dumps(__SCREAMING_SNAKE_CASE     ,			indent=2      )      )
          lowercase_	:				List[str]													=	{
              'Version': '2012-10-17',
              'Statement': [
                  {
                      'Effect': 'Allow',
                      'Action': [
                          'sagemaker:*',
                          'ecr:GetDownloadUrlForLayer',
                          'ecr:BatchGetImage',
                          'ecr:BatchCheckLayerAvailability',
                          'ecr:GetAuthorizationToken',
                          'cloudwatch:PutMetricData',
                          'cloudwatch:GetMetricData',
                          'cloudwatch:GetMetricStatistics',
                          'cloudwatch:ListMetrics',
                          'logs:CreateLogGroup',
                          'logs:CreateLogStream',
                          'logs:DescribeLogStreams',
                          'logs:PutLogEvents',
                          'logs:GetLogEvents',
                          's3:CreateBucket',
                          's3:ListBucket',
                          's3:GetBucketLocation',
                          's3:GetObject',
                          's3:PutObject',
                      ],
                      'Resource': '*',
                  }
              ],
          }
          # attach policy to role
          iam_client.put_role_policy(
              RoleName=__SCREAMING_SNAKE_CASE     ,			PolicyName=F'''{role_name}_policy_permission'''     ,			PolicyDocument=json.dumps(__SCREAMING_SNAKE_CASE     ,			indent=2      )     ,			)
     except iam_client.exceptions.EntityAlreadyExistsException:
          print(F'''role {role_name} already exists. Using existing one'''      )



def 			lowercase__(       __SCREAMING_SNAKE_CASE							:     Optional[Any]      ):
     lowercase_	:				int													=	botoa.client('iam'      )
     return iam_client.get_role(RoleName=__SCREAMING_SNAKE_CASE      )["Role"]["Arn"]







def 			lowercase__(       ):
     lowercase_	:				Union[str, Any]													=	_ask_options(
         'How do you want to authorize?'     ,			['AWS Profile', 'Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) ']     ,			__SCREAMING_SNAKE_CASE     ,			)
     lowercase_	:				Any													=	None
     if credentials_configuration == 0:
          lowercase_	:				Dict													=	_ask_field('Enter your AWS Profile name: [default] '     ,			default='default'      )
          lowercase_	:				Dict													=	aws_profile
     else:
          print(
              'Note you will need to provide AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY when you launch you training script with,'
              '`accelerate launch --aws_access_key_id XXX --aws_secret_access_key YYY`'      )
          lowercase_	:				Tuple													=	_ask_field('AWS Access Key ID: '      )
          lowercase_	:				Any													=	aws_access_key_id

          lowercase_	:				Optional[int]													=	_ask_field('AWS Secret Access Key: '      )
          lowercase_	:				Optional[int]													=	aws_secret_access_key

     lowercase_	:				Union[str, Any]													=	_ask_field('Enter your AWS Region: [us-east-1]'     ,			default='us-east-1'      )
     lowercase_	:				Union[str, Any]													=	aws_region

     lowercase_	:				int													=	_ask_options(
         'Do you already have an IAM Role for executing Amazon SageMaker Training Jobs?'     ,			['Provide IAM Role name', 'Create new IAM role using credentials']     ,			__SCREAMING_SNAKE_CASE     ,			)
     if role_management == 0:
          lowercase_	:				int													=	_ask_field('Enter your IAM role name: '      )
     else:
          lowercase_	:				List[Any]													=	'accelerate_sagemaker_execution_role'
          print(F'''Accelerate will create an iam role "{iam_role_name}" using the provided credentials'''      )
          _create_iam_role_for_sagemaker(__SCREAMING_SNAKE_CASE      )

     lowercase_	:				Any													=	_ask_field(
         'Do you want to use custom Docker image? [yes/NO]: '     ,			_convert_yes_no_to_bool     ,			default=__SCREAMING_SNAKE_CASE     ,			error_message='Please enter yes or no.'     ,			)
     lowercase_	:				Optional[int]													=	None
     if is_custom_docker_image:
          lowercase_	:				Tuple													=	_ask_field('Enter your Docker image: '     ,			lambda __SCREAMING_SNAKE_CASE      : str(__SCREAMING_SNAKE_CASE      ).lower()      )

     lowercase_	:				Optional[Any]													=	_ask_field(
         'Do you want to provide SageMaker input channels with data locations? [yes/NO]: '     ,			_convert_yes_no_to_bool     ,			default=__SCREAMING_SNAKE_CASE     ,			error_message='Please enter yes or no.'     ,			)
     lowercase_	:				Any													=	None
     if is_sagemaker_inputs_enabled:
          lowercase_	:				List[str]													=	_ask_field(
              'Enter the path to the SageMaker inputs TSV file with columns (channel_name, data_location): '     ,			lambda __SCREAMING_SNAKE_CASE      : str(__SCREAMING_SNAKE_CASE      ).lower()     ,			)

     lowercase_	:				Dict													=	_ask_field(
         'Do you want to enable SageMaker metrics? [yes/NO]: '     ,			_convert_yes_no_to_bool     ,			default=__SCREAMING_SNAKE_CASE     ,			error_message='Please enter yes or no.'     ,			)
     lowercase_	:				Any													=	None
     if is_sagemaker_metrics_enabled:
          lowercase_	:				int													=	_ask_field(
              'Enter the path to the SageMaker metrics TSV file with columns (metric_name, metric_regex): '     ,			lambda __SCREAMING_SNAKE_CASE      : str(__SCREAMING_SNAKE_CASE      ).lower()     ,			)

     lowercase_	:				Optional[int]													=	_ask_options(
         'What is the distributed mode?'     ,			['No distributed training', 'Data parallelism']     ,			_convert_sagemaker_distributed_mode     ,			)
     lowercase_	:				Optional[int]													=	{}
     lowercase_	:				Any													=	_ask_field(
         'Do you wish to optimize your script with torch dynamo?[yes/NO]:'     ,			_convert_yes_no_to_bool     ,			default=__SCREAMING_SNAKE_CASE     ,			error_message='Please enter yes or no.'     ,			)
     if use_dynamo:
          lowercase_	:				Union[str, Any]													=	'dynamo_'
          lowercase_	:				int													=	_ask_options(
              'Which dynamo backend would you like to use?'     ,			[x.lower() for x in DYNAMO_BACKENDS]     ,			_convert_dynamo_backend     ,			default=2     ,			)
          lowercase_	:				Optional[int]													=	_ask_field(
              'Do you want to customize the defaults sent to torch.compile? [yes/NO]: '     ,			_convert_yes_no_to_bool     ,			default=__SCREAMING_SNAKE_CASE     ,			error_message='Please enter yes or no.'     ,			)

          if use_custom_options:
               lowercase_	:				List[str]													=	_ask_options(
                   'Which mode do you want to use?'     ,			__SCREAMING_SNAKE_CASE     ,			lambda __SCREAMING_SNAKE_CASE      : TORCH_DYNAMO_MODES[int(__SCREAMING_SNAKE_CASE      )]     ,			default='default'     ,			)
               lowercase_	:				Union[str, Any]													=	_ask_field(
                   'Do you want the fullgraph mode or it is ok to break model into several subgraphs? [yes/NO]: '     ,			_convert_yes_no_to_bool     ,			default=__SCREAMING_SNAKE_CASE     ,			error_message='Please enter yes or no.'     ,			)
               lowercase_	:				Optional[Any]													=	_ask_field(
                   'Do you want to enable dynamic shape tracing? [yes/NO]: '     ,			_convert_yes_no_to_bool     ,			default=__SCREAMING_SNAKE_CASE     ,			error_message='Please enter yes or no.'     ,			)
     lowercase_	:				Any													=	'Which EC2 instance type you want to use for your training?'
     if distributed_type != SageMakerDistributedType.NO:
          lowercase_	:				int													=	_ask_options(
              __SCREAMING_SNAKE_CASE     ,			__SCREAMING_SNAKE_CASE     ,			lambda __SCREAMING_SNAKE_CASE      : SAGEMAKER_PARALLEL_EC2_INSTANCES[int(__SCREAMING_SNAKE_CASE      )]      )
     else:
          eca_instance_query += "? [ml.p3.2xlarge]:"
          lowercase_	:				int													=	_ask_field(__SCREAMING_SNAKE_CASE     ,			lambda __SCREAMING_SNAKE_CASE      : str(__SCREAMING_SNAKE_CASE      ).lower()     ,			default='ml.p3.2xlarge'      )

     lowercase_	:				List[str]													=	1
     if distributed_type in (SageMakerDistributedType.DATA_PARALLEL, SageMakerDistributedType.MODEL_PARALLEL):
          lowercase_	:				Dict													=	_ask_field(
              'How many machines do you want use? [1]: '     ,			__SCREAMING_SNAKE_CASE     ,			default=1     ,			)

     lowercase_	:				Union[str, Any]													=	_ask_options(
         'Do you wish to use FP16 or BF16 (mixed precision)?'     ,			['no', 'fp16', 'bf16', 'fp8']     ,			_convert_mixed_precision     ,			)

     if use_dynamo and mixed_precision == "no":
          print(
              'Torch dynamo used without mixed precision requires TF32 to be efficient. Accelerate will enable it by default when launching your scripts.'      )

     return SageMakerConfig(
         image_uri=__SCREAMING_SNAKE_CASE     ,			compute_environment=ComputeEnvironment.AMAZON_SAGEMAKER     ,			distributed_type=__SCREAMING_SNAKE_CASE     ,			use_cpu=__SCREAMING_SNAKE_CASE     ,			dynamo_config=__SCREAMING_SNAKE_CASE     ,			eca_instance_type=__SCREAMING_SNAKE_CASE     ,			profile=__SCREAMING_SNAKE_CASE     ,			region=__SCREAMING_SNAKE_CASE     ,			iam_role_name=__SCREAMING_SNAKE_CASE     ,			mixed_precision=__SCREAMING_SNAKE_CASE     ,			num_machines=__SCREAMING_SNAKE_CASE     ,			sagemaker_inputs_file=__SCREAMING_SNAKE_CASE     ,			sagemaker_metrics_file=__SCREAMING_SNAKE_CASE     ,			)


"""simple docstring"""







import inspect
import unittest

import numpy as np

from transformers import ViTConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
       import jax

       from transformers.models.vit.modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel







class        UpperCamelCase		(			unittest.TestCase  ):
   def __init__(			self     ,__UpperCamelCase     ,__UpperCamelCase=13     ,__UpperCamelCase=30     ,__UpperCamelCase=2     ,__UpperCamelCase=3     ,__UpperCamelCase=True     ,__UpperCamelCase=True     ,__UpperCamelCase=32     ,__UpperCamelCase=5     ,__UpperCamelCase=4     ,__UpperCamelCase=37     ,__UpperCamelCase="gelu"     ,__UpperCamelCase=0.1     ,__UpperCamelCase=0.1     ,__UpperCamelCase=10     ,__UpperCamelCase=0.02     ,)	->      Tuple:


        '''simple docstring'''


        lowercase_	:				Tuple													=	parent
        lowercase_	:				Union[str, Any]													=	batch_size
        lowercase_	:				int													=	image_size
        lowercase_	:				Tuple													=	patch_size
        lowercase_	:				Optional[int]													=	num_channels
        lowercase_	:				Union[str, Any]													=	is_training
        lowercase_	:				Dict													=	use_labels
        lowercase_	:				Optional[int]													=	hidden_size
        lowercase_	:				List[str]													=	num_hidden_layers
        lowercase_	:				Optional[Any]													=	num_attention_heads
        lowercase_	:				Optional[int]													=	intermediate_size
        lowercase_	:				Tuple													=	hidden_act
        lowercase_	:				int													=	hidden_dropout_prob
        lowercase_	:				str													=	attention_probs_dropout_prob
        lowercase_	:				str													=	type_sequence_label_size
        lowercase_	:				Optional[int]													=	initializer_range

        # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
        lowercase_	:				str													=	(image_size // patch_size) ** 2
        lowercase_	:				Optional[int]													=	num_patches + 1
   def 			_UpperCAmelCase   (			self			)	->      Tuple:


        '''simple docstring'''


        lowercase_	:				Tuple													=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

        lowercase_	:				List[Any]													=	ViTConfig(
            image_size=self.image_size     ,patch_size=self.patch_size     ,num_channels=self.num_channels     ,hidden_size=self.hidden_size     ,num_hidden_layers=self.num_hidden_layers     ,num_attention_heads=self.num_attention_heads     ,intermediate_size=self.intermediate_size     ,hidden_act=self.hidden_act     ,hidden_dropout_prob=self.hidden_dropout_prob     ,attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,is_decoder=__UpperCamelCase     ,initializer_range=self.initializer_range     ,)

        return config, pixel_values
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase			)	->      Dict:


        '''simple docstring'''


        lowercase_	:				List[Any]													=	FlaxViTModel(config=__UpperCamelCase			)
        lowercase_	:				Dict													=	model(__UpperCamelCase			)
        # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
        lowercase_	:				Union[str, Any]													=	(self.image_size, self.image_size)
        lowercase_	:				List[Any]													=	(self.patch_size, self.patch_size)
        lowercase_	:				List[str]													=	(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
        self.parent.assertEqual(result.last_hidden_state.shape     ,(self.batch_size, num_patches + 1, self.hidden_size)			)
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase			)	->      Union[str, Any]:


        '''simple docstring'''


        lowercase_	:				List[Any]													=	self.type_sequence_label_size
        lowercase_	:				str													=	FlaxViTForImageClassification(config=__UpperCamelCase			)
        lowercase_	:				Optional[Any]													=	model(__UpperCamelCase			)
        self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.type_sequence_label_size)			)

        # test greyscale images
        lowercase_	:				Union[str, Any]													=	1
        lowercase_	:				Optional[int]													=	FlaxViTForImageClassification(__UpperCamelCase			)

        lowercase_	:				List[Any]													=	floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
        lowercase_	:				str													=	model(__UpperCamelCase			)







   def 			_UpperCAmelCase   (			self			)	->      Any:


        '''simple docstring'''


        lowercase_	:				Optional[int]													=	self.prepare_config_and_inputs()
        (
            (
            lowercase_
        )							,  (
            lowercase_
        )							,  
        )	:				List[Any]													=	config_and_inputs
        lowercase_	:				Tuple													=	{'pixel_values': pixel_values}
        return config, inputs_dict







@require_flax
class        UpperCamelCase		(			lowercase_						,      unittest.TestCase  ):
   lowercase							=  (FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()
   def 			_UpperCAmelCase   (			self			)	->      None:


        '''simple docstring'''


        lowercase_	:				Optional[Any]													=	FlaxViTModelTester(self			)
        lowercase_	:				Union[str, Any]													=	ConfigTester(self     ,config_class=__UpperCamelCase     ,has_text_modality=__UpperCamelCase     ,hidden_size=37			)
   def 			_UpperCAmelCase   (			self			)	->      List[str]:


        '''simple docstring'''


        self.config_tester.run_common_tests()
   def 			_UpperCAmelCase   (			self			)	->      Optional[Any]:


        '''simple docstring'''


        lowercase_	:				Optional[Any]													=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__UpperCamelCase			)
   def 			_UpperCAmelCase   (			self			)	->      Optional[Any]:


        '''simple docstring'''


        lowercase_	:				List[Any]													=	self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*__UpperCamelCase			)
   def 			_UpperCAmelCase   (			self			)	->      List[Any]:


        '''simple docstring'''


        lowercase_							,  lowercase_	:				Any													=	self.model_tester.prepare_config_and_inputs_for_common()

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

             lowercase_	:				str													=	['pixel_values']
             self.assertListEqual(arg_names[:1]     ,__UpperCamelCase			)
   def 			_UpperCAmelCase   (			self			)	->      Optional[Any]:


        '''simple docstring'''


        lowercase_							,  lowercase_	:				str													=	self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
             with self.subTest(model_class.__name__			):
                  lowercase_	:				Optional[Any]													=	self._prepare_for_class(__UpperCamelCase     ,__UpperCamelCase			)
                  lowercase_	:				Dict													=	model_class(__UpperCamelCase			)

                  @jax.jit
                  def model_jitted(__UpperCamelCase     ,**__UpperCamelCase			):
                       return model(pixel_values=__UpperCamelCase     ,**__UpperCamelCase			)

                  with self.subTest('JIT Enabled'			):
                       lowercase_	:				Optional[int]													=	model_jitted(**__UpperCamelCase			).to_tuple()

                  with self.subTest('JIT Disabled'			):
                       with jax.disable_jit():
                            lowercase_	:				List[str]													=	model_jitted(**__UpperCamelCase			).to_tuple()

                  self.assertEqual(len(__UpperCamelCase			)     ,len(__UpperCamelCase			)			)
                  for jitted_output, output in zip(__UpperCamelCase     ,__UpperCamelCase			):
                       self.assertEqual(jitted_output.shape     ,output.shape			)







   @slow
   def 			_UpperCAmelCase   (			self			)	->      Optional[int]:


        '''simple docstring'''


        for model_class_name in self.all_model_classes:
             lowercase_	:				Optional[int]													=	model_class_name.from_pretrained('google/vit-base-patch16-224'			)
             lowercase_	:				int													=	model(np.ones((1, 3, 224, 224)			)			)
             self.assertIsNotNone(__UpperCamelCase			)

