





"""simple docstring"""






import tensorflow as tf

from ...tf_utils import shape_list



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






       def __init__(   self    ,  a__    ,  a__    ,  a__    ,  a__    ,  a__=1    ,  a__=False    ,  **a__     ):
              super().__init__(**a__     )

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






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






       @staticmethod
       def 							__A (   a__    ,  a__     ):
              _UpperCAmelCase     =     shape_list(a__     )
              _UpperCAmelCase     =     tf.range(lp_size[0]    ,  dtype=target.dtype     )
              _UpperCAmelCase     =     tf.stack([r, target]    ,  1     )
              return tf.gather_nd(a__    ,  a__     )






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

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

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

              return out













"""simple docstring"""






from datetime import datetime

import matplotlib.pyplot as plt
import torch

def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE							)		->       Any:






       """simple docstring"""




       for param in module.parameters():
              _UpperCAmelCase     =     False

def       __lowerCamelCase    (       )		->       Dict:






       """simple docstring"""




       _UpperCAmelCase     =     'cuda' if torch.cuda.is_available() else 'cpu'
       if torch.backends.mps.is_available() and torch.backends.mps.is_built():
              _UpperCAmelCase     =     'mps'
       if device == "mps":
              print(
                  'WARNING: MPS currently doesn\'t seem to work, and messes up backpropagation without any visible torch'
                  ' errors. I recommend using CUDA on a colab notebook or CPU instead if you\'re facing inexplicable issues'
                  ' with generations.'							)
       return device

def       __lowerCamelCase    (       SCREAMING_SNAKE_CASE							)		->       str:






       """simple docstring"""




       _UpperCAmelCase     =     plt.imshow(SCREAMING_SNAKE_CASE							)
       fig.axes.get_xaxis().set_visible(SCREAMING_SNAKE_CASE							)
       fig.axes.get_yaxis().set_visible(SCREAMING_SNAKE_CASE							)
       plt.show()


def       __lowerCamelCase    (       )		->       Dict:






       """simple docstring"""




       _UpperCAmelCase     =     datetime.now()
       _UpperCAmelCase     =     current_time.strftime('%H:%M:%S'							)
       return timestamp






