


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     UpperCamelCase    (      _a       ,		_a=1_0		)					->							int:
     '''simple docstring'''





     lowercase_       :Dict     =       []
     for _ in range(_a		):
          lrs.append(scheduler.get_lr()[0]		)
          scheduler.step()
     return lrs




def     UpperCamelCase    (      _a       ,		_a=1_0		)					->							List[Any]:
     '''simple docstring'''





     lowercase_       :Optional[Any]     =       []
     for step in range(_a		):
          lrs.append(scheduler.get_lr()[0]		)
          scheduler.step()
          if step == num_steps // 2:
               with tempfile.TemporaryDirectory() as tmpdirname:
                    lowercase_       :Any     =       os.path.join(_a       ,		'''schedule.bin'''		)
                    torch.save(scheduler.state_dict()       ,		_a		)

                    lowercase_       :Optional[Any]     =       torch.load(_a		)
                    scheduler.load_state_dict(_a		)
     return lrs






@require_torch
class      UpperCamelCase					(		unittest.TestCase     ):


   '''simple docstring'''







   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		):
        self.assertEqual(len(UpperCamelCase_		)   ,							len(UpperCamelCase_		)		)
        for a, b in zip(UpperCamelCase_   ,							UpperCamelCase_		):
             self.assertAlmostEqual(UpperCamelCase_   ,							UpperCamelCase_   ,							delta=UpperCamelCase_		)
   def 		UpperCamelCase  (     self		):
        lowercase_       :List[Any]     =       torch.tensor([0.1, -0.2, -0.1]   ,							requires_grad=UpperCamelCase_		)
        lowercase_       :Any     =       torch.tensor([0.4, 0.2, -0.5]		)
        lowercase_       :Tuple     =       nn.MSELoss()
        # No warmup, constant schedule, no gradient clipping
        lowercase_       :str     =       AdamW(params=[w]   ,							lr=2E-1   ,							weight_decay=0.0		)
        for _ in range(100		):
             lowercase_       :List[str]     =       criterion(UpperCamelCase_   ,							UpperCamelCase_		)
             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 		UpperCamelCase  (     self		):
        lowercase_       :Optional[int]     =       torch.tensor([0.1, -0.2, -0.1]   ,							requires_grad=UpperCamelCase_		)
        lowercase_       :str     =       torch.tensor([0.4, 0.2, -0.5]		)
        lowercase_       :List[Any]     =       nn.MSELoss()
        # No warmup, constant schedule, no gradient clipping
        lowercase_       :Optional[Any]     =       Adafactor(
            params=[w]   ,							lr=1E-2   ,							eps=(1E-30, 1E-3)   ,							clip_threshold=1.0   ,							decay_rate=-0.8   ,							betaa=UpperCamelCase_   ,							weight_decay=0.0   ,							relative_step=UpperCamelCase_   ,							scale_parameter=UpperCamelCase_   ,							warmup_init=UpperCamelCase_   ,							)
        for _ in range(1000		):
             lowercase_       :List[str]     =       criterion(UpperCamelCase_   ,							UpperCamelCase_		)
             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      UpperCamelCase					(		unittest.TestCase     ):


   '''simple docstring'''







   lowercase					:	List[Any]									=nn.Linear(50						,      50     ) if is_torch_available() else None
   lowercase					:	List[Any]									=AdamW(m.parameters()						,      lr=10.0     ) if is_torch_available() else None
   lowercase					:	List[str]									=10
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_=None		):
        self.assertEqual(len(UpperCamelCase_		)   ,							len(UpperCamelCase_		)		)
        for a, b in zip(UpperCamelCase_   ,							UpperCamelCase_		):
             self.assertAlmostEqual(UpperCamelCase_   ,							UpperCamelCase_   ,							delta=UpperCamelCase_   ,							msg=UpperCamelCase_		)




   def 		UpperCamelCase  (     self		):
        lowercase_       :Tuple     =       {'''num_warmup_steps''': 2, '''num_training_steps''': 10}
        # schedulers doct format
        # function: (sched_args_dict, expected_learning_rates)
        lowercase_       :Any     =       {
            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():
             lowercase_       :int     =       data

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

             lowercase_       :Optional[int]     =       scheduler_func(self.optimizer   ,							**UpperCamelCase_		)
             if scheduler_func.__name__ != "get_constant_schedule":
                  LambdaScheduleWrapper.wrap_scheduler(UpperCamelCase_		)  # wrap to test picklability of the schedule
             lowercase_       :int     =       unwrap_and_save_reload_schedule(UpperCamelCase_   ,							self.num_steps		)
             self.assertListEqual(UpperCamelCase_   ,							UpperCamelCase_   ,							msg=f"failed for {scheduler_func} in save and reload"		)






class      UpperCamelCase					:


   '''simple docstring'''







   def __init__(     self   ,							UpperCamelCase_		):
        lowercase_       :Dict     =       fn
   def __call__(     self   ,							*UpperCamelCase_   ,							**UpperCamelCase_		):
        return self.fn(*UpperCamelCase_   ,							**UpperCamelCase_		)




   @classmethod
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_		):
        lowercase_       :Union[str, Any]     =       list(map(self   ,							scheduler.lr_lambdas		)		)





import unittest

