





import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf


if is_tf_available():
							import tensorflow as tf
							from tensorflow.python.eager import context
							from tensorflow.python.framework import ops

							from transformers import GradientAccumulator, create_optimizer





@require_tf
class 				a			(      unittest.TestCase			):






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										self.assertEqual(len(_lowerCAmelCase					)			, len(_lowerCAmelCase					)					)
										for a, b in zip(_lowerCAmelCase			, _lowerCAmelCase					):
																self.assertAlmostEqual(_lowerCAmelCase			, _lowerCAmelCase			, delta=_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   GradientAccumulator()
										accumulator([tf.constant([1.0, 2.0]					)]					)
										accumulator([tf.constant([-2.0, 1.0]					)]					)
										accumulator([tf.constant([-1.0, 2.0]					)]					)
										with self.assertRaises(_lowerCAmelCase					):
																accumulator([tf.constant([1.0, 1.0]					), tf.constant([2.0, 2.0]					)]					)
										self.assertEqual(accumulator.step			, 3					)
										self.assertEqual(len(accumulator.gradients					)			, 1					)
										self.assertListAlmostEqual(accumulator.gradients[0].numpy().tolist()			, [-2.0, 5.0]			, tol=1e-2					)
										accumulator.reset()
										self.assertEqual(accumulator.step			, 0					)
										self.assertListAlmostEqual(accumulator.gradients[0].numpy().tolist()			, [0.0, 0.0]			, tol=1e-2					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   None
										ops.enable_eager_execution_internal()
										__SCREAMING_SNAKE_CASE:				List[str]           =   tf.config.list_physical_devices('''CPU'''					)
										if len(_lowerCAmelCase					) == 1:
																tf.config.set_logical_device_configuration(
																    physical_devices[0]			, [tf.config.LogicalDeviceConfiguration(), tf.config.LogicalDeviceConfiguration()]					)
										__SCREAMING_SNAKE_CASE:				Dict           =   tf.config.list_logical_devices(device_type='''CPU'''					)
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   tf.distribute.MirroredStrategy(devices=devices[:2]					)

										with strategy.scope():
																__SCREAMING_SNAKE_CASE:				Tuple           =   GradientAccumulator()
																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   tf.Variable([4.0, 3.0]					)
																__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE:				List[Any]           =   create_optimizer(5e-5			, 10			, 5					)
																__SCREAMING_SNAKE_CASE:				int           =   tf.Variable([0.0, 0.0]			, trainable=_lowerCAmelCase					)

										def accumulate_on_replica(_lowerCAmelCase					):
																accumulator([gradient]					)

										def apply_on_replica():
																optimizer.apply_gradients(list(zip(accumulator.gradients			, [variable]					)					)					)

										@tf.function
										def accumulate(_lowerCAmelCase			, _lowerCAmelCase					):
																with strategy.scope():
																						__SCREAMING_SNAKE_CASE:				Tuple           =   strategy.experimental_local_results(_lowerCAmelCase					)
																						local_variables[0].assign(_lowerCAmelCase					)
																						local_variables[1].assign(_lowerCAmelCase					)
																						strategy.run(_lowerCAmelCase			, args=(gradient_placeholder,)					)

										@tf.function
										def apply_grad():
																with strategy.scope():
																						strategy.run(_lowerCAmelCase					)

										def _check_local_values(_lowerCAmelCase			, _lowerCAmelCase					):
																__SCREAMING_SNAKE_CASE:				int           =   strategy.experimental_local_results(accumulator._gradients[0]					)
																self.assertListAlmostEqual(values[0].value()			, _lowerCAmelCase			, tol=1e-2					)
																self.assertListAlmostEqual(values[1].value()			, _lowerCAmelCase			, tol=1e-2					)

										accumulate([1.0, 2.0]			, [-1.0, 1.0]					)
										accumulate([3.0, -1.0]			, [-1.0, -1.0]					)
										accumulate([-2.0, 2.0]			, [3.0, -2.0]					)
										self.assertEqual(accumulator.step			, 3					)
										_check_local_values([2.0, 3.0]			, [1.0, -2.0]					)
										apply_grad()
										self.assertListAlmostEqual(variable.value()			, [4.0, 3.0]			, tol=1e-2					)
										accumulator.reset()
										self.assertEqual(accumulator.step			, 0					)
										_check_local_values([0.0, 0.0]			, [0.0, 0.0]					)










import unittest
from dataclasses import dataclass

import pytest

from accelerate.commands.config.config_args import SageMakerConfig
from accelerate.utils import ComputeEnvironment
from accelerate.utils.launch import _convert_nargs_to_dict





@dataclass
class 				a			(      __lowercase			):
				SCREAMING_SNAKE_CASE__     :   Any		      =     ComputeEnvironment.AMAZON_SAGEMAKER
				SCREAMING_SNAKE_CASE__     :   Optional[int]		      =     True
				SCREAMING_SNAKE_CASE__     :   int		      =     '''ml.p3.2xlarge'''
				SCREAMING_SNAKE_CASE__     :   Union[str, Any]		      =     '''accelerate_sagemaker_execution_role'''
				SCREAMING_SNAKE_CASE__     :   List[Any]		      =     '''hf-sm'''
				SCREAMING_SNAKE_CASE__     :   Any		      =     '''us-east-1'''
				SCREAMING_SNAKE_CASE__     :   Any		      =     1
				SCREAMING_SNAKE_CASE__     :   Dict		      =     '''accelerate-sagemaker-1'''
				SCREAMING_SNAKE_CASE__     :   Dict		      =     '''1.6'''
				SCREAMING_SNAKE_CASE__     :   int		      =     '''4.4'''
				SCREAMING_SNAKE_CASE__     :   int		      =     '''train.py'''
				SCREAMING_SNAKE_CASE__     :   Tuple		      =     [
				    '''--model_name_or_path''',
				    '''bert''',
				    '''--do_train''',
				    '''False''',
				    '''--epochs''',
				    '''3''',
				    '''--learning_rate''',
				    '''5e-5''',
				    '''--max_steps''',
				    '''50.5''',
				]
				SCREAMING_SNAKE_CASE__     :   List[str]		      =     [
				    '''--model_name_or_path''',
				    '''bert''',
				    '''--do_train''',
				    '''--do_test''',
				    '''False''',
				    '''--do_predict''',
				    '''--epochs''',
				    '''3''',
				    '''--learning_rate''',
				    '''5e-5''',
				    '''--max_steps''',
				    '''50.5''',
				]





class 				a			(      unittest.TestCase			):






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   _convert_nargs_to_dict(MockLaunchConfig.success_training_script_args					)
										assert isinstance(converted_args['''model_name_or_path''']			, _lowerCAmelCase					)
										assert isinstance(converted_args['''do_train''']			, _lowerCAmelCase					)
										assert isinstance(converted_args['''epochs''']			, _lowerCAmelCase					)
										assert isinstance(converted_args['''learning_rate''']			, _lowerCAmelCase					)
										assert isinstance(converted_args['''max_steps''']			, _lowerCAmelCase					)

										with pytest.raises(_lowerCAmelCase					):
																_convert_nargs_to_dict(MockLaunchConfig.fail_training_script_args					)



