




import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import UperNetForSemanticSegmentation
	from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from PIL import Image

	from transformers import AutoImageProcessor


class    lowerCamelCase    :






			'''simple docstring'''





			def __init__( self			:   Optional[int]       ,						lowerCamelCase_			:   Optional[Any]       ,						lowerCamelCase_			:   Optional[Any]=13       ,						lowerCamelCase_			:   Union[str, Any]=32       ,						lowerCamelCase_			:   Optional[Any]=3       ,						lowerCamelCase_			:   Dict=4       ,						lowerCamelCase_			:   Optional[int]=[10, 20, 30, 40]       ,						lowerCamelCase_			:   Optional[Any]=[2, 2, 3, 2]       ,						lowerCamelCase_			:   int=True       ,						lowerCamelCase_			:   Optional[Any]=True       ,						lowerCamelCase_			:   int=37       ,						lowerCamelCase_			:   str="gelu"       ,						lowerCamelCase_			:   int=10       ,						lowerCamelCase_			:   int=0.0_2       ,						lowerCamelCase_			:   Dict=["stage2", "stage3", "stage4"]       ,						lowerCamelCase_			:   int=3       ,						lowerCamelCase_			:   Union[str, Any]=None       ,						)						->   Tuple:
								__magic_name__							:	Optional[Any]    							=							parent
								__magic_name__							:	Union[str, Any]    							=							batch_size
								__magic_name__							:	Union[str, Any]    							=							image_size
								__magic_name__							:	int    							=							num_channels
								__magic_name__							:	List[str]    							=							num_stages
								__magic_name__							:	Dict    							=							hidden_sizes
								__magic_name__							:	str    							=							depths
								__magic_name__							:	List[Any]    							=							is_training
								__magic_name__							:	int    							=							use_labels
								__magic_name__							:	Dict    							=							intermediate_size
								__magic_name__							:	Tuple    							=							hidden_act
								__magic_name__							:	List[Any]    							=							type_sequence_label_size
								__magic_name__							:	List[Any]    							=							initializer_range
								__magic_name__							:	List[str]    							=							out_features
								__magic_name__							:	str    							=							num_labels
								__magic_name__							:	List[Any]    							=							scope
								__magic_name__							:	Any    							=							num_stages

			def 							UpperCAmelCase__       ( self			:   List[Any]  )						->   Optional[Any]:
								__magic_name__							:	Optional[int]    							=							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

								__magic_name__							:	List[Any]    							=							None
								if self.use_labels:
													__magic_name__							:	str    							=							ids_tensor([self.batch_size]       ,						self.type_sequence_label_size  )

								__magic_name__							:	List[Any]    							=							self.get_config()

								return config, pixel_values, labels

			def 							UpperCAmelCase__       ( self			:   Any  )						->   Union[str, Any]:
								return ConvNextConfig(
								    num_channels=self.num_channels       ,						num_stages=self.num_stages       ,						hidden_sizes=self.hidden_sizes       ,						depths=self.depths       ,						is_training=self.is_training       ,						intermediate_size=self.intermediate_size       ,						hidden_act=self.hidden_act       ,						out_features=self.out_features       ,						)

			def 							UpperCAmelCase__       ( self			:   List[str]  )						->   Union[str, Any]:
								return UperNetConfig(
								    backbone_config=self.get_backbone_config()       ,						hidden_size=512       ,						pool_scales=[1, 2, 3, 6]       ,						use_auxiliary_head=lowerCamelCase_       ,						auxiliary_loss_weight=0.4       ,						auxiliary_in_channels=40       ,						auxiliary_channels=256       ,						auxiliary_num_convs=1       ,						auxiliary_concat_input=lowerCamelCase_       ,						loss_ignore_index=255       ,						num_labels=self.num_labels       ,						)

			def 							UpperCAmelCase__       ( self			:   str       ,						lowerCamelCase_			:   str       ,						lowerCamelCase_			:   Dict       ,						lowerCamelCase_			:   Any  )						->   List[str]:
								__magic_name__							:	Any    							=							UperNetForSemanticSegmentation(config=lowerCamelCase_  )
								model.to(lowerCamelCase_  )
								model.eval()
								__magic_name__							:	int    							=							model(lowerCamelCase_  )
								self.parent.assertEqual(
								    result.logits.shape       ,						(self.batch_size, self.num_labels, self.image_size, self.image_size)  )


			def 							UpperCAmelCase__       ( self			:   List[Any]  )						->   Optional[int]:
								__magic_name__							:	Tuple    							=							self.prepare_config_and_inputs()
								(
								    __magic_name__    
								)							:	Dict    							=							config_and_inputs
								__magic_name__							:	Tuple    							=							{'''pixel_values''': pixel_values}
								return config, inputs_dict






