





'''simple docstring'''

import tensorflow as tf

from ...tf_utils import shape_list
class     __lowerCAmelCase					(   tf.keras.layers.Layer			):






       """simple docstring"""





       def __init__(							self       :		str	,			lowerCAmelCase__       :		Union[str, Any]	,			lowerCAmelCase__       :		str	,			lowerCAmelCase__       :		Any	,			lowerCAmelCase__       :		Dict	,			lowerCAmelCase__       :		List[str]=1	,			lowerCAmelCase__       :		Optional[int]=False	,			**lowerCAmelCase__       :		List[Any]					)   ->					Any:

            '''simple docstring'''




            super().__init__(**_snake_case					)

            _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 snake_case__						(							self       :		Optional[Any]	,			lowerCAmelCase__       :		str					)   ->					Optional[int]:

            '''simple docstring'''




            if self.n_clusters > 0:
                 _UpperCamelCase       =	self.add_weight(
                     shape=(self.n_clusters, self.d_embed)	,			initializer='''zeros'''	,			trainable=_snake_case	,			name='''cluster_weight'''					)
                 _UpperCamelCase       =	self.add_weight(
                     shape=(self.n_clusters,)	,			initializer='''zeros'''	,			trainable=_snake_case	,			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=_snake_case	,			name=f"""out_projs_._{i}"""	,			)
                           self.out_projs.append(_snake_case					)
                      else:
                           self.out_projs.append(_snake_case					)
                      _UpperCamelCase       =	self.add_weight(
                          shape=(self.vocab_size, self.d_embed)	,			initializer='''zeros'''	,			trainable=_snake_case	,			name=f"""out_layers_._{i}_._weight"""	,			)
                      _UpperCamelCase       =	self.add_weight(
                          shape=(self.vocab_size,)	,			initializer='''zeros'''	,			trainable=_snake_case	,			name=f"""out_layers_._{i}_._bias"""	,			)
                      self.out_layers.append((weight, bias)					)
            else:
                 for i in range(len(self.cutoffs					)					):
                      _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=_snake_case	,			name=f"""out_projs_._{i}"""					)
                      self.out_projs.append(_snake_case					)
                      _UpperCamelCase       =	self.add_weight(
                          shape=(r_idx - l_idx, d_emb_i)	,			initializer='''zeros'''	,			trainable=_snake_case	,			name=f"""out_layers_._{i}_._weight"""	,			)
                      _UpperCamelCase       =	self.add_weight(
                          shape=(r_idx - l_idx,)	,			initializer='''zeros'''	,			trainable=_snake_case	,			name=f"""out_layers_._{i}_._bias"""	,			)
                      self.out_layers.append((weight, bias)					)
            super().build(_snake_case					)





       @staticmethod
       def snake_case__						(							lowerCAmelCase__       :		Optional[Any]	,			lowerCAmelCase__       :		List[Any]	,			lowerCAmelCase__       :		List[Any]	,			lowerCAmelCase__       :		Tuple=None					)   ->					Tuple:

            '''simple docstring'''




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





       @staticmethod
       def snake_case__						(							lowerCAmelCase__       :		Tuple	,			lowerCAmelCase__       :		Optional[Any]					)   ->					Any:

            '''simple docstring'''




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







       def snake_case__						(							self       :		Optional[Any]	,			lowerCAmelCase__       :		Optional[int]	,			lowerCAmelCase__       :		str	,			lowerCAmelCase__       :		Union[str, Any]=True	,			lowerCAmelCase__       :		Dict=False					)   ->					Union[str, Any]:

            '''simple docstring'''




            _UpperCamelCase       =	0
            if self.n_clusters == 0:
                 _UpperCamelCase       =	self._logit(_snake_case	,			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=_snake_case	,			logits=_snake_case					)
                 _UpperCamelCase       =	tf.nn.log_softmax(_snake_case	,			axis=-1					)
            else:
                 _UpperCamelCase       =	shape_list(_snake_case					)
                 _UpperCamelCase       =	[]
                 _UpperCamelCase       =	tf.zeros(hidden_sizes[:2]					)
                 for i in range(len(self.cutoffs					)					):
                      _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(_snake_case					)
                           _UpperCamelCase       =	tf.boolean_mask(_snake_case	,			_snake_case					) - 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(_snake_case	,			_snake_case	,			_snake_case	,			self.out_projs[0]					)
                           _UpperCamelCase       =	tf.nn.log_softmax(_snake_case					)
                           out.append(head_logprob[..., : self.cutoffs[0]]					)
                           if target is not None:
                                _UpperCamelCase       =	tf.boolean_mask(_snake_case	,			_snake_case					)
                                _UpperCamelCase       =	self._gather_logprob(_snake_case	,			_snake_case					)
                      else:
                           _UpperCamelCase       =	self._logit(_snake_case	,			_snake_case	,			_snake_case	,			self.out_projs[i]					)
                           _UpperCamelCase       =	tf.nn.log_softmax(_snake_case					)
                           _UpperCamelCase       =	self.cutoffs[0] + i - 1  # No probability for the head cluster
                           _UpperCamelCase       =	head_logprob[..., cluster_prob_idx, None] + tail_logprob
                           out.append(_snake_case					)
                           if target is not None:
                                _UpperCamelCase       =	tf.boolean_mask(_snake_case	,			_snake_case					)
                                _UpperCamelCase       =	tf.boolean_mask(_snake_case	,			_snake_case					)
                                _UpperCamelCase       =	self._gather_logprob(_snake_case	,			_snake_case					)
                                cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
                      if target is not None:
                           loss += tf.scatter_nd(_snake_case	,			-cur_logprob	,			shape_list(_snake_case					)					)
                 _UpperCamelCase       =	tf.concat(_snake_case	,			axis=-1					)

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

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

            return out










'''simple docstring'''

from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowercase__						:      Optional[Any]	=    {
    'configuration_autoformer': [
        'AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'AutoformerConfig',
    ],
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowercase__						:      List[Any]	=    [
       'AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
       'AutoformerForPrediction',
       'AutoformerModel',
       'AutoformerPreTrainedModel',
   ]


if TYPE_CHECKING:
   from .configuration_autoformer import (
       AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
       AutoformerConfig,
   )

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_autoformer import (
          AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
          AutoformerForPrediction,
          AutoformerModel,
          AutoformerPreTrainedModel,
      )

else:
   import sys

   lowercase__						:      Tuple	=    _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



