
import tensorflow as tf

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

 """simple docstring"""

 def __init__(						self			:			str					, SCREAMING_SNAKE_CASE__			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			Optional[int]					, SCREAMING_SNAKE_CASE__			:			Tuple					, SCREAMING_SNAKE_CASE__			:			str					, SCREAMING_SNAKE_CASE__			:			Dict=1					, SCREAMING_SNAKE_CASE__			:			str=False					, **SCREAMING_SNAKE_CASE__			:			int   ):






      """simple docstring"""



      super().__init__(**SCREAMING_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      __lowerCAmelCase  (						self			:			Tuple					, SCREAMING_SNAKE_CASE__			:			Tuple   ):






      """simple docstring"""



      if self.n_clusters > 0:
           UpperCamelCase						    =					self.add_weight(
               shape=(self.n_clusters, self.d_embed)					, initializer='zeros'					, trainable=SCREAMING_SNAKE_CASE__					, name='cluster_weight'   )
           UpperCamelCase						    =					self.add_weight(
               shape=(self.n_clusters,)					, initializer='zeros'					, trainable=SCREAMING_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=SCREAMING_SNAKE_CASE__					, name=F'out_projs_._{i}'					, )
                     self.out_projs.append(SCREAMING_SNAKE_CASE__   )
                else:
                     self.out_projs.append(SCREAMING_SNAKE_CASE__   )
                UpperCamelCase						    =					self.add_weight(
                    shape=(self.vocab_size, self.d_embed)					, initializer='zeros'					, trainable=SCREAMING_SNAKE_CASE__					, name=F'out_layers_._{i}_._weight'					, )
                UpperCamelCase						    =					self.add_weight(
                    shape=(self.vocab_size,)					, initializer='zeros'					, trainable=SCREAMING_SNAKE_CASE__					, 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=SCREAMING_SNAKE_CASE__					, name=F'out_projs_._{i}'   )
                self.out_projs.append(SCREAMING_SNAKE_CASE__   )
                UpperCamelCase						    =					self.add_weight(
                    shape=(r_idx - l_idx, d_emb_i)					, initializer='zeros'					, trainable=SCREAMING_SNAKE_CASE__					, name=F'out_layers_._{i}_._weight'					, )
                UpperCamelCase						    =					self.add_weight(
                    shape=(r_idx - l_idx,)					, initializer='zeros'					, trainable=SCREAMING_SNAKE_CASE__					, name=F'out_layers_._{i}_._bias'					, )
                self.out_layers.append((weight, bias)   )
      super().build(SCREAMING_SNAKE_CASE__   )

 @staticmethod
 def      __lowerCAmelCase  (						SCREAMING_SNAKE_CASE__			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			Tuple					, SCREAMING_SNAKE_CASE__			:			int					, SCREAMING_SNAKE_CASE__			:			Optional[Any]=None   ):






      """simple docstring"""



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

 @staticmethod
 def      __lowerCAmelCase  (						SCREAMING_SNAKE_CASE__			:			List[Any]					, SCREAMING_SNAKE_CASE__			:			Optional[Any]   ):






      """simple docstring"""



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






 def      __lowerCAmelCase  (						self			:			Dict					, SCREAMING_SNAKE_CASE__			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]					, SCREAMING_SNAKE_CASE__			:			int=True					, SCREAMING_SNAKE_CASE__			:			str=False   ):






      """simple docstring"""



      UpperCamelCase						    =					0
      if self.n_clusters == 0:
           UpperCamelCase						    =					self._logit(SCREAMING_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=SCREAMING_SNAKE_CASE__					, logits=SCREAMING_SNAKE_CASE__   )
           UpperCamelCase						    =					tf.nn.log_softmax(SCREAMING_SNAKE_CASE__					, axis=-1   )
      else:
           UpperCamelCase						    =					shape_list(SCREAMING_SNAKE_CASE__   )
           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(SCREAMING_SNAKE_CASE__   )
                     UpperCamelCase						    =					tf.boolean_mask(SCREAMING_SNAKE_CASE__					, SCREAMING_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(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, self.out_projs[0]   )
                     UpperCamelCase						    =					tf.nn.log_softmax(SCREAMING_SNAKE_CASE__   )
                     out.append(head_logprob[..., : self.cutoffs[0]]   )
                     if target is not None:
                          UpperCamelCase						    =					tf.boolean_mask(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )
                          UpperCamelCase						    =					self._gather_logprob(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )
                else:
                     UpperCamelCase						    =					self._logit(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, self.out_projs[i]   )
                     UpperCamelCase						    =					tf.nn.log_softmax(SCREAMING_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(SCREAMING_SNAKE_CASE__   )
                     if target is not None:
                          UpperCamelCase						    =					tf.boolean_mask(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )
                          UpperCamelCase						    =					tf.boolean_mask(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )
                          UpperCamelCase						    =					self._gather_logprob(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )
                          cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
                if target is not None:
                     loss += tf.scatter_nd(SCREAMING_SNAKE_CASE__					, -cur_logprob					, shape_list(SCREAMING_SNAKE_CASE__   )   )
           UpperCamelCase						    =					tf.concat(SCREAMING_SNAKE_CASE__					, axis=-1   )

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

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

      return out




# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

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


_snake_case      = {
    '''configuration_efficientnet''': [
        '''EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''EfficientNetConfig''',
        '''EfficientNetOnnxConfig''',
    ]
}

try:
 if not is_vision_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _snake_case      = ['''EfficientNetImageProcessor''']

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _snake_case      = [
     '''EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''EfficientNetForImageClassification''',
     '''EfficientNetModel''',
     '''EfficientNetPreTrainedModel''',
 ]

if TYPE_CHECKING:
 from .configuration_efficientnet import (
     EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP,
     EfficientNetConfig,
     EfficientNetOnnxConfig,
 )

 try:
  if not is_vision_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .image_processing_efficientnet import EfficientNetImageProcessor

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_efficientnet import (
      EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST,
      EfficientNetForImageClassification,
      EfficientNetModel,
      EfficientNetPreTrainedModel,
  )

else:
 import sys

 _snake_case      = _LazyModule(__name__, globals()['''__file__'''], _import_structure)


