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






   '''simple docstring'''







   def __init__(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ = 1.0 ,   lowerCamelCase__ = None ,   ):

      '''simple docstring'''







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







   def __call__(			self ,   lowerCamelCase__   ):

      '''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`.
         UpperCamelCase       =       tf.cast(lowerCamelCase__ ,   tf.floataa   )
         UpperCamelCase       =       tf.cast(self.warmup_steps ,   tf.floataa   )
         UpperCamelCase       =       global_step_float / warmup_steps_float
         UpperCamelCase       =       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   ):

      '''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.9_9_9,      _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,      ):
   UpperCamelCase       =       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:
      UpperCamelCase       =       WarmUp(
          initial_learning_rate=_UpperCAmelCase,      decay_schedule_fn=_UpperCAmelCase,      warmup_steps=_UpperCAmelCase,      )
   if weight_decay_rate > 0.0:
      UpperCamelCase       =       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:
      UpperCamelCase       =       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 						lowercase__	(  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__ ,   ):

      '''simple docstring'''







      super().__init__(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   **lowerCamelCase__   )
      UpperCamelCase       =       weight_decay_rate
      UpperCamelCase       =       include_in_weight_decay
      UpperCamelCase       =       exclude_from_weight_decay







   @classmethod
   def       UpperCAmelCase				(			cls ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =       {'''WarmUp''': WarmUp}
      return super(lowerCamelCase__ ,   cls   ).from_config(lowerCamelCase__ ,   custom_objects=lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   ):

      '''simple docstring'''







      super(lowerCamelCase__ ,   self   )._prepare_local(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   )
      UpperCamelCase       =       tf.constant(
          self.weight_decay_rate ,   name='''adam_weight_decay_rate'''   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =       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 ,   lowerCamelCase__ ,   lowerCamelCase__=None ,   **lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase     , UpperCamelCase       =       list(zip(*lowerCamelCase__   )   )
      return super(lowerCamelCase__ ,   self   ).apply_gradients(zip(lowerCamelCase__ ,   lowerCamelCase__   ) ,   name=lowerCamelCase__ ,   **lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   ):

      '''simple docstring'''







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

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

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







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=None   ):

      '''simple docstring'''







      UpperCamelCase     , UpperCamelCase       =       self._get_lr(var.device ,   var.dtype.base_dtype ,   lowerCamelCase__   )
      UpperCamelCase       =       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 ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=None   ):

      '''simple docstring'''







      UpperCamelCase     , UpperCamelCase       =       self._get_lr(var.device ,   var.dtype.base_dtype ,   lowerCamelCase__   )
      UpperCamelCase       =       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   ):

      '''simple docstring'''







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







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''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(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 						lowercase__	(  snake_case_   ):






   '''simple docstring'''







   def __init__(			self   ):

      '''simple docstring'''







      UpperCamelCase       =       []
      UpperCamelCase       =       None







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      if self._accum_steps is None:
         UpperCamelCase       =       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   ):

      '''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 ,   lowerCamelCase__   ):

      '''simple docstring'''







      if not self._gradients:
         UpperCamelCase       =       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   ):

      '''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(lowerCamelCase__   )   )







'''simple docstring'''


import tensorflow as tf

from ...tf_utils import shape_list


