




import unittest

from parameterized import parameterized

from transformers import LlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
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 LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel, LlamaTokenizer







class 						lowercase							:

     def __init__(  self		,     _a		,     _a=13		,     _a=7		,     _a=True		,     _a=True		,     _a=False		,     _a=True		,     _a=99		,     _a=32		,     _a=5		,     _a=4		,     _a=37		,     _a="gelu"		,     _a=0.1		,     _a=0.1		,     _a=512		,     _a=16		,     _a=2		,     _a=0.02		,     _a=3		,     _a=4		,     _a=None		,     )						->							List[str]:
           _A							:     Optional[Any]			=						parent
           _A							:     int			=						batch_size
           _A							:     List[Any]			=						seq_length
           _A							:     List[Any]			=						is_training
           _A							:     str			=						use_input_mask
           _A							:     List[Any]			=						use_token_type_ids
           _A							:     Any			=						use_labels
           _A							:     Dict			=						vocab_size
           _A							:     List[str]			=						hidden_size
           _A							:     Optional[int]			=						num_hidden_layers
           _A							:     List[str]			=						num_attention_heads
           _A							:     Tuple			=						intermediate_size
           _A							:     Optional[int]			=						hidden_act
           _A							:     Optional[int]			=						hidden_dropout_prob
           _A							:     Any			=						attention_probs_dropout_prob
           _A							:     str			=						max_position_embeddings
           _A							:     Any			=						type_vocab_size
           _A							:     Optional[Any]			=						type_sequence_label_size
           _A							:     Optional[Any]			=						initializer_range
           _A							:     Any			=						num_labels
           _A							:     int			=						num_choices
           _A							:     Any			=						scope

     def 					a__	(  self				)						->							List[Any]:
           _A							:     Any			=						ids_tensor([self.batch_size, self.seq_length]		,     self.vocab_size				)

           _A							:     Union[str, Any]			=						None
           if self.use_input_mask:
                 _A							:     Union[str, Any]			=						random_attention_mask([self.batch_size, self.seq_length]				)

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

           _A							:     Dict			=						None
           _A							:     List[Any]			=						None
           _A							:     Dict			=						None
           if self.use_labels:
                 _A							:     Union[str, Any]			=						ids_tensor([self.batch_size]		,     self.type_sequence_label_size				)
                 _A							:     Dict			=						ids_tensor([self.batch_size, self.seq_length]		,     self.num_labels				)
                 _A							:     List[str]			=						ids_tensor([self.batch_size]		,     self.num_choices				)

           _A							:     Any			=						self.get_config()

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

     def 					a__	(  self				)						->							List[Any]:
           return LlamaConfig(
               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=_a		,     initializer_range=self.initializer_range		,     )

     def 					a__	(  self		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a				)						->							Optional[int]:
           _A							:     Optional[int]			=						LlamaModel(config=_a				)
           model.to(_a				)
           model.eval()
           _A							:     int			=						model(_a		,     attention_mask=_a				)
           _A							:     str			=						model(_a				)
           self.parent.assertEqual(result.last_hidden_state.shape		,     (self.batch_size, self.seq_length, self.hidden_size)				)

     def 					a__	(  self		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a		,     )						->							List[str]:
           _A							:     Union[str, Any]			=						True
           _A							:     Any			=						LlamaModel(_a				)
           model.to(_a				)
           model.eval()
           _A							:     int			=						model(
               _a		,     attention_mask=_a		,     encoder_hidden_states=_a		,     encoder_attention_mask=_a		,     )
           _A							:     Optional[int]			=						model(
               _a		,     attention_mask=_a		,     encoder_hidden_states=_a		,     )
           _A							:     Optional[int]			=						model(_a		,     attention_mask=_a				)
           self.parent.assertEqual(result.last_hidden_state.shape		,     (self.batch_size, self.seq_length, self.hidden_size)				)

     def 					a__	(  self		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a		,     )						->							int:
           _A							:     List[Any]			=						LlamaForCausalLM(config=_a				)
           model.to(_a				)
           model.eval()
           _A							:     Optional[Any]			=						model(_a		,     attention_mask=_a		,     labels=_a				)
           self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.seq_length, self.vocab_size)				)

     def 					a__	(  self		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a		,     _a		,     )						->							Any:
           _A							:     Optional[Any]			=						True
           _A							:     Tuple			=						True
           _A							:     str			=						LlamaForCausalLM(config=_a				)
           model.to(_a				)
           model.eval()

           # first forward pass
           _A							:     Union[str, Any]			=						model(
               _a		,     attention_mask=_a		,     encoder_hidden_states=_a		,     encoder_attention_mask=_a		,     use_cache=_a		,     )
           _A							:     str			=						outputs.past_key_values

           # create hypothetical multiple next token and extent to next_input_ids
           _A							:     Tuple			=						ids_tensor((self.batch_size, 3)		,     config.vocab_size				)
           _A							:     Optional[int]			=						ids_tensor((self.batch_size, 3)		,     vocab_size=2				)

           # append to next input_ids and
           _A							:     Dict			=						torch.cat([input_ids, next_tokens]		,     dim=-1				)
           _A							:     List[Any]			=						torch.cat([input_mask, next_mask]		,     dim=-1				)

           _A							:     Dict			=						model(
               _a		,     attention_mask=_a		,     encoder_hidden_states=_a		,     encoder_attention_mask=_a		,     output_hidden_states=_a		,     )["""hidden_states"""][0]
           _A							:     Optional[int]			=						model(
               _a		,     attention_mask=_a		,     encoder_hidden_states=_a		,     encoder_attention_mask=_a		,     past_key_values=_a		,     output_hidden_states=_a		,     )["""hidden_states"""][0]

           # select random slice
           _A							:     Optional[Any]			=						ids_tensor((1,)		,     output_from_past.shape[-1]				).item()
           _A							:     Union[str, Any]			=						output_from_no_past[:, -3:, random_slice_idx].detach()
           _A							:     Tuple			=						output_from_past[:, :, random_slice_idx].detach()

           self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]				)

           # test that outputs are equal for slice
           self.parent.assertTrue(torch.allclose(_a		,     _a		,     atol=1e-3				)				)



     def 					a__	(  self				)						->							int:
           _A							:     List[Any]			=						self.prepare_config_and_inputs()
           (
               (
               _A
           )			,   (
               _A
           )			,   (
               _A
           )			,   (
               _A
           )			,   (
               _A
           )			,   (
               _A
           )			,   (
               _A
           )			,   
           )							:     Any			=						config_and_inputs
           _A							:     List[Any]			=						{"""input_ids""": input_ids, """attention_mask""": input_mask}
           return config, inputs_dict







