import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import (
          MODEL_FOR_PRETRAINING_MAPPING,
          MegatronBertForCausalLM,
          MegatronBertForMaskedLM,
          MegatronBertForMultipleChoice,
          MegatronBertForNextSentencePrediction,
          MegatronBertForPreTraining,
          MegatronBertForQuestionAnswering,
          MegatronBertForSequenceClassification,
          MegatronBertForTokenClassification,
          MegatronBertModel,
      )





class      A							:


   def __init__(self     :						Optional[Any]       ,  __UpperCAmelCase     :						Union[str, Any]       ,  __UpperCAmelCase     :						Optional[Any]=1_3       ,  __UpperCAmelCase     :						Optional[int]=7       ,  __UpperCAmelCase     :						Optional[int]=True       ,  __UpperCAmelCase     :						Optional[Any]=True       ,  __UpperCAmelCase     :						Optional[Any]=True       ,  __UpperCAmelCase     :						Optional[int]=True       ,  __UpperCAmelCase     :						Union[str, Any]=9_9       ,  __UpperCAmelCase     :						Union[str, Any]=6_4       ,  __UpperCAmelCase     :						Dict=3_2       ,  __UpperCAmelCase     :						int=5       ,  __UpperCAmelCase     :						Optional[int]=4       ,  __UpperCAmelCase     :						Optional[int]=3_7       ,  __UpperCAmelCase     :						List[str]="gelu"       ,  __UpperCAmelCase     :						str=0.1       ,  __UpperCAmelCase     :						Union[str, Any]=0.1       ,  __UpperCAmelCase     :						List[str]=5_1_2       ,  __UpperCAmelCase     :						str=1_6       ,  __UpperCAmelCase     :						List[Any]=2       ,  __UpperCAmelCase     :						List[str]=0.02       ,  __UpperCAmelCase     :						Optional[int]=3       ,  __UpperCAmelCase     :						Any=4       ,  __UpperCAmelCase     :						Tuple=None       ,  )						->				List[str]:
         """simple docstring"""



         UpperCAmelCase__								=  parent
         UpperCAmelCase__								=  batch_size
         UpperCAmelCase__								=  seq_length
         UpperCAmelCase__								=  is_training
         UpperCAmelCase__								=  use_input_mask
         UpperCAmelCase__								=  use_token_type_ids
         UpperCAmelCase__								=  use_labels
         UpperCAmelCase__								=  vocab_size
         UpperCAmelCase__								=  hidden_size
         UpperCAmelCase__								=  embedding_size
         UpperCAmelCase__								=  num_hidden_layers
         UpperCAmelCase__								=  num_attention_heads
         UpperCAmelCase__								=  intermediate_size
         UpperCAmelCase__								=  hidden_act
         UpperCAmelCase__								=  hidden_dropout_prob
         UpperCAmelCase__								=  attention_probs_dropout_prob
         UpperCAmelCase__								=  max_position_embeddings
         UpperCAmelCase__								=  type_vocab_size
         UpperCAmelCase__								=  type_sequence_label_size
         UpperCAmelCase__								=  initializer_range
         UpperCAmelCase__								=  num_labels
         UpperCAmelCase__								=  num_choices
         UpperCAmelCase__								=  scope


   def 			lowercase_      (self     :						Optional[int]		)						->				Any:
         """simple docstring"""



         UpperCAmelCase__								=  ids_tensor([self.batch_size, self.seq_length]       ,  self.vocab_size		)

         UpperCAmelCase__								=  None
         if self.use_input_mask:
               UpperCAmelCase__								=  random_attention_mask([self.batch_size, self.seq_length]		)

         UpperCAmelCase__								=  None
         if self.use_token_type_ids:
               UpperCAmelCase__								=  ids_tensor([self.batch_size, self.seq_length]       ,  self.type_vocab_size		)

         UpperCAmelCase__								=  None
         UpperCAmelCase__								=  None
         UpperCAmelCase__								=  None
         if self.use_labels:
               UpperCAmelCase__								=  ids_tensor([self.batch_size]       ,  self.type_sequence_label_size		)
               UpperCAmelCase__								=  ids_tensor([self.batch_size, self.seq_length]       ,  self.num_labels		)
               UpperCAmelCase__								=  ids_tensor([self.batch_size]       ,  self.num_choices		)

         UpperCAmelCase__								=  self.get_config()

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


   def 			lowercase_      (self     :						int		)						->				Dict:
         """simple docstring"""



         return MegatronBertConfig(
             vocab_size=self.vocab_size       ,  hidden_size=self.hidden_size       ,  num_hidden_layers=self.num_hidden_layers       ,  num_attention_heads=self.num_attention_heads       ,  intermediate_size=self.intermediate_size       ,  embedding_size=self.embedding_size       ,  hidden_act=self.hidden_act       ,  hidden_dropout_prob=self.hidden_dropout_prob       ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,  max_position_embeddings=self.max_position_embeddings       ,  type_vocab_size=self.type_vocab_size       ,  is_decoder=lowerCAmelCase__       ,  initializer_range=self.initializer_range       ,  )


   def 			lowercase_      (self     :						int       ,  __UpperCAmelCase     :						int       ,  __UpperCAmelCase     :						Any       ,  __UpperCAmelCase     :						Optional[Any]       ,  __UpperCAmelCase     :						Any       ,  __UpperCAmelCase     :						int       ,  __UpperCAmelCase     :						Optional[int]       ,  __UpperCAmelCase     :						Optional[Any]		)						->				List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  MegatronBertModel(config=lowerCAmelCase__		)
         model.to(lowerCAmelCase__		)
         model.eval()
         UpperCAmelCase__								=  model(lowerCAmelCase__       ,  attention_mask=lowerCAmelCase__       ,  token_type_ids=lowerCAmelCase__		)
         UpperCAmelCase__								=  model(lowerCAmelCase__       ,  token_type_ids=lowerCAmelCase__		)
         UpperCAmelCase__								=  model(lowerCAmelCase__		)

         self.parent.assertEqual(result.last_hidden_state.shape       ,  (self.batch_size, self.seq_length, self.hidden_size)		)
         self.parent.assertEqual(result.pooler_output.shape       ,  (self.batch_size, self.hidden_size)		)


   def 			lowercase_      (self     :						str       ,  __UpperCAmelCase     :						Optional[Any]       ,  __UpperCAmelCase     :						List[str]       ,  __UpperCAmelCase     :						List[Any]       ,  __UpperCAmelCase     :						List[str]       ,  __UpperCAmelCase     :						Optional[int]       ,  __UpperCAmelCase     :						str       ,  __UpperCAmelCase     :						Optional[Any]		)						->				List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  MegatronBertForMaskedLM(config=lowerCAmelCase__		)
         model.to(lowerCAmelCase__		)
         model.eval()
         UpperCAmelCase__								=  model(lowerCAmelCase__       ,  attention_mask=lowerCAmelCase__       ,  token_type_ids=lowerCAmelCase__       ,  labels=lowerCAmelCase__		)
         self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.seq_length, self.vocab_size)		)


   def 			lowercase_      (self     :						List[Any]       ,  __UpperCAmelCase     :						List[str]       ,  __UpperCAmelCase     :						Dict       ,  __UpperCAmelCase     :						Optional[int]       ,  __UpperCAmelCase     :						Union[str, Any]       ,  __UpperCAmelCase     :						Optional[Any]       ,  __UpperCAmelCase     :						str       ,  __UpperCAmelCase     :						Tuple		)						->				Optional[int]:
         """simple docstring"""



         UpperCAmelCase__								=  MegatronBertForCausalLM(config=lowerCAmelCase__		)
         model.to(lowerCAmelCase__		)
         model.eval()
         UpperCAmelCase__								=  model(lowerCAmelCase__       ,  attention_mask=lowerCAmelCase__       ,  token_type_ids=lowerCAmelCase__       ,  labels=lowerCAmelCase__		)
         self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.seq_length, self.vocab_size)		)


   def 			lowercase_      (self     :						List[Any]       ,  __UpperCAmelCase     :						Dict       ,  __UpperCAmelCase     :						List[str]       ,  __UpperCAmelCase     :						Dict       ,  __UpperCAmelCase     :						Optional[Any]       ,  __UpperCAmelCase     :						int       ,  __UpperCAmelCase     :						Tuple       ,  __UpperCAmelCase     :						List[str]		)						->				Optional[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  MegatronBertForNextSentencePrediction(config=lowerCAmelCase__		)
         model.to(lowerCAmelCase__		)
         model.eval()
         UpperCAmelCase__								=  model(
             lowerCAmelCase__       ,  attention_mask=lowerCAmelCase__       ,  token_type_ids=lowerCAmelCase__       ,  labels=lowerCAmelCase__       ,  )
         self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, 2)		)


   def 			lowercase_      (self     :						List[str]       ,  __UpperCAmelCase     :						Optional[int]       ,  __UpperCAmelCase     :						List[Any]       ,  __UpperCAmelCase     :						Union[str, Any]       ,  __UpperCAmelCase     :						List[str]       ,  __UpperCAmelCase     :						str       ,  __UpperCAmelCase     :						str       ,  __UpperCAmelCase     :						Tuple		)						->				Optional[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  MegatronBertForPreTraining(config=lowerCAmelCase__		)
         model.to(lowerCAmelCase__		)
         model.eval()
         UpperCAmelCase__								=  model(
             lowerCAmelCase__       ,  attention_mask=lowerCAmelCase__       ,  token_type_ids=lowerCAmelCase__       ,  labels=lowerCAmelCase__       ,  next_sentence_label=lowerCAmelCase__       ,  )
         self.parent.assertEqual(result.prediction_logits.shape       ,  (self.batch_size, self.seq_length, self.vocab_size)		)
         self.parent.assertEqual(result.seq_relationship_logits.shape       ,  (self.batch_size, 2)		)


   def 			lowercase_      (self     :						Any       ,  __UpperCAmelCase     :						Union[str, Any]       ,  __UpperCAmelCase     :						int       ,  __UpperCAmelCase     :						Dict       ,  __UpperCAmelCase     :						Any       ,  __UpperCAmelCase     :						Tuple       ,  __UpperCAmelCase     :						int       ,  __UpperCAmelCase     :						Tuple		)						->				Dict:
         """simple docstring"""



         UpperCAmelCase__								=  MegatronBertForQuestionAnswering(config=lowerCAmelCase__		)
         model.to(lowerCAmelCase__		)
         model.eval()
         UpperCAmelCase__								=  model(
             lowerCAmelCase__       ,  attention_mask=lowerCAmelCase__       ,  token_type_ids=lowerCAmelCase__       ,  start_positions=lowerCAmelCase__       ,  end_positions=lowerCAmelCase__       ,  )
         self.parent.assertEqual(result.start_logits.shape       ,  (self.batch_size, self.seq_length)		)
         self.parent.assertEqual(result.end_logits.shape       ,  (self.batch_size, self.seq_length)		)


   def 			lowercase_      (self     :						Optional[Any]       ,  __UpperCAmelCase     :						Optional[int]       ,  __UpperCAmelCase     :						Union[str, Any]       ,  __UpperCAmelCase     :						Dict       ,  __UpperCAmelCase     :						Tuple       ,  __UpperCAmelCase     :						Optional[int]       ,  __UpperCAmelCase     :						Union[str, Any]       ,  __UpperCAmelCase     :						Tuple		)						->				List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self.num_labels
         UpperCAmelCase__								=  MegatronBertForSequenceClassification(lowerCAmelCase__		)
         model.to(lowerCAmelCase__		)
         model.eval()
         UpperCAmelCase__								=  model(lowerCAmelCase__       ,  attention_mask=lowerCAmelCase__       ,  token_type_ids=lowerCAmelCase__       ,  labels=lowerCAmelCase__		)
         self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.num_labels)		)


   def 			lowercase_      (self     :						List[Any]       ,  __UpperCAmelCase     :						List[Any]       ,  __UpperCAmelCase     :						Any       ,  __UpperCAmelCase     :						Dict       ,  __UpperCAmelCase     :						Tuple       ,  __UpperCAmelCase     :						Any       ,  __UpperCAmelCase     :						Tuple       ,  __UpperCAmelCase     :						Optional[int]		)						->				Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self.num_labels
         UpperCAmelCase__								=  MegatronBertForTokenClassification(config=lowerCAmelCase__		)
         model.to(lowerCAmelCase__		)
         model.eval()
         UpperCAmelCase__								=  model(lowerCAmelCase__       ,  attention_mask=lowerCAmelCase__       ,  token_type_ids=lowerCAmelCase__       ,  labels=lowerCAmelCase__		)
         self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.seq_length, self.num_labels)		)


   def 			lowercase_      (self     :						int       ,  __UpperCAmelCase     :						int       ,  __UpperCAmelCase     :						Optional[int]       ,  __UpperCAmelCase     :						List[Any]       ,  __UpperCAmelCase     :						List[str]       ,  __UpperCAmelCase     :						Tuple       ,  __UpperCAmelCase     :						str       ,  __UpperCAmelCase     :						List[str]		)						->				Optional[int]:
         """simple docstring"""



         UpperCAmelCase__								=  self.num_choices
         UpperCAmelCase__								=  MegatronBertForMultipleChoice(config=lowerCAmelCase__		)
         model.to(lowerCAmelCase__		)
         model.eval()
         UpperCAmelCase__								=  input_ids.unsqueeze(1		).expand(-1       ,  self.num_choices       ,  -1		).contiguous()
         UpperCAmelCase__								=  token_type_ids.unsqueeze(1		).expand(-1       ,  self.num_choices       ,  -1		).contiguous()
         UpperCAmelCase__								=  input_mask.unsqueeze(1		).expand(-1       ,  self.num_choices       ,  -1		).contiguous()
         UpperCAmelCase__								=  model(
             lowerCAmelCase__       ,  attention_mask=lowerCAmelCase__       ,  token_type_ids=lowerCAmelCase__       ,  labels=lowerCAmelCase__       ,  )
         self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.num_choices)		)





   def 			lowercase_      (self     :						List[Any]		)						->				Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self.prepare_config_and_inputs()
         (
             UpperCAmelCase__    
         )								=  config_and_inputs
         UpperCAmelCase__								=  {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
         return config, inputs_dict





@require_torch
class      A							(				UpperCAmelCase_		,				UpperCAmelCase_		,				unittest.TestCase     ):
   __UpperCAmelCase : Any      =			(
       (
           MegatronBertModel,
           MegatronBertForMaskedLM,
           MegatronBertForCausalLM,
           MegatronBertForMultipleChoice,
           MegatronBertForNextSentencePrediction,
           MegatronBertForPreTraining,
           MegatronBertForQuestionAnswering,
           MegatronBertForSequenceClassification,
           MegatronBertForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   __UpperCAmelCase : Optional[int]      =			(
       {
           '''feature-extraction''': MegatronBertModel,
           '''fill-mask''': MegatronBertForMaskedLM,
           '''question-answering''': MegatronBertForQuestionAnswering,
           '''text-classification''': MegatronBertForSequenceClassification,
           '''text-generation''': MegatronBertForCausalLM,
           '''token-classification''': MegatronBertForTokenClassification,
           '''zero-shot''': MegatronBertForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   __UpperCAmelCase : List[Any]      =			True
   # test_resize_embeddings = False
   __UpperCAmelCase : str      =			False


   def 			lowercase_      (self     :						Any       ,  __UpperCAmelCase     :						Union[str, Any]       ,  __UpperCAmelCase     :						Union[str, Any]       ,  __UpperCAmelCase     :						Optional[int]=False		)						->				Dict:
         """simple docstring"""



         UpperCAmelCase__								=  super()._prepare_for_class(lowerCAmelCase__       ,  lowerCAmelCase__       ,  return_labels=lowerCAmelCase__		)

         if return_labels:
               if model_class in get_values(lowerCAmelCase__		):
                     UpperCAmelCase__								=  torch.zeros(
                         (self.model_tester.batch_size, self.model_tester.seq_length)       ,  dtype=torch.long       ,  device=lowerCAmelCase__		)
                     UpperCAmelCase__								=  torch.zeros(
                         self.model_tester.batch_size       ,  dtype=torch.long       ,  device=lowerCAmelCase__		)
         return inputs_dict


   def 			lowercase_      (self     :						Optional[Any]		)						->				List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  MegatronBertModelTester(self		)
         UpperCAmelCase__								=  ConfigTester(self       ,  config_class=lowerCAmelCase__       ,  hidden_size=3_7		)


   def 			lowercase_      (self     :						Dict		)						->				int:
         """simple docstring"""



         self.config_tester.run_common_tests()


   def 			lowercase_      (self     :						int		)						->				Dict:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_model(*lowerCAmelCase__		)


   def 			lowercase_      (self     :						Any		)						->				Dict:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowerCAmelCase__		)


   def 			lowercase_      (self     :						int		)						->				List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowerCAmelCase__		)


   def 			lowercase_      (self     :						List[Any]		)						->				str:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowerCAmelCase__		)


   def 			lowercase_      (self     :						List[str]		)						->				int:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowerCAmelCase__		)


   def 			lowercase_      (self     :						Optional[Any]		)						->				int:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowerCAmelCase__		)


   def 			lowercase_      (self     :						List[Any]		)						->				List[str]:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowerCAmelCase__		)





   def 			lowercase_      (self     :						Any		)						->				Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowerCAmelCase__		)




