




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







							def __init__(  self	:  Optional[int]							,      lowerCAmelCase_	:  List[Any]							,      lowerCAmelCase_	:  Optional[int]							,      lowerCAmelCase_	:  int							,      lowerCAmelCase_	:  Any = 1.0							,      lowerCAmelCase_	:  Union[str, Any] = None							,      )      ->							Tuple:


														"""simple docstring"""


														super().__init__()
														_a     					=    initial_learning_rate
														_a     					=    warmup_steps
														_a     					=    power
														_a     					=    decay_schedule_fn
														_a     					=    name







							def __call__(  self	:  int							,      lowerCAmelCase_	:  Dict )      ->							List[Any]:


														"""simple docstring"""


														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     					=    tf.cast(_snake_case							,      tf.floataa )
																					_a     					=    tf.cast(self.warmup_steps							,      tf.floataa )
																					_a     					=    global_step_float / warmup_steps_float
																					_a     					=    self.initial_learning_rate * tf.math.pow(_snake_case							,      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=_snake_case							,      )







							def        __lowerCAmelCase      (  self	:  Optional[Any] )      ->							Optional[int]:


														"""simple docstring"""


														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 		snake_case_   (UpperCamelCase	:  float			,						UpperCamelCase	:  int			,						UpperCamelCase	:  int			,						UpperCamelCase	:  float = 0.0			,						UpperCamelCase	:  float = 0.9			,						UpperCamelCase	:  float = 0.999			,						UpperCamelCase	:  float = 1e-8			,						UpperCamelCase	:  Optional[float] = None			,						UpperCamelCase	:  Optional[float] = None			,						UpperCamelCase	:  float = 0.0			,						UpperCamelCase	:  float = 1.0			,						UpperCamelCase	:  Optional[List[str]] = None			,						):





							'''simple docstring'''


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







							def __init__(  self	:  Dict							,      lowerCAmelCase_	:  str = 0.0_0_1							,      lowerCAmelCase_	:  Any = 0.9							,      lowerCAmelCase_	:  str = 0.9_9_9							,      lowerCAmelCase_	:  Union[str, Any] = 1e-7							,      lowerCAmelCase_	:  Dict = False							,      lowerCAmelCase_	:  Tuple = 0.0							,      lowerCAmelCase_	:  List[Any] = None							,      lowerCAmelCase_	:  List[Any] = None							,      lowerCAmelCase_	:  str = "AdamWeightDecay"							,      **lowerCAmelCase_	:  List[str]							,      )      ->							List[str]:


														"""simple docstring"""


														super().__init__(_snake_case							,      _snake_case							,      _snake_case							,      _snake_case							,      _snake_case							,      _snake_case							,      **_snake_case )
														_a     					=    weight_decay_rate
														_a     					=    include_in_weight_decay
														_a     					=    exclude_from_weight_decay







							@classmethod
							def        __lowerCAmelCase      (  cls	:  int							,      lowerCAmelCase_	:  Tuple )      ->							Dict:


														"""simple docstring"""


														_a     					=    {"WarmUp": WarmUp}
														return super(_snake_case							,      cls ).from_config(_snake_case							,      custom_objects=_snake_case )







							def        __lowerCAmelCase      (  self	:  Union[str, Any]							,      lowerCAmelCase_	:  int							,      lowerCAmelCase_	:  Any							,      lowerCAmelCase_	:  str )      ->							int:


														"""simple docstring"""


														super(_snake_case							,      self )._prepare_local(_snake_case							,      _snake_case							,      _snake_case )
														_a     					=    tf.constant(
														    self.weight_decay_rate							,      name='''adam_weight_decay_rate''' )







							def        __lowerCAmelCase      (  self	:  str							,      lowerCAmelCase_	:  str							,      lowerCAmelCase_	:  Optional[int]							,      lowerCAmelCase_	:  int )      ->							List[Any]:


														"""simple docstring"""


														_a     					=    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	:  Union[str, Any]							,      lowerCAmelCase_	:  int							,      lowerCAmelCase_	:  Any=None							,      **lowerCAmelCase_	:  int )      ->							str:


														"""simple docstring"""


														_a     					=    list(zip(*_snake_case ) )
														return super(_snake_case							,      self ).apply_gradients(zip(_snake_case							,      _snake_case )							,      name=_snake_case							,      **_snake_case )







							def        __lowerCAmelCase      (  self	:  str							,      lowerCAmelCase_	:  int							,      lowerCAmelCase_	:  Tuple							,      lowerCAmelCase_	:  Any )      ->							Dict:


														"""simple docstring"""


														if apply_state is None:
																					return self._decayed_lr_t[var_dtype], {}

														_a     					=    apply_state or {}
														_a     					=    apply_state.get((var_device, var_dtype) )
														if coefficients is None:
																					_a     					=    self._fallback_apply_state(_snake_case							,      _snake_case )
																					_a     					=    coefficients

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







							def        __lowerCAmelCase      (  self	:  int							,      lowerCAmelCase_	:  int							,      lowerCAmelCase_	:  List[Any]							,      lowerCAmelCase_	:  Any=None )      ->							Tuple:


														"""simple docstring"""


														_a     					=    self._get_lr(var.device							,      var.dtype.base_dtype							,      _snake_case )
														_a     					=    self._decay_weights_op(_snake_case							,      _snake_case							,      _snake_case )
														with tf.control_dependencies([decay] ):
																					return super(_snake_case							,      self )._resource_apply_dense(_snake_case							,      _snake_case							,      **_snake_case )







							def        __lowerCAmelCase      (  self	:  List[str]							,      lowerCAmelCase_	:  List[str]							,      lowerCAmelCase_	:  Optional[int]							,      lowerCAmelCase_	:  Dict							,      lowerCAmelCase_	:  List[Any]=None )      ->							List[str]:


														"""simple docstring"""


														_a     					=    self._get_lr(var.device							,      var.dtype.base_dtype							,      _snake_case )
														_a     					=    self._decay_weights_op(_snake_case							,      _snake_case							,      _snake_case )
														with tf.control_dependencies([decay] ):
																					return super(_snake_case							,      self )._resource_apply_sparse(_snake_case							,      _snake_case							,      _snake_case							,      **_snake_case )







							def        __lowerCAmelCase      (  self	:  Optional[Any] )      ->							Any:


														"""simple docstring"""


														_a     					=    super().get_config()
														config.update({'''weight_decay_rate''': self.weight_decay_rate} )
														return config







							def        __lowerCAmelCase      (  self	:  Any							,      lowerCAmelCase_	:  int )      ->							Any:


														"""simple docstring"""


														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(_snake_case							,      _snake_case ) is not None:
																																			return True

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