@require_torch
class 						lowercase							(   UpperCamelCase__,UpperCamelCase__,UpperCamelCase__,unittest.TestCase		):
     _a							 =						(LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
     _a							 =						(LlamaForCausalLM,) if is_torch_available() else ()
     _a							 =						(
         {
             "feature-extraction": LlamaModel,
             "text-classification": LlamaForSequenceClassification,
             "text-generation": LlamaForCausalLM,
             "zero-shot": LlamaForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     _a							 =						False
     _a							 =						False

     def 					a__	(  self				)						->							List[str]:
           _A							:     Union[str, Any]			=						LlamaModelTester(self				)
           _A							:     Dict			=						ConfigTester(self		,     config_class=_a		,     hidden_size=37				)

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

     def 					a__	(  self				)						->							Union[str, Any]:
           _A							:     Dict			=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*_a				)

     def 					a__	(  self				)						->							List[Any]:
           _A							:     Union[str, Any]			=						self.model_tester.prepare_config_and_inputs()
           for type in ["absolute", "relative_key", "relative_key_query"]:
                 _A							:     Optional[Any]			=						type
                 self.model_tester.create_and_check_model(*_a				)

     def 					a__	(  self				)						->							List[Any]:
           _A			,   _A							:     Any			=						self.model_tester.prepare_config_and_inputs_for_common()
           _A							:     List[str]			=						3
           _A							:     Optional[Any]			=						input_dict["""input_ids"""]
           _A							:     Tuple			=						input_ids.ne(1				).to(_a				)
           _A							:     Any			=						ids_tensor([self.model_tester.batch_size]		,     self.model_tester.type_sequence_label_size				)
           _A							:     Tuple			=						LlamaForSequenceClassification(_a				)
           model.to(_a				)
           model.eval()
           _A							:     int			=						model(_a		,     attention_mask=_a		,     labels=_a				)
           self.assertEqual(result.logits.shape		,     (self.model_tester.batch_size, self.model_tester.num_labels)				)

     def 					a__	(  self				)						->							Tuple:
           _A			,   _A							:     Union[str, Any]			=						self.model_tester.prepare_config_and_inputs_for_common()
           _A							:     int			=						3
           _A							:     Union[str, Any]			=						"""single_label_classification"""
           _A							:     Union[str, Any]			=						input_dict["""input_ids"""]
           _A							:     int			=						input_ids.ne(1				).to(_a				)
           _A							:     str			=						ids_tensor([self.model_tester.batch_size]		,     self.model_tester.type_sequence_label_size				)
           _A							:     Optional[Any]			=						LlamaForSequenceClassification(_a				)
           model.to(_a				)
           model.eval()
           _A							:     Union[str, Any]			=						model(_a		,     attention_mask=_a		,     labels=_a				)
           self.assertEqual(result.logits.shape		,     (self.model_tester.batch_size, self.model_tester.num_labels)				)

     def 					a__	(  self				)						->							str:
           _A			,   _A							:     Optional[int]			=						self.model_tester.prepare_config_and_inputs_for_common()
           _A							:     List[str]			=						3
           _A							:     Tuple			=						"""multi_label_classification"""
           _A							:     List[str]			=						input_dict["""input_ids"""]
           _A							:     Union[str, Any]			=						input_ids.ne(1				).to(_a				)
           _A							:     List[str]			=						ids_tensor(
               [self.model_tester.batch_size, config.num_labels]		,     self.model_tester.type_sequence_label_size				).to(torch.float				)
           _A							:     Optional[int]			=						LlamaForSequenceClassification(_a				)
           model.to(_a				)
           model.eval()
           _A							:     Tuple			=						model(_a		,     attention_mask=_a		,     labels=_a				)
           self.assertEqual(result.logits.shape		,     (self.model_tester.batch_size, self.model_tester.num_labels)				)

     @unittest.skip("""LLaMA buffers include complex numbers, which breaks this test"""				)
     def 					a__	(  self				)						->							Dict:
           pass



     @parameterized.expand([("""linear""",), ("""dynamic""",)]				)
     def 					a__	(  self		,     _a				)						->							Any:
           _A			,   _A							:     Dict			=						self.model_tester.prepare_config_and_inputs_for_common()
           _A							:     Dict			=						ids_tensor([1, 10]		,     config.vocab_size				)
           _A							:     Tuple			=						ids_tensor([1, int(config.max_position_embeddings * 1.5				)]		,     config.vocab_size				)

           set_seed(42				)  # Fixed seed at init time so the two models get the same random weights
           _A							:     str			=						LlamaModel(_a				)
           original_model.to(_a				)
           original_model.eval()
           _A							:     Dict			=						original_model(_a				).last_hidden_state
           _A							:     Union[str, Any]			=						original_model(_a				).last_hidden_state

           set_seed(42				)  # Fixed seed at init time so the two models get the same random weights
           _A							:     Optional[int]			=						{"""type""": scaling_type, """factor""": 10.0}
           _A							:     Optional[Any]			=						LlamaModel(_a				)
           scaled_model.to(_a				)
           scaled_model.eval()
           _A							:     Tuple			=						scaled_model(_a				).last_hidden_state
           _A							:     Any			=						scaled_model(_a				).last_hidden_state

           # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
           # maximum sequence length, so the outputs for the short input should match.
           if scaling_type == "dynamic":
                 self.assertTrue(torch.allclose(_a		,     _a		,     atol=1e-5				)				)
           else:
                 self.assertFalse(torch.allclose(_a		,     _a		,     atol=1e-5				)				)

           # The output should be different for long inputs
           self.assertFalse(torch.allclose(_a		,     _a		,     atol=1e-5				)				)







@require_torch
class 						lowercase							(   unittest.TestCase		):

     @unittest.skip("""Logits are not exactly the same, once we fix the instabalities somehow, will update!"""				)
     @slow
     def 					a__	(  self				)						->							Dict:
           _A							:     str			=						[1, 306, 4658, 278, 6593, 310, 2834, 338]
           _A							:     Tuple			=						LlamaForCausalLM.from_pretrained("""meta-llama/Llama-2-7b-hf"""		,     device_map="""auto"""				)
           _A							:     int			=						model(torch.tensor([input_ids]				)				)
           # Expected mean on dim = -1
           _A							:     Optional[int]			=						torch.tensor([[-6.6550, -4.1227, -4.9859, -3.2406, 0.8262, -3.0033, 1.2964, -3.3699]]				)
           torch.testing.assert_close(out.mean(-1				)		,     _a		,     atol=1e-2		,     rtol=1e-2				)
           # slicing logits[0, 0, 0:30]
           # fmt: off
           _A							:     Any			=						torch.tensor([-12.8281, -7.4453, -0.4639, -8.0625, -7.2500, -8.0000, -6.4883, -7.7695, -7.8438, -7.0312, -6.2188, -7.1328, -1.8496, 1.9961, -8.6250, -6.7227, -12.8281, -6.9492, -7.0742, -7.7852, -7.5820, -7.9062, -6.9375, -7.9805, -8.3438, -8.1562, -8.0469, -7.6250, -7.7422, -7.3398,]				)
           # fmt: on
           torch.testing.assert_close(out[0, 0, :30]		,     _a		,     atol=1e-5		,     rtol=1e-5				)

     @unittest.skip("""Logits are not exactly the same, once we fix the instabalities somehow, will update!"""				)
     @slow
     def 					a__	(  self				)						->							Any:
           _A							:     int			=						[1, 306, 4658, 278, 6593, 310, 2834, 338]
           _A							:     str			=						LlamaForCausalLM.from_pretrained("""meta-llama/Llama-2-13b-hf"""		,     device_map="""auto"""				)
           _A							:     Dict			=						model(torch.tensor(_a				)				)
           # Expected mean on dim = -1
           _A							:     Dict			=						torch.tensor([[-2.0622, -1.2794, -1.1638, -0.9788, -1.4603, -1.0238, -1.7893, -1.4411]]				)
           torch.testing.assert_close(out.mean(-1				)		,     _a		,     atol=1e-2		,     rtol=1e-2				)
           # slicing logits[0, 0, 0:30]
           # fmt: off
           _A							:     str			=						torch.tensor([-8.1406, -8.0547, 2.7461, -1.2344, -0.1448, -1.8262, -1.0020, -1.8154, -1.6895, -1.8516, -2.3574, -0.9277, 3.7598, 6.5742, -1.2998, -0.1177, -8.1406, -2.9688, -2.9199, -3.1699, -3.5254, -2.3555, -2.7988, -3.4141, -2.8262, -4.5195, -3.3379, -3.3164, -2.7832, -3.0273]				)
           # fmt: on
           torch.testing.assert_close(out[0, 0, :30]		,     _a		,     atol=1e-5		,     rtol=1e-5				)

     @unittest.skip("""Logits are not exactly the same, once we fix the instabalities somehow, will update!"""				)
     @slow
     def 					a__	(  self				)						->							Optional[Any]:
           _A							:     Tuple			=						[1, 306, 4658, 278, 6593, 310, 2834, 338]
           _A							:     Tuple			=						LlamaForCausalLM.from_pretrained("""meta-llama/Llama-2-13b-chat-hf"""		,     device_map="""auto"""				)
           _A							:     Optional[int]			=						model(torch.tensor(_a				)				)
           # Expected mean on dim = -1
           _A							:     Tuple			=						torch.tensor([[-0.8562, -1.8520, -0.7551, -0.4162, -1.5161, -1.2038, -2.4823, -2.3254]]				)
           torch.testing.assert_close(out.mean(-1				)		,     _a		,     atol=1e-2		,     rtol=1e-2				)
           # slicing logits[0, 0, 0:30]
           # fmt: off
           _A							:     List[Any]			=						torch.tensor([-2.2227, 4.8828, 0.9023, -0.4578, -0.7871, -0.1033, -0.6221, -0.5786, -0.7803, -1.0674, -1.2920, -0.1570, 0.8008, 2.0723, -0.9497, 0.2771, -2.2227, -0.7612, -1.4346, -1.2061, -1.6426, -0.3000, -0.7139, -1.1934, -1.8691, -1.6973, -1.5947, -1.2705, -0.3523, -0.5513]				)
           # fmt: on
           torch.testing.assert_close(out.mean(-1				)		,     _a		,     atol=1e-2		,     rtol=1e-2				)

     @unittest.skip(
         """Logits are not exactly the same, once we fix the instabalities somehow, will update! Also it is gonna be a `too_slow` test"""				)
     @slow
     def 					a__	(  self				)						->							Optional[Any]:
           _A							:     Optional[int]			=						[1, 306, 4658, 278, 6593, 310, 2834, 338]
           _A							:     List[str]			=						LlamaForCausalLM.from_pretrained("""meta-llama/Llama-2-70b-hf"""		,     device_map="""auto"""				)
           _A							:     int			=						model(torch.tensor(_a				)				)

           _A							:     Tuple			=						torch.tensor(
               [[-4.2327, -3.3360, -4.6665, -4.7631, -1.8180, -3.4170, -1.4211, -3.1810]]		,     dtype=torch.floataa				)
           torch.testing.assert_close(out.mean(-1				)		,     _a		,     atol=1e-2		,     rtol=1e-2				)
           # fmt: off
           _A							:     Any			=						torch.tensor([-9.4922, -3.9551, 1.7998, -5.6758, -5.1055, -5.8984, -4.8320, -6.8086, -6.5391, -5.6172, -5.5820, -5.5352, 1.7881, 3.6289, -6.5117, -3.4785, -9.5000, -6.0352, -6.8125, -6.0195, -6.6836, -5.4727, -6.2812, -6.0391, -7.3398, -7.4297, -7.4844, -6.5820, -5.8789, -5.5312]				)
           # fmt: on
           torch.testing.assert_close(out[0, 0, :30]		,     _a		,     atol=1e-5		,     rtol=1e-5				)



     @unittest.skip("""Model is curently gated"""				)
     @slow
     def 					a__	(  self				)						->							Any:
           _A							:     Any			=						"""Simply put, the theory of relativity states that 1) the laws of physics are the same everywhere in the universe and 2) the passage of time and the length of objects can vary depending on the observer\'s frame of reference.\n\nThe first part of the theory, that the laws of physics are the same everywhere, is known as the \"princi"""
           _A							:     List[Any]			=						"""Simply put, the theory of relativity states that """
           _A							:     Dict			=						LlamaTokenizer.from_pretrained("""meta-llama/Llama-2-13b-chat-hf"""				)
           _A							:     List[str]			=						tokenizer.encode(_a		,     return_tensors="""pt"""				)
           _A							:     Optional[Any]			=						LlamaForCausalLM.from_pretrained(
               """meta-llama/Llama-2-13b-chat-hf"""		,     device_map="""sequential"""		,     use_safetensors=_a				)

           # greedy generation outputs
           _A							:     Optional[Any]			=						model.generate(_a		,     max_new_tokens=64		,     top_p=_a		,     temperature=1		,     do_sample=_a				)
           _A							:     Union[str, Any]			=						tokenizer.decode(generated_ids[0]		,     skip_special_tokens=_a				)
           self.assertEqual(_a		,     _a				)






import re
from typing import Callable, List, Optional, Union

import tensorflow as tf


try:
     from tensorflow.keras.optimizers.legacy import Adam
except ImportError:
     from tensorflow.keras.optimizers import Adam






class SCREAMING_SNAKE_CASE__     (		tf.keras.optimizers.schedules.LearningRateSchedule    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__ = 1.0,		lowerCamelCase__ = None,		):
        super().__init__()
        A :       Union[str, Any]          =   initial_learning_rate
        A :       List[Any]          =   warmup_steps
        A :       int          =   power
        A :       Optional[int]          =   decay_schedule_fn
        A :       int          =   name






    def __call__(       self,		lowerCamelCase__		):
        with tf.name_scope(self.name or """WarmUp"""		) as name:
            # Implements polynomial warmup. i.e., if global_step < warmup_steps, the
            # learning rate will be `global_step/num_warmup_steps * init_lr`.
            A :       str          =   tf.cast(lowerCamelCase__,		tf.floataa		)
            A :       List[Any]          =   tf.cast(self.warmup_steps,		tf.floataa		)
            A :       Dict          =   global_step_float / warmup_steps_float
            A :       Union[str, Any]          =   self.initial_learning_rate * tf.math.pow(lowerCamelCase__,		self.power		)
            return tf.cond(
                global_step_float < warmup_steps_float,		lambda: warmup_learning_rate,		lambda: self.decay_schedule_fn(step - self.warmup_steps		),		name=lowerCamelCase__,		)






    def 			_lowerCAmelCase					(       self		):
        return {
            "initial_learning_rate": self.initial_learning_rate,
            "decay_schedule_fn": self.decay_schedule_fn,
            "warmup_steps": self.warmup_steps,
            "power": self.power,
            "name": self.name,
        }




def 	__UpperCamelCase	(  _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase = 0.0      ,     _lowerCAmelCase = 0.9      ,     _lowerCAmelCase = 0.999      ,     _lowerCAmelCase = 1e-8      ,     _lowerCAmelCase = None      ,     _lowerCAmelCase = None      ,     _lowerCAmelCase = 0.0      ,     _lowerCAmelCase = 1.0      ,     _lowerCAmelCase = None      ,     )       ->  Union[str, Any]:



    """simple docstring"""






    A :       Optional[int]          =   tf.keras.optimizers.schedules.PolynomialDecay(
        initial_learning_rate=_lowerCAmelCase      ,     decay_steps=num_train_steps - num_warmup_steps      ,     end_learning_rate=init_lr * min_lr_ratio      ,     power=_lowerCAmelCase      ,     )
    if num_warmup_steps:
        A :       Dict          =   WarmUp(
            initial_learning_rate=_lowerCAmelCase      ,     decay_schedule_fn=_lowerCAmelCase      ,     warmup_steps=_lowerCAmelCase      ,     )
    if weight_decay_rate > 0.0:
        A :       str          =   AdamWeightDecay(
            learning_rate=_lowerCAmelCase      ,     weight_decay_rate=_lowerCAmelCase      ,     beta_a=_lowerCAmelCase      ,     beta_a=_lowerCAmelCase      ,     epsilon=_lowerCAmelCase      ,     clipnorm=_lowerCAmelCase      ,     global_clipnorm=_lowerCAmelCase      ,     exclude_from_weight_decay=["""LayerNorm""", """layer_norm""", """bias"""]      ,     include_in_weight_decay=_lowerCAmelCase      ,     )
    else:
        A :       Optional[int]          =   tf.keras.optimizers.Adam(
            learning_rate=_lowerCAmelCase      ,     beta_a=_lowerCAmelCase      ,     beta_a=_lowerCAmelCase      ,     epsilon=_lowerCAmelCase      ,     clipnorm=_lowerCAmelCase      ,     global_clipnorm=_lowerCAmelCase      ,     )
    # We return the optimizer and the LR scheduler in order to better track the
    # evolution of the LR independently of the optimizer.
    return optimizer, lr_schedule






class SCREAMING_SNAKE_CASE__     (		SCREAMING_SNAKE_CASE__    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__ = 0.001,		lowerCamelCase__ = 0.9,		lowerCamelCase__ = 0.999,		lowerCamelCase__ = 1e-7,		lowerCamelCase__ = False,		lowerCamelCase__ = 0.0,		lowerCamelCase__ = None,		lowerCamelCase__ = None,		lowerCamelCase__ = "AdamWeightDecay",		**lowerCamelCase__,		):
        super().__init__(lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		**lowerCamelCase__		)
        A :       int          =   weight_decay_rate
        A :       Any          =   include_in_weight_decay
        A :       Dict          =   exclude_from_weight_decay






    @classmethod
    def 			_lowerCAmelCase					(       cls,		lowerCamelCase__		):
        A :       Tuple          =   {"""WarmUp""": WarmUp}
        return super(lowerCamelCase__,		cls		).from_config(lowerCamelCase__,		custom_objects=lowerCamelCase__		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		):
        super(lowerCamelCase__,		self		)._prepare_local(lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		)
        A :       List[str]          =   tf.constant(
            self.weight_decay_rate,		name="""adam_weight_decay_rate"""		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		):
        A :       Any          =   self._do_use_weight_decay(var.name		)
        if do_decay:
            return var.assign_sub(
                learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)]["""weight_decay_rate"""],		use_locking=self._use_locking,		)
        return tf.no_op()






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__=None,		**lowerCamelCase__		):
        A     ,							A :       Dict          =   list(zip(*lowerCamelCase__		)		)
        return super(lowerCamelCase__,		self		).apply_gradients(zip(lowerCamelCase__,		lowerCamelCase__		),		name=lowerCamelCase__,		**lowerCamelCase__		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		):
        if apply_state is None:
            return self._decayed_lr_t[var_dtype], {}

        A :       Union[str, Any]          =   apply_state or {}
        A :       Optional[int]          =   apply_state.get((var_device, var_dtype)		)
        if coefficients is None:
            A :       Dict          =   self._fallback_apply_state(lowerCamelCase__,		lowerCamelCase__		)
            A :       List[Any]          =   coefficients

        return coefficients["lr_t"], {"apply_state": apply_state}






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__=None		):
        A     ,							A :       str          =   self._get_lr(var.device,		var.dtype.base_dtype,		lowerCamelCase__		)
        A :       Any          =   self._decay_weights_op(lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		)
        with tf.control_dependencies([decay]		):
            return super(lowerCamelCase__,		self		)._resource_apply_dense(lowerCamelCase__,		lowerCamelCase__,		**lowerCamelCase__		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__=None		):
        A     ,							A :       Tuple          =   self._get_lr(var.device,		var.dtype.base_dtype,		lowerCamelCase__		)
        A :       Optional[Any]          =   self._decay_weights_op(lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		)
        with tf.control_dependencies([decay]		):
            return super(lowerCamelCase__,		self		)._resource_apply_sparse(lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		**lowerCamelCase__		)






    def 			_lowerCAmelCase					(       self		):
        A :       Dict          =   super().get_config()
        config.update({"""weight_decay_rate""": self.weight_decay_rate}		)
        return config






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        if self.weight_decay_rate == 0:
            return False

        if self._include_in_weight_decay:
            for r in self._include_in_weight_decay:
                if re.search(lowerCamelCase__,		lowerCamelCase__		) is not None:
                    return True

        if self._exclude_from_weight_decay:
            for r in self._exclude_from_weight_decay:
                if re.search(lowerCamelCase__,		lowerCamelCase__		) is not None:
                    return False
        return True






