'''simple docstring'''

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  (						a__   ,						unittest.TestCase ):





       '''simple docstring'''



       SCREAMING_SNAKE_CASE 		=   UnCLIPImageVariationPipeline
       SCREAMING_SNAKE_CASE 		=   IMAGE_VARIATION_PARAMS - {"height", "width", "guidance_scale"}
       SCREAMING_SNAKE_CASE 		=   IMAGE_VARIATION_BATCH_PARAMS

       SCREAMING_SNAKE_CASE 		=   [
           "generator",
           "return_dict",
           "decoder_num_inference_steps",
           "super_res_num_inference_steps",
       ]
       SCREAMING_SNAKE_CASE 		=   False

       @property
       def _lowerCAmelCase(						self			)  ->  Tuple:
        return 32

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

       @property
       def _lowerCAmelCase(						self			)  ->  str:
        return self.time_input_dim

       @property
       def _lowerCAmelCase(						self			)  ->  int:
        return self.time_input_dim * 4

       @property
       def _lowerCAmelCase(						self			)  ->  Dict:
        return 100

       @property
       def _lowerCAmelCase(						self			)  ->  List[Any]:
        lowercase__    :   Dict						= CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''			)
        return tokenizer

       @property
       def _lowerCAmelCase(						self			)  ->  Any:
        torch.manual_seed(0			)
        lowercase__    :   Optional[Any]						= 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(__lowerCAmelCase			)

       @property
       def _lowerCAmelCase(						self			)  ->  Tuple:
        torch.manual_seed(0			)
        lowercase__    :   Optional[Any]						= 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(__lowerCAmelCase			)

       @property
       def _lowerCAmelCase(						self			)  ->  Optional[Any]:
        torch.manual_seed(0			)

        lowercase__    :   str						= {
            '''clip_embeddings_dim''': self.text_embedder_hidden_size,
            '''time_embed_dim''': self.time_embed_dim,
            '''cross_attention_dim''': self.cross_attention_dim,
        }

        lowercase__    :   Union[str, Any]						= UnCLIPTextProjModel(**__lowerCAmelCase			)
        return model

       @property
       def _lowerCAmelCase(						self			)  ->  Optional[int]:
        torch.manual_seed(0			)

        lowercase__    :   Any						= {
            '''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''',
        }

        lowercase__    :   str						= UNetaDConditionModel(**__lowerCAmelCase			)
        return model

       @property
       def _lowerCAmelCase(						self			)  ->  Dict:
        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 _lowerCAmelCase(						self			)  ->  Tuple:
        torch.manual_seed(0			)

        lowercase__    :   Optional[int]						= UNetaDModel(**self.dummy_super_res_kwargs			)
        return model

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

        lowercase__    :   Union[str, Any]						= UNetaDModel(**self.dummy_super_res_kwargs			)
        return model

       def _lowerCAmelCase(						self			)  ->  Optional[Any]:
        lowercase__    :   List[Any]						= self.dummy_decoder
        lowercase__    :   Dict						= self.dummy_text_proj
        lowercase__    :   str						= self.dummy_text_encoder
        lowercase__    :   Dict						= self.dummy_tokenizer
        lowercase__    :   Union[str, Any]						= self.dummy_super_res_first
        lowercase__    :   Optional[int]						= self.dummy_super_res_last

        lowercase__    :   Optional[int]						= UnCLIPScheduler(
            variance_type='''learned_range'''     ,       prediction_type='''epsilon'''     ,       num_train_timesteps=1000     ,       )

        lowercase__    :   Optional[Any]						= UnCLIPScheduler(
            variance_type='''fixed_small_log'''     ,       prediction_type='''epsilon'''     ,       num_train_timesteps=1000     ,       )

        lowercase__    :   Any						= CLIPImageProcessor(crop_size=32     ,       size=32			)

        lowercase__    :   List[str]						= 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 _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase=0     ,       __lowerCAmelCase=True			)  ->  List[str]:
        lowercase__    :   Dict						= floats_tensor((1, 3, 32, 32)     ,       rng=random.Random(__lowerCAmelCase			)			).to(__lowerCAmelCase			)
        if str(__lowerCAmelCase			).startswith('''mps'''			):
         lowercase__    :   int						= torch.manual_seed(__lowerCAmelCase			)
        else:
         lowercase__    :   int						= torch.Generator(device=__lowerCAmelCase			).manual_seed(__lowerCAmelCase			)

        if pil_image:
         lowercase__    :   Optional[int]						= input_image * 0.5 + 0.5
         lowercase__    :   Optional[Any]						= input_image.clamp(0     ,       1			)
         lowercase__    :   List[Any]						= input_image.cpu().permute(0     ,       2     ,       3     ,       1			).float().numpy()
         lowercase__    :   int						= DiffusionPipeline.numpy_to_pil(__lowerCAmelCase			)[0]

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

       def _lowerCAmelCase(						self			)  ->  Any:
        lowercase__    :   Optional[int]						= '''cpu'''

        lowercase__    :   Any						= self.get_dummy_components()

        lowercase__    :   Dict						= self.pipeline_class(**__lowerCAmelCase			)
        lowercase__    :   Optional[int]						= pipe.to(__lowerCAmelCase			)

        pipe.set_progress_bar_config(disable=__lowerCAmelCase			)

        lowercase__    :   Optional[int]						= self.get_dummy_inputs(__lowerCAmelCase     ,       pil_image=__lowerCAmelCase			)

        lowercase__    :   Dict						= pipe(**__lowerCAmelCase			)
        lowercase__    :   Dict						= output.images

        lowercase__    :   Union[str, Any]						= self.get_dummy_inputs(__lowerCAmelCase     ,       pil_image=__lowerCAmelCase			)

        lowercase__    :   Tuple						= pipe(
            **__lowerCAmelCase     ,       return_dict=__lowerCAmelCase     ,       )[0]

        lowercase__    :   str						= image[0, -3:, -3:, -1]
        lowercase__    :   List[str]						= image_from_tuple[0, -3:, -3:, -1]

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

        lowercase__    :   List[Any]						= np.array(
            [
                0.9_9_9_7,
                0.0_0_0_2,
                0.9_9_9_7,
                0.9_9_9_7,
                0.9_9_6_9,
                0.0_0_2_3,
                0.9_9_9_7,
                0.9_9_6_9,
                0.9_9_7_0,
            ]			)

        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 _lowerCAmelCase(						self			)  ->  Tuple:
        lowercase__    :   Union[str, Any]						= '''cpu'''

        lowercase__    :   List[Any]						= self.get_dummy_components()

        lowercase__    :   List[str]						= self.pipeline_class(**__lowerCAmelCase			)
        lowercase__    :   List[Any]						= pipe.to(__lowerCAmelCase			)

        pipe.set_progress_bar_config(disable=__lowerCAmelCase			)

        lowercase__    :   Any						= self.get_dummy_inputs(__lowerCAmelCase     ,       pil_image=__lowerCAmelCase			)

        lowercase__    :   Optional[Any]						= pipe(**__lowerCAmelCase			)
        lowercase__    :   Tuple						= output.images

        lowercase__    :   Dict						= self.get_dummy_inputs(__lowerCAmelCase     ,       pil_image=__lowerCAmelCase			)

        lowercase__    :   Optional[int]						= pipe(
            **__lowerCAmelCase     ,       return_dict=__lowerCAmelCase     ,       )[0]

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

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

        lowercase__    :   Optional[Any]						= np.array([0.9_9_9_7, 0.0_0_0_3, 0.9_9_9_7, 0.9_9_9_7, 0.9_9_7_0, 0.0_0_2_4, 0.9_9_9_7, 0.9_9_7_1, 0.9_9_7_1]			)

        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 _lowerCAmelCase(						self			)  ->  str:
        lowercase__    :   Optional[int]						= '''cpu'''

        lowercase__    :   Optional[int]						= self.get_dummy_components()

        lowercase__    :   str						= self.pipeline_class(**__lowerCAmelCase			)
        lowercase__    :   Tuple						= pipe.to(__lowerCAmelCase			)

        pipe.set_progress_bar_config(disable=__lowerCAmelCase			)

        lowercase__    :   Dict						= self.get_dummy_inputs(__lowerCAmelCase     ,       pil_image=__lowerCAmelCase			)
        lowercase__    :   Union[str, Any]						= [
            pipeline_inputs['''image'''],
            pipeline_inputs['''image'''],
        ]

        lowercase__    :   Optional[Any]						= pipe(**__lowerCAmelCase			)
        lowercase__    :   int						= output.images

        lowercase__    :   Any						= self.get_dummy_inputs(__lowerCAmelCase     ,       pil_image=__lowerCAmelCase			)
        lowercase__    :   Union[str, Any]						= [
            tuple_pipeline_inputs['''image'''],
            tuple_pipeline_inputs['''image'''],
        ]

        lowercase__    :   str						= pipe(
            **__lowerCAmelCase     ,       return_dict=__lowerCAmelCase     ,       )[0]

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

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

        lowercase__    :   str						= np.array(
            [
                0.9_9_9_7,
                0.9_9_8_9,
                0.0_0_0_8,
                0.0_0_2_1,
                0.9_9_6_0,
                0.0_0_1_8,
                0.0_0_1_4,
                0.0_0_0_2,
                0.9_9_3_3,
            ]			)

        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 _lowerCAmelCase(						self			)  ->  Dict:
        lowercase__    :   Union[str, Any]						= torch.device('''cpu'''			)



        class       UpperCAmelCase  :





               '''simple docstring'''



               SCREAMING_SNAKE_CASE 		=   1


        lowercase__    :   str						= self.get_dummy_components()

        lowercase__    :   Union[str, Any]						= self.pipeline_class(**__lowerCAmelCase			)
        lowercase__    :   Optional[Any]						= pipe.to(__lowerCAmelCase			)

        pipe.set_progress_bar_config(disable=__lowerCAmelCase			)

        lowercase__    :   Optional[Any]						= torch.Generator(device=__lowerCAmelCase			).manual_seed(0			)
        lowercase__    :   str						= pipe.decoder.dtype
        lowercase__    :   Tuple						= 1

        lowercase__    :   Tuple						= (
            batch_size,
            pipe.decoder.config.in_channels,
            pipe.decoder.config.sample_size,
            pipe.decoder.config.sample_size,
        )
        lowercase__    :   List[str]						= pipe.prepare_latents(
            __lowerCAmelCase     ,       dtype=__lowerCAmelCase     ,       device=__lowerCAmelCase     ,       generator=__lowerCAmelCase     ,       latents=__lowerCAmelCase     ,       scheduler=DummyScheduler()			)

        lowercase__    :   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,
        )
        lowercase__    :   Any						= pipe.prepare_latents(
            __lowerCAmelCase     ,       dtype=__lowerCAmelCase     ,       device=__lowerCAmelCase     ,       generator=__lowerCAmelCase     ,       latents=__lowerCAmelCase     ,       scheduler=DummyScheduler()			)

        lowercase__    :   Any						= self.get_dummy_inputs(__lowerCAmelCase     ,       pil_image=__lowerCAmelCase			)

        lowercase__    :   Union[str, Any]						= pipe(
            **__lowerCAmelCase     ,       decoder_latents=__lowerCAmelCase     ,       super_res_latents=__lowerCAmelCase			).images

        lowercase__    :   List[str]						= self.get_dummy_inputs(__lowerCAmelCase     ,       pil_image=__lowerCAmelCase			)
        # Don't pass image, instead pass embedding
        lowercase__    :   List[str]						= pipeline_inputs.pop('''image'''			)
        lowercase__    :   str						= pipe.image_encoder(__lowerCAmelCase			).image_embeds

        lowercase__    :   List[Any]						= pipe(
            **__lowerCAmelCase     ,       decoder_latents=__lowerCAmelCase     ,       super_res_latents=__lowerCAmelCase     ,       image_embeddings=__lowerCAmelCase     ,       ).images

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

       @skip_mps
       def _lowerCAmelCase(						self			)  ->  List[str]:
        lowercase__    :   Dict						= torch_device == '''cpu'''

        # Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
        lowercase__    :   List[Any]						= 1E-2

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

       @skip_mps
       def _lowerCAmelCase(						self			)  ->  int:
        lowercase__    :   int						= torch_device == '''cpu'''
        lowercase__    :   Dict						= True
        lowercase__    :   List[str]						= [
            '''decoder_num_inference_steps''',
            '''super_res_num_inference_steps''',
        ]

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

       def _lowerCAmelCase(						self			)  ->  Dict:
        lowercase__    :   List[Any]						= [
            '''decoder_num_inference_steps''',
            '''super_res_num_inference_steps''',
        ]

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

       @skip_mps
       def _lowerCAmelCase(						self			)  ->  int:
        return super().test_dict_tuple_outputs_equivalent()

       @skip_mps
       def _lowerCAmelCase(						self			)  ->  Optional[Any]:
        return super().test_save_load_local()


       @skip_mps
       def _lowerCAmelCase(						self			)  ->  List[str]:
        return super().test_save_load_optional_components()



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





       '''simple docstring'''



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


       def _lowerCAmelCase(						self			)  ->  Optional[int]:
        lowercase__    :   Dict						= load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png'''			)
        lowercase__    :   str						= load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/unclip/karlo_v1_alpha_cat_variation_fp16.npy'''			)

        lowercase__    :   Union[str, Any]						= UnCLIPImageVariationPipeline.from_pretrained(
            '''kakaobrain/karlo-v1-alpha-image-variations'''     ,       torch_dtype=torch.floataa			)
        lowercase__    :   Union[str, Any]						= pipeline.to(__lowerCAmelCase			)
        pipeline.set_progress_bar_config(disable=__lowerCAmelCase			)

        lowercase__    :   List[str]						= torch.Generator(device='''cpu'''			).manual_seed(0			)
        lowercase__    :   Any						= pipeline(
            __lowerCAmelCase     ,       generator=__lowerCAmelCase     ,       output_type='''np'''     ,       )

        lowercase__    :   str						= output.images[0]

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

        assert_mean_pixel_difference(__lowerCAmelCase     ,       __lowerCAmelCase     ,       15			)


