




import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import (
    DiffusionPipeline,
    UnCLIPImageVariationPipeline,
    UnCLIPScheduler,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, load_image, require_torch_gpu, skip_mps

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()




class 		UpperCamelCase__	(lowerCAmelCase__    ,				unittest.TestCase     ):

    '''simple docstring'''



    lowerCamelCase_      :	int	     =      UnCLIPImageVariationPipeline
    lowerCamelCase_      :	List[str]	     =      IMAGE_VARIATION_PARAMS - {"""height""", """width""", """guidance_scale"""}
    lowerCamelCase_      :	Dict	     =      IMAGE_VARIATION_BATCH_PARAMS

    lowerCamelCase_      :	Any	     =      [
        """generator""",
        """return_dict""",
        """decoder_num_inference_steps""",
        """super_res_num_inference_steps""",
    ]
    lowerCamelCase_      :	Any	     =      False







    @property
    def     _lowercase     (    self     )     ->      List[Any]:
           return 32







    @property
    def     _lowercase     (    self     )     ->      Union[str, Any]:
           return 32







    @property
    def     _lowercase     (    self     )     ->      Optional[Any]:
           return self.time_input_dim







    @property
    def     _lowercase     (    self     )     ->      Optional[Any]:
           return self.time_input_dim * 4







    @property
    def     _lowercase     (    self     )     ->      str:
           return 100







    @property
    def     _lowercase     (    self     )     ->      Tuple:
           lowerCamelCase   :      Optional[Any]			 =					CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"     )
           return tokenizer







    @property
    def     _lowercase     (    self     )     ->      str:
           torch.manual_seed(0     )
           lowerCamelCase   :      Dict			 =					CLIPTextConfig(
               bos_token_id=0     ,  eos_token_id=2     ,  hidden_size=self.text_embedder_hidden_size     ,  projection_dim=self.text_embedder_hidden_size     ,  intermediate_size=37     ,  layer_norm_eps=1e-05     ,  num_attention_heads=4     ,  num_hidden_layers=5     ,  pad_token_id=1     ,  vocab_size=1000     ,  )
           return CLIPTextModelWithProjection(UpperCamelCase__     )







    @property
    def     _lowercase     (    self     )     ->      Dict:
           torch.manual_seed(0     )
           lowerCamelCase   :      List[str]			 =					CLIPVisionConfig(
               hidden_size=self.text_embedder_hidden_size     ,  projection_dim=self.text_embedder_hidden_size     ,  num_hidden_layers=5     ,  num_attention_heads=4     ,  image_size=32     ,  intermediate_size=37     ,  patch_size=1     ,  )
           return CLIPVisionModelWithProjection(UpperCamelCase__     )







    @property
    def     _lowercase     (    self     )     ->      Tuple:
           torch.manual_seed(0     )

           lowerCamelCase   :      str			 =					{
               "clip_embeddings_dim": self.text_embedder_hidden_size,
               "time_embed_dim": self.time_embed_dim,
               "cross_attention_dim": self.cross_attention_dim,
           }

           lowerCamelCase   :      str			 =					UnCLIPTextProjModel(**UpperCamelCase__     )
           return model







    @property
    def     _lowercase     (    self     )     ->      Tuple:
           torch.manual_seed(0     )

           lowerCamelCase   :      Tuple			 =					{
               "sample_size": 32,
               # RGB in channels
               "in_channels": 3,
               # Out channels is double in channels because predicts mean and variance
               "out_channels": 6,
               "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
               "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
               "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
               "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
               "layers_per_block": 1,
               "cross_attention_dim": self.cross_attention_dim,
               "attention_head_dim": 4,
               "resnet_time_scale_shift": "scale_shift",
               "class_embed_type": "identity",
           }

           lowerCamelCase   :      Tuple			 =					UNetaDConditionModel(**UpperCamelCase__     )
           return model







    @property
    def     _lowercase     (    self     )     ->      Tuple:
           return {
               "sample_size": 64,
               "layers_per_block": 1,
               "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
               "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
               "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
               "in_channels": 6,
               "out_channels": 3,
           }







    @property
    def     _lowercase     (    self     )     ->      List[str]:
           torch.manual_seed(0     )

           lowerCamelCase   :      Tuple			 =					UNetaDModel(**self.dummy_super_res_kwargs     )
           return model







    @property
    def     _lowercase     (    self     )     ->      Tuple:
           # seeded differently to get different unet than `self.dummy_super_res_first`
           torch.manual_seed(1     )

           lowerCamelCase   :      Any			 =					UNetaDModel(**self.dummy_super_res_kwargs     )
           return model







    def     _lowercase     (    self     )     ->      List[Any]:
           lowerCamelCase   :      List[Any]			 =					self.dummy_decoder
           lowerCamelCase   :      Union[str, Any]			 =					self.dummy_text_proj
           lowerCamelCase   :      Union[str, Any]			 =					self.dummy_text_encoder
           lowerCamelCase   :      List[str]			 =					self.dummy_tokenizer
           lowerCamelCase   :      Optional[Any]			 =					self.dummy_super_res_first
           lowerCamelCase   :      Tuple			 =					self.dummy_super_res_last

           lowerCamelCase   :      Union[str, Any]			 =					UnCLIPScheduler(
               variance_type="learned_range"     ,  prediction_type="epsilon"     ,  num_train_timesteps=1000     ,  )

           lowerCamelCase   :      Optional[int]			 =					UnCLIPScheduler(
               variance_type="fixed_small_log"     ,  prediction_type="epsilon"     ,  num_train_timesteps=1000     ,  )

           lowerCamelCase   :      Union[str, Any]			 =					CLIPImageProcessor(crop_size=32     ,  size=32     )

           lowerCamelCase   :      Any			 =					self.dummy_image_encoder

           return {
               "decoder": decoder,
               "text_encoder": text_encoder,
               "tokenizer": tokenizer,
               "text_proj": text_proj,
               "feature_extractor": feature_extractor,
               "image_encoder": image_encoder,
               "super_res_first": super_res_first,
               "super_res_last": super_res_last,
               "decoder_scheduler": decoder_scheduler,
               "super_res_scheduler": super_res_scheduler,
           }







    def     _lowercase     (    self     ,  UpperCamelCase__     ,  UpperCamelCase__=0     ,  UpperCamelCase__=True     )     ->      Optional[Any]:
           lowerCamelCase   :      Union[str, Any]			 =					floats_tensor((1, 3, 32, 32)     ,  rng=random.Random(UpperCamelCase__     )     ).to(UpperCamelCase__     )
           if str(UpperCamelCase__     ).startswith("mps"     ):
                  lowerCamelCase   :      List[Any]			 =					torch.manual_seed(UpperCamelCase__     )
           else:
                  lowerCamelCase   :      Any			 =					torch.Generator(device=UpperCamelCase__     ).manual_seed(UpperCamelCase__     )

           if pil_image:
                  lowerCamelCase   :      Tuple			 =					input_image * 0.5 + 0.5
                  lowerCamelCase   :      Any			 =					input_image.clamp(0     ,  1     )
                  lowerCamelCase   :      Dict			 =					input_image.cpu().permute(0     ,  2     ,  3     ,  1     ).float().numpy()
                  lowerCamelCase   :      Union[str, Any]			 =					DiffusionPipeline.numpy_to_pil(UpperCamelCase__     )[0]

           return {
               "image": input_image,
               "generator": generator,
               "decoder_num_inference_steps": 2,
               "super_res_num_inference_steps": 2,
               "output_type": "np",
           }







    def     _lowercase     (    self     )     ->      Tuple:
           lowerCamelCase   :      List[str]			 =					"cpu"

           lowerCamelCase   :      List[str]			 =					self.get_dummy_components()

           lowerCamelCase   :      Union[str, Any]			 =					self.pipeline_class(**UpperCamelCase__     )
           lowerCamelCase   :      Any			 =					pipe.to(UpperCamelCase__     )

           pipe.set_progress_bar_config(disable=UpperCamelCase__     )

           lowerCamelCase   :      str			 =					self.get_dummy_inputs(UpperCamelCase__     ,  pil_image=UpperCamelCase__     )

           lowerCamelCase   :      Optional[int]			 =					pipe(**UpperCamelCase__     )
           lowerCamelCase   :      Tuple			 =					output.images

           lowerCamelCase   :      str			 =					self.get_dummy_inputs(UpperCamelCase__     ,  pil_image=UpperCamelCase__     )

           lowerCamelCase   :      Optional[int]			 =					pipe(
               **UpperCamelCase__     ,  return_dict=UpperCamelCase__     ,  )[0]

           lowerCamelCase   :      List[Any]			 =					image[0, -3:, -3:, -1]
           lowerCamelCase   :      List[Any]			 =					image_from_tuple[0, -3:, -3:, -1]

           assert image.shape == (1, 64, 64, 3)

           lowerCamelCase   :      Union[str, Any]			 =					np.array(
               [
                   0.9997,
                   0.0002,
                   0.9997,
                   0.9997,
                   0.9969,
                   0.0023,
                   0.9997,
                   0.9969,
                   0.9970,
               ]     )

           assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1e-2
           assert np.abs(image_from_tuple_slice.flatten() - expected_slice     ).max() < 1e-2







    def     _lowercase     (    self     )     ->      int:
           lowerCamelCase   :      int			 =					"cpu"

           lowerCamelCase   :      Tuple			 =					self.get_dummy_components()

           lowerCamelCase   :      Union[str, Any]			 =					self.pipeline_class(**UpperCamelCase__     )
           lowerCamelCase   :      List[Any]			 =					pipe.to(UpperCamelCase__     )

           pipe.set_progress_bar_config(disable=UpperCamelCase__     )

           lowerCamelCase   :      int			 =					self.get_dummy_inputs(UpperCamelCase__     ,  pil_image=UpperCamelCase__     )

           lowerCamelCase   :      Any			 =					pipe(**UpperCamelCase__     )
           lowerCamelCase   :      List[str]			 =					output.images

           lowerCamelCase   :      Dict			 =					self.get_dummy_inputs(UpperCamelCase__     ,  pil_image=UpperCamelCase__     )

           lowerCamelCase   :      str			 =					pipe(
               **UpperCamelCase__     ,  return_dict=UpperCamelCase__     ,  )[0]

           lowerCamelCase   :      int			 =					image[0, -3:, -3:, -1]
           lowerCamelCase   :      int			 =					image_from_tuple[0, -3:, -3:, -1]

           assert image.shape == (1, 64, 64, 3)

           lowerCamelCase   :      Dict			 =					np.array([0.9997, 0.0003, 0.9997, 0.9997, 0.9970, 0.0024, 0.9997, 0.9971, 0.9971]     )

           assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1e-2
           assert np.abs(image_from_tuple_slice.flatten() - expected_slice     ).max() < 1e-2







    def     _lowercase     (    self     )     ->      Optional[int]:
           lowerCamelCase   :      int			 =					"cpu"

           lowerCamelCase   :      Tuple			 =					self.get_dummy_components()

           lowerCamelCase   :      Optional[int]			 =					self.pipeline_class(**UpperCamelCase__     )
           lowerCamelCase   :      Any			 =					pipe.to(UpperCamelCase__     )

           pipe.set_progress_bar_config(disable=UpperCamelCase__     )

           lowerCamelCase   :      List[Any]			 =					self.get_dummy_inputs(UpperCamelCase__     ,  pil_image=UpperCamelCase__     )
           lowerCamelCase   :      List[Any]			 =					[
               pipeline_inputs["image"],
               pipeline_inputs["image"],
           ]

           lowerCamelCase   :      Tuple			 =					pipe(**UpperCamelCase__     )
           lowerCamelCase   :      str			 =					output.images

           lowerCamelCase   :      Tuple			 =					self.get_dummy_inputs(UpperCamelCase__     ,  pil_image=UpperCamelCase__     )
           lowerCamelCase   :      Optional[Any]			 =					[
               tuple_pipeline_inputs["image"],
               tuple_pipeline_inputs["image"],
           ]

           lowerCamelCase   :      Optional[int]			 =					pipe(
               **UpperCamelCase__     ,  return_dict=UpperCamelCase__     ,  )[0]

           lowerCamelCase   :      Union[str, Any]			 =					image[0, -3:, -3:, -1]
           lowerCamelCase   :      Any			 =					image_from_tuple[0, -3:, -3:, -1]

           assert image.shape == (2, 64, 64, 3)

           lowerCamelCase   :      int			 =					np.array(
               [
                   0.9997,
                   0.9989,
                   0.0008,
                   0.0021,
                   0.9960,
                   0.0018,
                   0.0014,
                   0.0002,
                   0.9933,
               ]     )

           assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1e-2
           assert np.abs(image_from_tuple_slice.flatten() - expected_slice     ).max() < 1e-2







    def     _lowercase     (    self     )     ->      List[Any]:
           lowerCamelCase   :      Optional[Any]			 =					torch.device("cpu"     )




           class 		UpperCamelCase__	:

               '''simple docstring'''



               lowerCamelCase_      :	Any	     =      1




           lowerCamelCase   :      Dict			 =					self.get_dummy_components()

           lowerCamelCase   :      List[str]			 =					self.pipeline_class(**UpperCamelCase__     )
           lowerCamelCase   :      Dict			 =					pipe.to(UpperCamelCase__     )

           pipe.set_progress_bar_config(disable=UpperCamelCase__     )

           lowerCamelCase   :      Dict			 =					torch.Generator(device=UpperCamelCase__     ).manual_seed(0     )
           lowerCamelCase   :      Dict			 =					pipe.decoder.dtype
           lowerCamelCase   :      str			 =					1

           lowerCamelCase   :      str			 =					(
               batch_size,
               pipe.decoder.config.in_channels,
               pipe.decoder.config.sample_size,
               pipe.decoder.config.sample_size,
           )
           lowerCamelCase   :      Optional[Any]			 =					pipe.prepare_latents(
               UpperCamelCase__     ,  dtype=UpperCamelCase__     ,  device=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  latents=UpperCamelCase__     ,  scheduler=DummyScheduler()     )

           lowerCamelCase   :      Union[str, Any]			 =					(
               batch_size,
               pipe.super_res_first.config.in_channels // 2,
               pipe.super_res_first.config.sample_size,
               pipe.super_res_first.config.sample_size,
           )
           lowerCamelCase   :      Tuple			 =					pipe.prepare_latents(
               UpperCamelCase__     ,  dtype=UpperCamelCase__     ,  device=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  latents=UpperCamelCase__     ,  scheduler=DummyScheduler()     )

           lowerCamelCase   :      Optional[Any]			 =					self.get_dummy_inputs(UpperCamelCase__     ,  pil_image=UpperCamelCase__     )

           lowerCamelCase   :      Dict			 =					pipe(
               **UpperCamelCase__     ,  decoder_latents=UpperCamelCase__     ,  super_res_latents=UpperCamelCase__     ).images

           lowerCamelCase   :      Any			 =					self.get_dummy_inputs(UpperCamelCase__     ,  pil_image=UpperCamelCase__     )
           # Don't pass image, instead pass embedding
           lowerCamelCase   :      List[Any]			 =					pipeline_inputs.pop("image"     )
           lowerCamelCase   :      int			 =					pipe.image_encoder(UpperCamelCase__     ).image_embeds

           lowerCamelCase   :      Dict			 =					pipe(
               **UpperCamelCase__     ,  decoder_latents=UpperCamelCase__     ,  super_res_latents=UpperCamelCase__     ,  image_embeddings=UpperCamelCase__     ,  ).images

           # make sure passing text embeddings manually is identical
           assert np.abs(img_out_a - img_out_a     ).max() < 1e-4







    @skip_mps
    def     _lowercase     (    self     )     ->      str:
           lowerCamelCase   :      Union[str, Any]			 =					torch_device == "cpu"

           # Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
           lowerCamelCase   :      Optional[Any]			 =					1e-2

           self._test_attention_slicing_forward_pass(
               test_max_difference=UpperCamelCase__     ,  expected_max_diff=UpperCamelCase__     )







    @skip_mps
    def     _lowercase     (    self     )     ->      Union[str, Any]:
           lowerCamelCase   :      Union[str, Any]			 =					torch_device == "cpu"
           lowerCamelCase   :      Tuple			 =					True
           lowerCamelCase   :      Optional[int]			 =					[
               "decoder_num_inference_steps",
               "super_res_num_inference_steps",
           ]

           self._test_inference_batch_single_identical(
               test_max_difference=UpperCamelCase__     ,  relax_max_difference=UpperCamelCase__     ,  additional_params_copy_to_batched_inputs=UpperCamelCase__     ,  )







    def     _lowercase     (    self     )     ->      str:
           lowerCamelCase   :      List[Any]			 =					[
               "decoder_num_inference_steps",
               "super_res_num_inference_steps",
           ]

           if torch_device == "mps":
                  # TODO: MPS errors with larger batch sizes
                  lowerCamelCase   :      Optional[int]			 =					[2, 3]
                  self._test_inference_batch_consistent(
                      batch_sizes=UpperCamelCase__     ,  additional_params_copy_to_batched_inputs=UpperCamelCase__     ,  )
           else:
                  self._test_inference_batch_consistent(
                      additional_params_copy_to_batched_inputs=UpperCamelCase__     )







    @skip_mps
    def     _lowercase     (    self     )     ->      Dict:
           return super().test_dict_tuple_outputs_equivalent()







    @skip_mps
    def     _lowercase     (    self     )     ->      Union[str, Any]:
           return super().test_save_load_local()







    @skip_mps
    def     _lowercase     (    self     )     ->      Union[str, Any]:
           return super().test_save_load_optional_components()




@slow
@require_torch_gpu
class 		UpperCamelCase__	(unittest.TestCase     ):

    '''simple docstring'''



    def     _lowercase     (    self     )     ->      Optional[Any]:
           # clean up the VRAM after each test
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()







    def     _lowercase     (    self     )     ->      List[Any]:
           lowerCamelCase   :      Optional[Any]			 =					load_image(
               "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png"     )
           lowerCamelCase   :      List[str]			 =					load_numpy(
               "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
               "/unclip/karlo_v1_alpha_cat_variation_fp16.npy"     )

           lowerCamelCase   :      Optional[int]			 =					UnCLIPImageVariationPipeline.from_pretrained(
               "kakaobrain/karlo-v1-alpha-image-variations"     ,  torch_dtype=torch.floataa     )
           lowerCamelCase   :      Tuple			 =					pipeline.to(UpperCamelCase__     )
           pipeline.set_progress_bar_config(disable=UpperCamelCase__     )

           lowerCamelCase   :      List[Any]			 =					torch.Generator(device="cpu"     ).manual_seed(0     )
           lowerCamelCase   :      Tuple			 =					pipeline(
               UpperCamelCase__     ,  generator=UpperCamelCase__     ,  output_type="np"     ,  )

           lowerCamelCase   :      Union[str, Any]			 =					output.images[0]

           assert image.shape == (256, 256, 3)

           assert_mean_pixel_difference(UpperCamelCase__     ,  UpperCamelCase__     ,  15     )











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__		)		)