import numpy as np

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

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
   from transformers.models.bert.modeling_flax_bert import (
       FlaxBertForMaskedLM,
       FlaxBertForMultipleChoice,
       FlaxBertForNextSentencePrediction,
       FlaxBertForPreTraining,
       FlaxBertForQuestionAnswering,
       FlaxBertForSequenceClassification,
       FlaxBertForTokenClassification,
       FlaxBertModel,
   )






class      UpperCamelCase					(		unittest.TestCase     ):


   '''simple docstring'''







   def __init__(     self   ,							UpperCamelCase_   ,							UpperCamelCase_=13   ,							UpperCamelCase_=7   ,							UpperCamelCase_=True   ,							UpperCamelCase_=True   ,							UpperCamelCase_=True   ,							UpperCamelCase_=True   ,							UpperCamelCase_=99   ,							UpperCamelCase_=32   ,							UpperCamelCase_=5   ,							UpperCamelCase_=4   ,							UpperCamelCase_=37   ,							UpperCamelCase_="gelu"   ,							UpperCamelCase_=0.1   ,							UpperCamelCase_=0.1   ,							UpperCamelCase_=512   ,							UpperCamelCase_=16   ,							UpperCamelCase_=2   ,							UpperCamelCase_=0.02   ,							UpperCamelCase_=4   ,							):
        lowercase_       :Union[str, Any]     =       parent
        lowercase_       :int     =       batch_size
        lowercase_       :int     =       seq_length
        lowercase_       :str     =       is_training
        lowercase_       :Dict     =       use_attention_mask
        lowercase_       :List[Any]     =       use_token_type_ids
        lowercase_       :str     =       use_labels
        lowercase_       :str     =       vocab_size
        lowercase_       :Optional[int]     =       hidden_size
        lowercase_       :Dict     =       num_hidden_layers
        lowercase_       :List[str]     =       num_attention_heads
        lowercase_       :int     =       intermediate_size
        lowercase_       :Union[str, Any]     =       hidden_act
        lowercase_       :Optional[int]     =       hidden_dropout_prob
        lowercase_       :Tuple     =       attention_probs_dropout_prob
        lowercase_       :int     =       max_position_embeddings
        lowercase_       :List[Any]     =       type_vocab_size
        lowercase_       :Any     =       type_sequence_label_size
        lowercase_       :Tuple     =       initializer_range
        lowercase_       :Any     =       num_choices
   def 		UpperCamelCase  (     self		):
        lowercase_       :Optional[int]     =       ids_tensor([self.batch_size, self.seq_length]   ,							self.vocab_size		)

        lowercase_       :List[str]     =       None
        if self.use_attention_mask:
             lowercase_       :Any     =       random_attention_mask([self.batch_size, self.seq_length]		)

        lowercase_       :List[Any]     =       None
        if self.use_token_type_ids:
             lowercase_       :Union[str, Any]     =       ids_tensor([self.batch_size, self.seq_length]   ,							self.type_vocab_size		)

        lowercase_       :Union[str, Any]     =       BertConfig(
            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   ,							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=UpperCamelCase_   ,							initializer_range=self.initializer_range   ,							)

        return config, input_ids, token_type_ids, attention_mask
   def 		UpperCamelCase  (     self		):
        lowercase_       :Optional[Any]     =       self.prepare_config_and_inputs()
        lowercase_    ,      lowercase_    ,      lowercase_    ,      lowercase_       :List[str]     =       config_and_inputs
        lowercase_       :Union[str, Any]     =       {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
        return config, inputs_dict




   def 		UpperCamelCase  (     self		):
        lowercase_       :Dict     =       self.prepare_config_and_inputs()
        lowercase_    ,      lowercase_    ,      lowercase_    ,      lowercase_       :Dict     =       config_and_inputs

        lowercase_       :str     =       True
        lowercase_       :Dict     =       floats_tensor([self.batch_size, self.seq_length, self.hidden_size]		)
        lowercase_       :Optional[Any]     =       ids_tensor([self.batch_size, self.seq_length]   ,							vocab_size=2		)

        return (
            config,
            input_ids,
            attention_mask,
            encoder_hidden_states,
            encoder_attention_mask,
        )






@require_flax
class      UpperCamelCase					(		lowercase__						,      unittest.TestCase     ):


   '''simple docstring'''







   lowercase					:	int									=True

   lowercase					:	Dict									=(
       (
           FlaxBertModel,
           FlaxBertForPreTraining,
           FlaxBertForMaskedLM,
           FlaxBertForMultipleChoice,
           FlaxBertForQuestionAnswering,
           FlaxBertForNextSentencePrediction,
           FlaxBertForSequenceClassification,
           FlaxBertForTokenClassification,
           FlaxBertForQuestionAnswering,
       )
       if is_flax_available()
       else ()
   )
   def 		UpperCamelCase  (     self		):
        lowercase_       :Union[str, Any]     =       FlaxBertModelTester(self		)




   @slow
   def 		UpperCamelCase  (     self		):
        # Only check this for base model, not necessary for all model classes.
        # This will also help speed-up tests.
        lowercase_       :Dict     =       FlaxBertModel.from_pretrained('''bert-base-cased'''		)
        lowercase_       :str     =       model(np.ones((1, 1)		)		)
        self.assertIsNotNone(UpperCamelCase_		)