@require_torch
class    lowerCamelCase    (		_lowerCamelCase				,_lowerCamelCase				,unittest.TestCase						):






			'''simple docstring'''





			UpperCamelCase__								=(UperNetForSemanticSegmentation,) if is_torch_available() else ()
			UpperCamelCase__								={'''image-segmentation''': UperNetForSemanticSegmentation} if is_torch_available() else {}
			UpperCamelCase__								=False
			UpperCamelCase__								=False
			UpperCamelCase__								=False
			UpperCamelCase__								=False
			UpperCamelCase__								=False
			UpperCamelCase__								=False

			def 							UpperCAmelCase__       ( self			:   Optional[Any]  )						->   Any:
								__magic_name__							:	Dict    							=							UperNetModelTester(self  )
								__magic_name__							:	int    							=							ConfigTester(self       ,						config_class=lowerCamelCase_       ,						has_text_modality=lowerCamelCase_       ,						hidden_size=37  )

			def 							UpperCAmelCase__       ( self			:   List[Any]  )						->   Union[str, Any]:
								self.create_and_test_config_common_properties()
								self.config_tester.create_and_test_config_to_json_string()
								self.config_tester.create_and_test_config_to_json_file()
								self.config_tester.create_and_test_config_from_and_save_pretrained()
								self.config_tester.create_and_test_config_with_num_labels()
								self.config_tester.check_config_can_be_init_without_params()
								self.config_tester.check_config_arguments_init()

			def 							UpperCAmelCase__       ( self			:   List[Any]  )						->   Union[str, Any]:
								return

			def 							UpperCAmelCase__       ( self			:   Optional[Any]  )						->   Dict:
								__magic_name__							:	Optional[int]    							=							self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
													__magic_name__							:	List[Any]    							=							model_class(lowerCamelCase_  )
													__magic_name__							:	List[str]    							=							inspect.signature(model.forward  )
													# signature.parameters is an OrderedDict => so arg_names order is deterministic
													__magic_name__							:	Optional[int]    							=							[*signature.parameters.keys()]

													__magic_name__							:	int    							=							['''pixel_values''']
													self.assertListEqual(arg_names[:1]       ,						lowerCamelCase_  )

			def 							UpperCAmelCase__       ( self			:   List[str]  )						->   List[Any]:
								__magic_name__							:	Union[str, Any]    							=							self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_semantic_segmentation(*lowerCamelCase_  )

			@unittest.skip(reason='''UperNet does not use inputs_embeds'''  )
			def 							UpperCAmelCase__       ( self			:   int  )						->   List[str]:
								pass

			@unittest.skip(reason='''UperNet does not support input and output embeddings'''  )
			def 							UpperCAmelCase__       ( self			:   str  )						->   List[Any]:
								pass

			@unittest.skip(reason='''UperNet does not have a base model'''  )
			def 							UpperCAmelCase__       ( self			:   Optional[Any]  )						->   Optional[Any]:
								pass

			@unittest.skip(reason='''UperNet does not have a base model'''  )
			def 							UpperCAmelCase__       ( self			:   int  )						->   List[Any]:
								pass

			@require_torch_multi_gpu
			@unittest.skip(reason='''UperNet has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''  )
			def 							UpperCAmelCase__       ( self			:   Union[str, Any]  )						->   int:
								pass

			@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''  )
			def 							UpperCAmelCase__       ( self			:   List[Any]  )						->   List[str]:
								pass

			def 							UpperCAmelCase__       ( self			:   Any  )						->   Any:
								def check_hidden_states_output(lowerCamelCase_			:   List[str]       ,						lowerCamelCase_			:   Union[str, Any]       ,						lowerCamelCase_			:   Union[str, Any]  ):
													__magic_name__							:	Optional[Any]    							=							model_class(lowerCamelCase_  )
													model.to(lowerCamelCase_  )
													model.eval()

													with torch.no_grad():
																		__magic_name__							:	Tuple    							=							model(**self._prepare_for_class(lowerCamelCase_       ,						lowerCamelCase_  )  )

													__magic_name__							:	List[str]    							=							outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

													__magic_name__							:	Any    							=							self.model_tester.num_stages
													self.assertEqual(len(lowerCamelCase_  )       ,						expected_num_stages + 1  )

													# ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
													self.assertListEqual(
													    list(hidden_states[0].shape[-2:]  )       ,						[self.model_tester.image_size // 4, self.model_tester.image_size // 4]       ,						)

								__magic_name__							:	Union[str, Any]    							=							self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
													__magic_name__							:	int    							=							True
													check_hidden_states_output(lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_  )

													# check that output_hidden_states also work using config
													del inputs_dict["output_hidden_states"]
													__magic_name__							:	str    							=							True

													check_hidden_states_output(lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_  )

			def 							UpperCAmelCase__       ( self			:   Tuple  )						->   List[Any]:
								__magic_name__							:	List[Any]    							=							self.model_tester.prepare_config_and_inputs_for_common()

								__magic_name__							:	str    							=							_config_zero_init(lowerCamelCase_  )
								__magic_name__							:	List[Any]    							=							_config_zero_init(configs_no_init.backbone_config  )
								for model_class in self.all_model_classes:
													__magic_name__							:	List[Any]    							=							model_class(config=lowerCamelCase_  )
													for name, param in model.named_parameters():
																		if param.requires_grad:
																							self.assertIn(
																							    ((param.data.mean() * 1E9).round() / 1E9).item()       ,						[0.0, 1.0]       ,						msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''       ,						)

			@unittest.skip(reason='''UperNet does not have tied weights'''  )
			def 							UpperCAmelCase__       ( self			:   Any  )						->   List[str]:
								pass


			@slow
			def 							UpperCAmelCase__       ( self			:   List[str]  )						->   Union[str, Any]:
								for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													__magic_name__							:	List[Any]    							=							UperNetForSemanticSegmentation.from_pretrained(lowerCamelCase_  )
													self.assertIsNotNone(lowerCamelCase_  )







def 		lowercase__   (		):





					'''simple docstring'''






					__magic_name__							:	Optional[Any]    							=							hf_hub_download(
					    repo_id='''hf-internal-testing/fixtures_ade20k'''							,repo_type='''dataset'''							,filename='''ADE_val_00000001.jpg'''				)
					__magic_name__							:	List[str]    							=							Image.open(__A				).convert('''RGB'''				)
					return image


@require_torch
@require_vision
@slow
class    lowerCamelCase    (		unittest.TestCase						):






			'''simple docstring'''





			def 							UpperCAmelCase__       ( self			:   List[Any]  )						->   List[Any]:
								__magic_name__							:	Optional[int]    							=							AutoImageProcessor.from_pretrained('''openmmlab/upernet-swin-tiny'''  )
								__magic_name__							:	str    							=							UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-swin-tiny'''  ).to(lowerCamelCase_  )

								__magic_name__							:	List[Any]    							=							prepare_img()
								__magic_name__							:	Union[str, Any]    							=							processor(images=lowerCamelCase_       ,						return_tensors='''pt'''  ).to(lowerCamelCase_  )

								with torch.no_grad():
													__magic_name__							:	int    							=							model(**lowerCamelCase_  )

								__magic_name__							:	str    							=							torch.Size((1, model.config.num_labels, 512, 512)  )
								self.assertEqual(outputs.logits.shape       ,						lowerCamelCase_  )

								__magic_name__							:	Tuple    							=							torch.tensor(
								    [[-7.5_9_5_8, -7.5_9_5_8, -7.4_3_0_2], [-7.5_9_5_8, -7.5_9_5_8, -7.4_3_0_2], [-7.4_7_9_7, -7.4_7_9_7, -7.3_0_6_8]]  ).to(lowerCamelCase_  )
								self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]       ,						lowerCamelCase_       ,						atol=1E-4  )  )


			def 							UpperCAmelCase__       ( self			:   Optional[int]  )						->   str:
								__magic_name__							:	List[str]    							=							AutoImageProcessor.from_pretrained('''openmmlab/upernet-convnext-tiny'''  )
								__magic_name__							:	Optional[int]    							=							UperNetForSemanticSegmentation.from_pretrained('''openmmlab/upernet-convnext-tiny'''  ).to(lowerCamelCase_  )

								__magic_name__							:	Tuple    							=							prepare_img()
								__magic_name__							:	List[str]    							=							processor(images=lowerCamelCase_       ,						return_tensors='''pt'''  ).to(lowerCamelCase_  )

								with torch.no_grad():
													__magic_name__							:	Optional[Any]    							=							model(**lowerCamelCase_  )

								__magic_name__							:	Optional[int]    							=							torch.Size((1, model.config.num_labels, 512, 512)  )
								self.assertEqual(outputs.logits.shape       ,						lowerCamelCase_  )

								__magic_name__							:	Tuple    							=							torch.tensor(
								    [[-8.8_1_1_0, -8.8_1_1_0, -8.6_5_2_1], [-8.8_1_1_0, -8.8_1_1_0, -8.6_5_2_1], [-8.7_7_4_6, -8.7_7_4_6, -8.6_1_3_0]]  ).to(lowerCamelCase_  )
								self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]       ,						lowerCamelCase_       ,						atol=1E-4  )  )









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






			'''simple docstring'''





			def __init__( self			:   Union[str, Any]       ,						lowerCamelCase_			:   float       ,						lowerCamelCase_			:   Callable       ,						lowerCamelCase_			:   int       ,						lowerCamelCase_			:   float = 1.0       ,						lowerCamelCase_			:   str = None       ,						)						->   Optional[int]:
								super().__init__()
								__magic_name__							:	Union[str, Any]    							=							initial_learning_rate
								__magic_name__							:	Optional[int]    							=							warmup_steps
								__magic_name__							:	Union[str, Any]    							=							power
								__magic_name__							:	List[Any]    							=							decay_schedule_fn
								__magic_name__							:	Dict    							=							name

			def __call__( self			:   Any       ,						lowerCamelCase_			:   Dict  )						->   Optional[Any]:
								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`.
													__magic_name__							:	Union[str, Any]    							=							tf.cast(lowerCamelCase_       ,						tf.floataa  )
													__magic_name__							:	Dict    							=							tf.cast(self.warmup_steps       ,						tf.floataa  )
													__magic_name__							:	Dict    							=							global_step_float / warmup_steps_float
													__magic_name__							:	List[str]    							=							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 							UpperCAmelCase__       ( self			:   Optional[int]  )						->   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 		lowercase__   (		__A:	float							,__A:	int							,__A:	int							,__A:	float = 0.0							,__A:	float = 0.9							,__A:	float = 0.999							,__A:	float = 1e-8							,__A:	Optional[float] = None							,__A:	Optional[float] = None							,__A:	float = 0.0							,__A:	float = 1.0							,__A:	Optional[List[str]] = None							,):





					'''simple docstring'''






					__magic_name__							:	Any    							=							tf.keras.optimizers.schedules.PolynomialDecay(
					    initial_learning_rate=__A							,decay_steps=num_train_steps - num_warmup_steps							,end_learning_rate=init_lr * min_lr_ratio							,power=__A							,)
					if num_warmup_steps:
										__magic_name__							:	Tuple    							=							WarmUp(
										    initial_learning_rate=__A							,decay_schedule_fn=__A							,warmup_steps=__A							,)
					if weight_decay_rate > 0.0:
										__magic_name__							:	Union[str, Any]    							=							AdamWeightDecay(
										    learning_rate=__A							,weight_decay_rate=__A							,beta_a=__A							,beta_a=__A							,epsilon=__A							,clipnorm=__A							,global_clipnorm=__A							,exclude_from_weight_decay=['''LayerNorm''', '''layer_norm''', '''bias''']							,include_in_weight_decay=__A							,)
					else:
										__magic_name__							:	str    							=							tf.keras.optimizers.Adam(
										    learning_rate=__A							,beta_a=__A							,beta_a=__A							,epsilon=__A							,clipnorm=__A							,global_clipnorm=__A							,)
					# 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    lowerCamelCase    (		_lowerCamelCase						):






			'''simple docstring'''





			def __init__( self			:   Optional[int]       ,						lowerCamelCase_			:   Union[float, tf.keras.optimizers.schedules.LearningRateSchedule] = 0.0_0_1       ,						lowerCamelCase_			:   float = 0.9       ,						lowerCamelCase_			:   float = 0.9_9_9       ,						lowerCamelCase_			:   float = 1E-7       ,						lowerCamelCase_			:   bool = False       ,						lowerCamelCase_			:   float = 0.0       ,						lowerCamelCase_			:   Optional[List[str]] = None       ,						lowerCamelCase_			:   Optional[List[str]] = None       ,						lowerCamelCase_			:   str = "AdamWeightDecay"       ,						**lowerCamelCase_			:   int       ,						)						->   List[str]:
								super().__init__(lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_       ,						**lowerCamelCase_  )
								__magic_name__							:	int    							=							weight_decay_rate
								__magic_name__							:	Tuple    							=							include_in_weight_decay
								__magic_name__							:	Union[str, Any]    							=							exclude_from_weight_decay

			@classmethod
			def 							UpperCAmelCase__       ( cls			:   int       ,						lowerCamelCase_			:   int  )						->   Optional[Any]:
								__magic_name__							:	Tuple    							=							{'''WarmUp''': WarmUp}
								return super(lowerCamelCase_       ,						cls  ).from_config(lowerCamelCase_       ,						custom_objects=lowerCamelCase_  )

			def 							UpperCAmelCase__       ( self			:   Union[str, Any]       ,						lowerCamelCase_			:   int       ,						lowerCamelCase_			:   Optional[int]       ,						lowerCamelCase_			:   Optional[int]  )						->   int:
								super(lowerCamelCase_       ,						self  )._prepare_local(lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_  )
								__magic_name__							:	Any    							=							tf.constant(
								    self.weight_decay_rate       ,						name='''adam_weight_decay_rate'''  )

			def 							UpperCAmelCase__       ( self			:   Optional[int]       ,						lowerCamelCase_			:   List[str]       ,						lowerCamelCase_			:   Optional[Any]       ,						lowerCamelCase_			:   Any  )						->   Optional[Any]:
								__magic_name__							:	Optional[int]    							=							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 							UpperCAmelCase__       ( self			:   str       ,						lowerCamelCase_			:   Optional[Any]       ,						lowerCamelCase_			:   int=None       ,						**lowerCamelCase_			:   Union[str, Any]  )						->   Optional[int]:
								__magic_name__							,    __magic_name__							:	Dict    							=							list(zip(*lowerCamelCase_  )  )
								return super(lowerCamelCase_       ,						self  ).apply_gradients(zip(lowerCamelCase_       ,						lowerCamelCase_  )       ,						name=lowerCamelCase_       ,						**lowerCamelCase_  )

			def 							UpperCAmelCase__       ( self			:   List[Any]       ,						lowerCamelCase_			:   Optional[int]       ,						lowerCamelCase_			:   int       ,						lowerCamelCase_			:   Union[str, Any]  )						->   Union[str, Any]:
								if apply_state is None:
													return self._decayed_lr_t[var_dtype], {}

								__magic_name__							:	str    							=							apply_state or {}
								__magic_name__							:	Optional[Any]    							=							apply_state.get((var_device, var_dtype)  )
								if coefficients is None:
													__magic_name__							:	List[str]    							=							self._fallback_apply_state(lowerCamelCase_       ,						lowerCamelCase_  )
													__magic_name__							:	List[Any]    							=							coefficients

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

			def 							UpperCAmelCase__       ( self			:   List[str]       ,						lowerCamelCase_			:   List[str]       ,						lowerCamelCase_			:   List[str]       ,						lowerCamelCase_			:   Dict=None  )						->   Optional[Any]:
								__magic_name__							,    __magic_name__							:	Union[str, Any]    							=							self._get_lr(var.device       ,						var.dtype.base_dtype       ,						lowerCamelCase_  )
								__magic_name__							:	Optional[Any]    							=							self._decay_weights_op(lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_  )
								with tf.control_dependencies([decay]  ):
													return super(lowerCamelCase_       ,						self  )._resource_apply_dense(lowerCamelCase_       ,						lowerCamelCase_       ,						**lowerCamelCase_  )

			def 							UpperCAmelCase__       ( self			:   Dict       ,						lowerCamelCase_			:   List[Any]       ,						lowerCamelCase_			:   List[Any]       ,						lowerCamelCase_			:   Optional[Any]       ,						lowerCamelCase_			:   Union[str, Any]=None  )						->   Any:
								__magic_name__							,    __magic_name__							:	Tuple    							=							self._get_lr(var.device       ,						var.dtype.base_dtype       ,						lowerCamelCase_  )
								__magic_name__							:	Optional[int]    							=							self._decay_weights_op(lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_  )
								with tf.control_dependencies([decay]  ):
													return super(lowerCamelCase_       ,						self  )._resource_apply_sparse(lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_       ,						**lowerCamelCase_  )

			def 							UpperCAmelCase__       ( self			:   int  )						->   Dict:
								__magic_name__							:	List[Any]    							=							super().get_config()
								config.update({'''weight_decay_rate''': self.weight_decay_rate}  )
								return config


			def 							UpperCAmelCase__       ( self			:   Dict       ,						lowerCamelCase_			:   Union[str, Any]  )						->   Union[str, 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    lowerCamelCase    (		_lowerCamelCase						):






			'''simple docstring'''





			def __init__( self			:   Optional[Any]  )						->   Any:
								__magic_name__							:	Optional[int]    							=							[]
								__magic_name__							:	Dict    							=							None

			@property
			def 							UpperCAmelCase__       ( self			:   Optional[int]  )						->   Any:
								if self._accum_steps is None:
													__magic_name__							:	Optional[Any]    							=							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 							UpperCAmelCase__       ( self			:   Optional[Any]  )						->   Optional[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			:   Optional[int]       ,						lowerCamelCase_			:   Any  )						->   List[str]:
								if not self._gradients:
													__magic_name__							:	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 							UpperCAmelCase__       ( self			:   str  )						->   Optional[int]:
								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_  )  )