'''simple docstring'''

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       UpperCAmelCase  (						tf.keras.optimizers.schedules.LearningRateSchedule ):





       '''simple docstring'''



       def __init__(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase = 1.0     ,       __lowerCAmelCase = None     ,       )  ->  Union[str, Any]:
        super().__init__()
        lowercase__    :   Any						= initial_learning_rate
        lowercase__    :   str						= warmup_steps
        lowercase__    :   Dict						= power
        lowercase__    :   int						= decay_schedule_fn
        lowercase__    :   Dict						= name

       def __call__(						self     ,       __lowerCAmelCase			)  ->  int:
        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`.
         lowercase__    :   Union[str, Any]						= tf.cast(__lowerCAmelCase     ,       tf.floataa			)
         lowercase__    :   Optional[int]						= tf.cast(self.warmup_steps     ,       tf.floataa			)
         lowercase__    :   Dict						= global_step_float / warmup_steps_float
         lowercase__    :   Dict						= 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			)  ->  List[str]:
        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	(       UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase		,					UpperCAmelCase = 0.0		,					UpperCAmelCase = 0.9		,					UpperCAmelCase = 0.9_9_9		,					UpperCAmelCase = 1E-8		,					UpperCAmelCase = None		,					UpperCAmelCase = None		,					UpperCAmelCase = 0.0		,					UpperCAmelCase = 1.0		,					UpperCAmelCase = None		,					):
 lowercase__    :   List[Any]						= tf.keras.optimizers.schedules.PolynomialDecay(
     initial_learning_rate=UpperCAmelCase		,					decay_steps=num_train_steps - num_warmup_steps		,					end_learning_rate=init_lr * min_lr_ratio		,					power=UpperCAmelCase		,					)
 if num_warmup_steps:
  lowercase__    :   Optional[Any]						= WarmUp(
      initial_learning_rate=UpperCAmelCase		,					decay_schedule_fn=UpperCAmelCase		,					warmup_steps=UpperCAmelCase		,					)
 if weight_decay_rate > 0.0:
  lowercase__    :   str						= AdamWeightDecay(
      learning_rate=UpperCAmelCase		,					weight_decay_rate=UpperCAmelCase		,					beta_a=UpperCAmelCase		,					beta_a=UpperCAmelCase		,					epsilon=UpperCAmelCase		,					clipnorm=UpperCAmelCase		,					global_clipnorm=UpperCAmelCase		,					exclude_from_weight_decay=['''LayerNorm''', '''layer_norm''', '''bias''']		,					include_in_weight_decay=UpperCAmelCase		,					)
 else:
  lowercase__    :   Any						= tf.keras.optimizers.Adam(
      learning_rate=UpperCAmelCase		,					beta_a=UpperCAmelCase		,					beta_a=UpperCAmelCase		,					epsilon=UpperCAmelCase		,					clipnorm=UpperCAmelCase		,					global_clipnorm=UpperCAmelCase		,					)
 # 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       UpperCAmelCase  (						a__ ):





       '''simple docstring'''



       def __init__(						self     ,       __lowerCAmelCase = 0.0_0_1     ,       __lowerCAmelCase = 0.9     ,       __lowerCAmelCase = 0.9_9_9     ,       __lowerCAmelCase = 1E-7     ,       __lowerCAmelCase = False     ,       __lowerCAmelCase = 0.0     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = None     ,       __lowerCAmelCase = "AdamWeightDecay"     ,       **__lowerCAmelCase     ,       )  ->  str:
        super().__init__(__lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       **__lowerCAmelCase			)
        lowercase__    :   int						= weight_decay_rate
        lowercase__    :   Dict						= include_in_weight_decay
        lowercase__    :   List[str]						= exclude_from_weight_decay

       @classmethod
       def _lowerCAmelCase(						cls     ,       __lowerCAmelCase			)  ->  str:
        lowercase__    :   int						= {'''WarmUp''': WarmUp}
        return super(__lowerCAmelCase     ,       cls			).from_config(__lowerCAmelCase     ,       custom_objects=__lowerCAmelCase			)

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase			)  ->  List[str]:
        super(__lowerCAmelCase     ,       self			)._prepare_local(__lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase			)
        lowercase__    :   Any						= tf.constant(
            self.weight_decay_rate     ,       name='''adam_weight_decay_rate'''			)

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase			)  ->  Tuple:
        lowercase__    :   str						= 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			)  ->  List[Any]:
        lowercase__			,							lowercase__    :   Any						= list(zip(*__lowerCAmelCase			)			)
        return super(__lowerCAmelCase     ,       self			).apply_gradients(zip(__lowerCAmelCase     ,       __lowerCAmelCase			)     ,       name=__lowerCAmelCase     ,       **__lowerCAmelCase			)

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase			)  ->  Optional[int]:
        if apply_state is None:
         return self._decayed_lr_t[var_dtype], {}

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

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

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase=None			)  ->  Any:
        lowercase__			,							lowercase__    :   Dict						= self._get_lr(var.device     ,       var.dtype.base_dtype     ,       __lowerCAmelCase			)
        lowercase__    :   List[str]						= 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			)  ->  List[str]:
        lowercase__			,							lowercase__    :   List[Any]						= self._get_lr(var.device     ,       var.dtype.base_dtype     ,       __lowerCAmelCase			)
        lowercase__    :   Union[str, 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			)  ->  List[str]:
        lowercase__    :   Optional[Any]						= super().get_config()
        config.update({'''weight_decay_rate''': self.weight_decay_rate}			)
        return config


       def _lowerCAmelCase(						self     ,       __lowerCAmelCase			)  ->  List[Any]:
        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       UpperCAmelCase  (						a__ ):





       '''simple docstring'''



       def __init__(						self			)  ->  List[Any]:
        lowercase__    :   List[str]						= []
        lowercase__    :   Optional[int]						= None

       @property
       def _lowerCAmelCase(						self			)  ->  Optional[Any]:
        if self._accum_steps is None:
         lowercase__    :   int						= 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			)  ->  Union[str, Any]:
        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			)  ->  Optional[Any]:
        if not self._gradients:
         lowercase__    :   str						= 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			)  ->  Optional[Any]:
        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			)			)