class       A   (		__SCREAMING_SNAKE_CASE   ):







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


														"""simple docstring"""


														_a     					=    []
														_a     					=    None







							@property
							def        __lowerCAmelCase      (  self	:  Union[str, Any] )      ->							Optional[Any]:


														"""simple docstring"""


														if self._accum_steps is None:
																					_a     					=    tf.Variable(
																					    tf.constant(0							,      dtype=tf.intaa )							,      trainable=_snake_case							,      synchronization=tf.VariableSynchronization.ON_READ							,      aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA							,      )

														return self._accum_steps.value()







							@property
							def        __lowerCAmelCase      (  self	:  int )      ->							Union[str, Any]:


														"""simple docstring"""


														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	:  Union[str, Any]							,      lowerCAmelCase_	:  str )      ->							Optional[int]:


														"""simple docstring"""


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

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

														self._accum_steps.assign_add(1 )







							def        __lowerCAmelCase      (  self	:  Dict )      ->							List[Any]:


														"""simple docstring"""


														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(_snake_case ) )







'''simple docstring'''


from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


_snake_case				:				Dict       							=		logging.get_logger(__name__)

_snake_case				:				Optional[Any]       							=		{
    'shi-labs/nat-mini-in1k-224': 'https://huggingface.co/shi-labs/nat-mini-in1k-224/resolve/main/config.json',
    # See all Nat models at https://huggingface.co/models?filter=nat
}

class       A   (		_a					,_a   ):
							lowercase_          =    'nat'

							lowercase_          =    {
							    'num_attention_heads': 'num_heads',
							    'num_hidden_layers': 'num_layers',
							}







							def __init__(  self	:  List[str]							,      lowerCAmelCase_	:  str=4							,      lowerCAmelCase_	:  Any=3							,      lowerCAmelCase_	:  List[Any]=64							,      lowerCAmelCase_	:  Dict=[3, 4, 6, 5]							,      lowerCAmelCase_	:  Dict=[2, 4, 8, 16]							,      lowerCAmelCase_	:  str=7							,      lowerCAmelCase_	:  Dict=3.0							,      lowerCAmelCase_	:  int=True							,      lowerCAmelCase_	:  Tuple=0.0							,      lowerCAmelCase_	:  str=0.0							,      lowerCAmelCase_	:  List[str]=0.1							,      lowerCAmelCase_	:  int="gelu"							,      lowerCAmelCase_	:  List[str]=0.0_2							,      lowerCAmelCase_	:  str=1e-5							,      lowerCAmelCase_	:  Tuple=0.0							,      lowerCAmelCase_	:  str=None							,      lowerCAmelCase_	:  int=None							,      **lowerCAmelCase_	:  List[Any]							,      )      ->							Any:


														"""simple docstring"""


														super().__init__(**lowerCAmelCase_ )

														_a     					=    patch_size
														_a     					=    num_channels
														_a     					=    embed_dim
														_a     					=    depths
														_a     					=    len(lowerCAmelCase_ )
														_a     					=    num_heads
														_a     					=    kernel_size
														_a     					=    mlp_ratio
														_a     					=    qkv_bias
														_a     					=    hidden_dropout_prob
														_a     					=    attention_probs_dropout_prob
														_a     					=    drop_path_rate
														_a     					=    hidden_act
														_a     					=    layer_norm_eps
														_a     					=    initializer_range
														# we set the hidden_size attribute in order to make Nat work with VisionEncoderDecoderModel
														# this indicates the channel dimension after the last stage of the model
														_a     					=    int(embed_dim * 2 ** (len(lowerCAmelCase_ ) - 1) )
														_a     					=    layer_scale_init_value
														_a     					=    ['''stem'''] + [F'stage{idx}' for idx in range(1							,      len(lowerCAmelCase_ ) + 1 )]
														_a					,							_a     					=    get_aligned_output_features_output_indices(
														    out_features=lowerCAmelCase_							,      out_indices=lowerCAmelCase_							,      stage_names=self.stage_names )