def    lowerCAmelCase_ (     __A							)     ->				str:







      '''simple docstring'''


      return torch.tensor(
          _UpperCAmelCase,  dtype=torch.long,  device=_UpperCAmelCase,  )


UpperCamelCase__											=					1e-4







@require_torch
@require_sentencepiece
@require_tokenizers
class      A							(				unittest.TestCase     ):


   @slow
   @unittest.skip("Model is not available."		)
   def 			lowercase_      (self     :						Any		)						->				Dict:
         """simple docstring"""



         UpperCAmelCase__								=  "nvidia/megatron-bert-uncased-345m"
         if "MYDIR" in os.environ:
               UpperCAmelCase__								=  os.path.join(os.environ["MYDIR"]       ,  lowerCAmelCase__		)
         UpperCAmelCase__								=  MegatronBertModel.from_pretrained(lowerCAmelCase__		)
         model.to(lowerCAmelCase__		)
         model.half()
         UpperCAmelCase__								=  _long_tensor([[1_0_1, 7_1_1_0, 1_0_0_5, 1_0_5_6, 2_0_2_3, 1_1_3_3_3, 1_7_4_1_3, 1_0_2_9, 1_0_2]]		)
         with torch.no_grad():
               UpperCAmelCase__								=  model(lowerCAmelCase__		)[0]
         UpperCAmelCase__								=  torch.Size((1, 9, 1_0_2_4)		)
         self.assertEqual(output.shape       ,  lowerCAmelCase__		)

         UpperCAmelCase__								=  [-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
         for ii in range(3		):
               for jj in range(3		):
                     UpperCAmelCase__								=  output[0, ii, jj]
                     UpperCAmelCase__								=  expected[3 * ii + jj]
                     UpperCAmelCase__								=  "ii={} jj={} a={} b={}".format(lowerCAmelCase__       ,  lowerCAmelCase__       ,  lowerCAmelCase__       ,  lowerCAmelCase__		)
                     self.assertTrue(math.isclose(lowerCAmelCase__       ,  lowerCAmelCase__       ,  rel_tol=lowerCAmelCase__       ,  abs_tol=lowerCAmelCase__		)       ,  msg=lowerCAmelCase__		)






import os
import tempfile
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
						import torch
						from torch import nn

						from transformers import (
						    Adafactor,
						    AdamW,
						    get_constant_schedule,
						    get_constant_schedule_with_warmup,
						    get_cosine_schedule_with_warmup,
						    get_cosine_with_hard_restarts_schedule_with_warmup,
						    get_inverse_sqrt_schedule,
						    get_linear_schedule_with_warmup,
						    get_polynomial_decay_schedule_with_warmup,
						)







def 			A_      (    _UpperCAmelCase				,	_UpperCAmelCase=10    ):
						SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       []
						for _ in range(_UpperCAmelCase    ):
												lrs.append(scheduler.get_lr()[0]    )
												scheduler.step()
						return lrs







def 			A_      (    _UpperCAmelCase				,	_UpperCAmelCase=10    ):
						SCREAMING_SNAKE_CASE_: List[str]  		=       []
						for step in range(_UpperCAmelCase    ):
												lrs.append(scheduler.get_lr()[0]    )
												scheduler.step()
												if step == num_steps // 2:
																		with tempfile.TemporaryDirectory() as tmpdirname:
																								SCREAMING_SNAKE_CASE_: Optional[int]  		=       os.path.join(_UpperCAmelCase				,	"schedule.bin"    )
																								torch.save(scheduler.state_dict()				,	_UpperCAmelCase    )

																								SCREAMING_SNAKE_CASE_: Optional[Any]  		=       torch.load(_UpperCAmelCase    )
																								scheduler.load_state_dict(_UpperCAmelCase    )
						return lrs



@require_torch
class   __lowercase				( unittest.TestCase	):




			"""simple docstring"""







			def        _SCREAMING_SNAKE_CASE			( self  :     Dict   ,						lowerCAmelCase__  :     int   ,						lowerCAmelCase__  :     Dict   ,						lowerCAmelCase__  :     Tuple):
									self.assertEqual(len(lowerCAmelCase__)   ,						len(lowerCAmelCase__))
									for a, b in zip(lowerCAmelCase__   ,						lowerCAmelCase__):
															self.assertAlmostEqual(lowerCAmelCase__   ,						lowerCAmelCase__   ,						delta=lowerCAmelCase__)



			def        _SCREAMING_SNAKE_CASE			( self  :     Dict):
									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       torch.tensor([0.1, -0.2, -0.1]   ,						requires_grad=lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: str  		=       torch.tensor([0.4, 0.2, -0.5])
									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       nn.MSELoss()
									# No warmup, constant schedule, no gradient clipping
									SCREAMING_SNAKE_CASE_: int  		=       AdamW(params=[w]   ,						lr=2E-1   ,						weight_decay=0.0)
									for _ in range(100):
															SCREAMING_SNAKE_CASE_: Dict  		=       criterion(lowerCAmelCase__   ,						lowerCAmelCase__)
															loss.backward()
															optimizer.step()
															w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
															w.grad.zero_()
									self.assertListAlmostEqual(w.tolist()   ,						[0.4, 0.2, -0.5]   ,						tol=1E-2)




			def        _SCREAMING_SNAKE_CASE			( self  :     Dict):
									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       torch.tensor([0.1, -0.2, -0.1]   ,						requires_grad=lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: int  		=       torch.tensor([0.4, 0.2, -0.5])
									SCREAMING_SNAKE_CASE_: Any  		=       nn.MSELoss()
									# No warmup, constant schedule, no gradient clipping
									SCREAMING_SNAKE_CASE_: int  		=       Adafactor(
									    params=[w]   ,						lr=1E-2   ,						eps=(1E-30, 1E-3)   ,						clip_threshold=1.0   ,						decay_rate=-0.8   ,						betaa=lowerCAmelCase__   ,						weight_decay=0.0   ,						relative_step=lowerCAmelCase__   ,						scale_parameter=lowerCAmelCase__   ,						warmup_init=lowerCAmelCase__   ,						)
									for _ in range(1000):
															SCREAMING_SNAKE_CASE_: List[Any]  		=       criterion(lowerCAmelCase__   ,						lowerCAmelCase__)
															loss.backward()
															optimizer.step()
															w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
															w.grad.zero_()
									self.assertListAlmostEqual(w.tolist()   ,						[0.4, 0.2, -0.5]   ,						tol=1E-2)







@require_torch
class   __lowercase				( unittest.TestCase	):




			"""simple docstring"""







			_UpperCAmelCase   :							Union[str, Any]										=     nn.Linear(50		,       50	) if is_torch_available() else None
			_UpperCAmelCase   :							List[Any]										=     AdamW(m.parameters()		,       lr=10.0	) if is_torch_available() else None
			_UpperCAmelCase   :							Optional[Any]										=     10



			def        _SCREAMING_SNAKE_CASE			( self  :     Optional[int]   ,						lowerCAmelCase__  :     List[Any]   ,						lowerCAmelCase__  :     int   ,						lowerCAmelCase__  :     int   ,						lowerCAmelCase__  :     Optional[Any]=None):
									self.assertEqual(len(lowerCAmelCase__)   ,						len(lowerCAmelCase__))
									for a, b in zip(lowerCAmelCase__   ,						lowerCAmelCase__):
															self.assertAlmostEqual(lowerCAmelCase__   ,						lowerCAmelCase__   ,						delta=lowerCAmelCase__   ,						msg=lowerCAmelCase__)




			def        _SCREAMING_SNAKE_CASE			( self  :     int):
									SCREAMING_SNAKE_CASE_: Dict  		=       {"num_warmup_steps": 2, "num_training_steps": 10}
									# schedulers doct format
									# function: (sched_args_dict, expected_learning_rates)
									SCREAMING_SNAKE_CASE_: Dict  		=       {
									    get_constant_schedule: ({}, [10.0] * self.num_steps),
									    get_constant_schedule_with_warmup: (
									        {"num_warmup_steps": 4},
									        [0.0, 2.5, 5.0, 7.5, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0],
									    ),
									    get_linear_schedule_with_warmup: (
									        {**common_kwargs},
									        [0.0, 5.0, 10.0, 8.75, 7.5, 6.25, 5.0, 3.75, 2.5, 1.25],
									    ),
									    get_cosine_schedule_with_warmup: (
									        {**common_kwargs},
									        [0.0, 5.0, 10.0, 9.61, 8.53, 6.91, 5.0, 3.08, 1.46, 0.38],
									    ),
									    get_cosine_with_hard_restarts_schedule_with_warmup: (
									        {**common_kwargs, "num_cycles": 2},
									        [0.0, 5.0, 10.0, 8.53, 5.0, 1.46, 10.0, 8.53, 5.0, 1.46],
									    ),
									    get_polynomial_decay_schedule_with_warmup: (
									        {**common_kwargs, "power": 2.0, "lr_end": 1E-7},
									        [0.0, 5.0, 10.0, 7.656, 5.625, 3.906, 2.5, 1.406, 0.625, 0.156],
									    ),
									    get_inverse_sqrt_schedule: (
									        {"num_warmup_steps": 2},
									        [0.0, 5.0, 10.0, 8.165, 7.071, 6.325, 5.774, 5.345, 5.0, 4.714],
									    ),
									}

									for scheduler_func, data in scheds.items():
															SCREAMING_SNAKE_CASE_		,    SCREAMING_SNAKE_CASE_: Tuple  		=       data

															SCREAMING_SNAKE_CASE_: List[Any]  		=       scheduler_func(self.optimizer   ,						**lowerCAmelCase__)
															self.assertEqual(len([scheduler.get_lr()[0]])   ,						1)
															SCREAMING_SNAKE_CASE_: int  		=       unwrap_schedule(lowerCAmelCase__   ,						self.num_steps)
															self.assertListAlmostEqual(
															    lowerCAmelCase__   ,						lowerCAmelCase__   ,						tol=1E-2   ,						msg=F"failed for {scheduler_func} in normal scheduler"   ,						)

															SCREAMING_SNAKE_CASE_: List[str]  		=       scheduler_func(self.optimizer   ,						**lowerCAmelCase__)
															if scheduler_func.__name__ != "get_constant_schedule":
																					LambdaScheduleWrapper.wrap_scheduler(lowerCAmelCase__)  # wrap to test picklability of the schedule
															SCREAMING_SNAKE_CASE_: Tuple  		=       unwrap_and_save_reload_schedule(lowerCAmelCase__   ,						self.num_steps)
															self.assertListEqual(lowerCAmelCase__   ,						lowerCAmelCase__   ,						msg=F"failed for {scheduler_func} in save and reload")







class   __lowercase				:




			"""simple docstring"""







			def __init__( self  :     str   ,						lowerCAmelCase__  :     List[str]):
									SCREAMING_SNAKE_CASE_: List[Any]  		=       fn



			def __call__( self  :     Optional[int]   ,						*lowerCAmelCase__  :     List[Any]   ,						**lowerCAmelCase__  :     Tuple):
									return self.fn(*lowerCAmelCase__   ,						**lowerCAmelCase__)




			@classmethod
			def        _SCREAMING_SNAKE_CASE			( self  :     Optional[int]   ,						lowerCAmelCase__  :     str):
									SCREAMING_SNAKE_CASE_: str  		=       list(map(self   ,						scheduler.lr_lambdas))