class 						lowercase__	(  tf.keras.layers.Layer   ):






   '''simple docstring'''







   def __init__(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=1 ,   lowerCamelCase__=False ,   **lowerCamelCase__   ):

      '''simple docstring'''







      super().__init__(**lowerCamelCase__   )

      UpperCamelCase       =  vocab_size
      UpperCamelCase       =  d_embed
      UpperCamelCase       =  d_proj

      UpperCamelCase       =  cutoffs + [vocab_size]
      UpperCamelCase       =  [0] + self.cutoffs
      UpperCamelCase       =  div_val

      UpperCamelCase       =  self.cutoffs[0]
      UpperCamelCase       =  len(self.cutoffs   ) - 1
      UpperCamelCase       =  self.shortlist_size + self.n_clusters
      UpperCamelCase       =  keep_order

      UpperCamelCase       =  []
      UpperCamelCase       =  []







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      if self.n_clusters > 0:
         UpperCamelCase       =  self.add_weight(
             shape=(self.n_clusters, self.d_embed) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name='''cluster_weight'''   )
         UpperCamelCase       =  self.add_weight(
             shape=(self.n_clusters,) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name='''cluster_bias'''   )

      if self.div_val == 1:
         for i in range(len(self.cutoffs   )   ):
            if self.d_proj != self.d_embed:
               UpperCamelCase       =  self.add_weight(
                   shape=(self.d_embed, self.d_proj) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_projs_._{i}' ,   )
               self.out_projs.append(lowerCamelCase__   )
            else:
               self.out_projs.append(lowerCamelCase__   )
            UpperCamelCase       =  self.add_weight(
                shape=(self.vocab_size, self.d_embed) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_layers_._{i}_._weight' ,   )
            UpperCamelCase       =  self.add_weight(
                shape=(self.vocab_size,) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_layers_._{i}_._bias' ,   )
            self.out_layers.append((weight, bias)   )
      else:
         for i in range(len(self.cutoffs   )   ):
            UpperCamelCase     , UpperCamelCase       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
            UpperCamelCase       =  self.d_embed // (self.div_val**i)

            UpperCamelCase       =  self.add_weight(
                shape=(d_emb_i, self.d_proj) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_projs_._{i}'   )
            self.out_projs.append(lowerCamelCase__   )
            UpperCamelCase       =  self.add_weight(
                shape=(r_idx - l_idx, d_emb_i) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_layers_._{i}_._weight' ,   )
            UpperCamelCase       =  self.add_weight(
                shape=(r_idx - l_idx,) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_layers_._{i}_._bias' ,   )
            self.out_layers.append((weight, bias)   )
      super().build(lowerCamelCase__   )







   @staticmethod
   def       UpperCAmelCase				(			lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=None   ):

      '''simple docstring'''







      UpperCamelCase       =  x
      if proj is not None:
         UpperCamelCase       =  tf.einsum('''ibd,ed->ibe''' ,   lowerCamelCase__ ,   lowerCamelCase__   )
      return tf.einsum('''ibd,nd->ibn''' ,   lowerCamelCase__ ,   lowerCamelCase__   ) + b







   @staticmethod
   def       UpperCAmelCase				(			lowerCamelCase__ ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  shape_list(lowerCamelCase__   )
      UpperCamelCase       =  tf.range(lp_size[0] ,   dtype=target.dtype   )
      UpperCamelCase       =  tf.stack([r, target] ,   1   )
      return tf.gather_nd(lowerCamelCase__ ,   lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=True ,   lowerCamelCase__=False   ):

      '''simple docstring'''







      UpperCamelCase       =  0
      if self.n_clusters == 0:
         UpperCamelCase       =  self._logit(lowerCamelCase__ ,   self.out_layers[0][0] ,   self.out_layers[0][1] ,   self.out_projs[0]   )
         if target is not None:
            UpperCamelCase       =  tf.nn.sparse_softmax_cross_entropy_with_logits(labels=lowerCamelCase__ ,   logits=lowerCamelCase__   )
         UpperCamelCase       =  tf.nn.log_softmax(lowerCamelCase__ ,   axis=-1   )
      else:
         UpperCamelCase       =  shape_list(lowerCamelCase__   )
         UpperCamelCase       =  []
         UpperCamelCase       =  tf.zeros(hidden_sizes[:2]   )
         for i in range(len(self.cutoffs   )   ):
            UpperCamelCase     , UpperCamelCase       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
            if target is not None:
               UpperCamelCase       =  (target >= l_idx) & (target < r_idx)
               UpperCamelCase       =  tf.where(lowerCamelCase__   )
               UpperCamelCase       =  tf.boolean_mask(lowerCamelCase__ ,   lowerCamelCase__   ) - l_idx

            if self.div_val == 1:
               UpperCamelCase       =  self.out_layers[0][0][l_idx:r_idx]
               UpperCamelCase       =  self.out_layers[0][1][l_idx:r_idx]
            else:
               UpperCamelCase       =  self.out_layers[i][0]
               UpperCamelCase       =  self.out_layers[i][1]

            if i == 0:
               UpperCamelCase       =  tf.concat([cur_W, self.cluster_weight] ,   0   )
               UpperCamelCase       =  tf.concat([cur_b, self.cluster_bias] ,   0   )

               UpperCamelCase       =  self._logit(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   self.out_projs[0]   )
               UpperCamelCase       =  tf.nn.log_softmax(lowerCamelCase__   )
               out.append(head_logprob[..., : self.cutoffs[0]]   )
               if target is not None:
                  UpperCamelCase       =  tf.boolean_mask(lowerCamelCase__ ,   lowerCamelCase__   )
                  UpperCamelCase       =  self._gather_logprob(lowerCamelCase__ ,   lowerCamelCase__   )
            else:
               UpperCamelCase       =  self._logit(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   self.out_projs[i]   )
               UpperCamelCase       =  tf.nn.log_softmax(lowerCamelCase__   )
               UpperCamelCase       =  self.cutoffs[0] + i - 1  # No probability for the head cluster
               UpperCamelCase       =  head_logprob[..., cluster_prob_idx, None] + tail_logprob
               out.append(lowerCamelCase__   )
               if target is not None:
                  UpperCamelCase       =  tf.boolean_mask(lowerCamelCase__ ,   lowerCamelCase__   )
                  UpperCamelCase       =  tf.boolean_mask(lowerCamelCase__ ,   lowerCamelCase__   )
                  UpperCamelCase       =  self._gather_logprob(lowerCamelCase__ ,   lowerCamelCase__   )
                  cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
            if target is not None:
               loss += tf.scatter_nd(lowerCamelCase__ ,   -cur_logprob ,   shape_list(lowerCamelCase__   )   )
         UpperCamelCase       =  tf.concat(lowerCamelCase__ ,   axis=-1   )

      if target is not None:
         if return_mean:
            UpperCamelCase       =  tf.reduce_mean(lowerCamelCase__   )
         # Add the training-time loss value to the layer using `self.add_loss()`.
         self.add_loss(lowerCamelCase__   )

         # Log the loss as a metric (we could log arbitrary metrics,
         # including different metrics for training and inference.
         self.add_metric(lowerCamelCase__ ,   name=self.name ,   aggregation='''mean''' if return_mean else ''''''   )

      return out





