




'''simple docstring'''


import tensorflow as tf

from ...tf_utils import shape_list





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







   def __init__(						self   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase=1   ,				UpperCamelCase=False   ,				**UpperCamelCase   )					->     List[str]:
     super().__init__(**UpperCamelCase   )

     __a								=       vocab_size
     __a								=       d_embed
     __a								=       d_proj

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

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

     __a								=       []
     __a								=       []







   def 		UpperCamelCase__      (						self   ,				UpperCamelCase   )					->     Dict:
     if self.n_clusters > 0:
       __a								=       self.add_weight(
           shape=(self.n_clusters, self.d_embed)   ,				initializer='zeros'   ,				trainable=UpperCamelCase   ,				name='cluster_weight'   )
       __a								=       self.add_weight(
           shape=(self.n_clusters,)   ,				initializer='zeros'   ,				trainable=UpperCamelCase   ,				name='cluster_bias'   )

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

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







   @staticmethod
   def 		UpperCamelCase__      (						UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase   ,				UpperCamelCase=None   )					->     Dict:
     __a								=       x
     if proj is not None:
       __a								=       tf.einsum('ibd,ed->ibe'   ,				UpperCamelCase   ,				UpperCamelCase   )
     return tf.einsum('ibd,nd->ibn'   ,				UpperCamelCase   ,				UpperCamelCase   ) + b







   @staticmethod
   def 		UpperCamelCase__      (						UpperCamelCase   ,				UpperCamelCase   )					->     List[str]:
     __a								=       shape_list(UpperCamelCase   )
     __a								=       tf.range(lp_size[0]   ,				dtype=target.dtype   )
     __a								=       tf.stack([r, target]   ,				1   )
     return tf.gather_nd(UpperCamelCase   ,				UpperCamelCase   )







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

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

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

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

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

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

     return out









'''simple docstring'''


from typing import TYPE_CHECKING

from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCAmelCase_					=   {
    "configuration_trajectory_transformer": [
        "TRAJECTORY_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "TrajectoryTransformerConfig",
    ],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_					=   [
      "TRAJECTORY_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
      "TrajectoryTransformerModel",
      "TrajectoryTransformerPreTrainedModel",
      "load_tf_weights_in_trajectory_transformer",
  ]


if TYPE_CHECKING:
  from .configuration_trajectory_transformer import (
      TRAJECTORY_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
      TrajectoryTransformerConfig,
  )

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_trajectory_transformer import (
        TRAJECTORY_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TrajectoryTransformerModel,
        TrajectoryTransformerPreTrainedModel,
        load_tf_weights_in_trajectory_transformer,
    )


else:
  import sys

  UpperCAmelCase_					=   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