class SCREAMING_SNAKE_CASE__     (		SCREAMING_SNAKE_CASE__    ):




    '''simple docstring'''






    def __init__(       self		):
        A :       List[str]          =   []
        A :       List[str]          =   None






    @property
    def 			_lowerCAmelCase					(       self		):
        if self._accum_steps is None:
            A :       str          =   tf.Variable(
                tf.constant(0,		dtype=tf.intaa		),		trainable=lowerCamelCase__,		synchronization=tf.VariableSynchronization.ON_READ,		aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,		)

        return self._accum_steps.value()






    @property
    def 			_lowerCAmelCase					(       self		):
        if not self._gradients:
            raise ValueError("""The accumulator should be called first to initialize the gradients"""		)
        return [gradient.value() if gradient is not None else gradient for gradient in self._gradients]






    def __call__(       self,		lowerCamelCase__		):
        if not self._gradients:
            A :       int          =   self.step  # Create the step variable.
            self._gradients.extend(
                [
                    tf.Variable(
                        tf.zeros_like(lowerCamelCase__		),		trainable=lowerCamelCase__,		synchronization=tf.VariableSynchronization.ON_READ,		aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,		)
                    if gradient is not None
                    else gradient
                    for gradient in gradients
                ]		)
        if len(lowerCamelCase__		) != len(self._gradients		):
            raise ValueError(f'''Expected {len(self._gradients		)} gradients, but got {len(lowerCamelCase__		)}'''		)

        for accum_gradient, gradient in zip(self._gradients,		lowerCamelCase__		):
            if accum_gradient is not None and gradient is not None:
                accum_gradient.assign_add(lowerCamelCase__		)

        self._accum_steps.assign_add(1		)






    def 			_lowerCAmelCase					(       self		):
        if not self._gradients:
            return
        self._accum_steps.assign(0		)
        for gradient in self._gradients:
            if gradient is not None:
                gradient.assign(tf.zeros_like(lowerCamelCase__		)		)


