'''simple docstring'''

import tensorflow as tf

from ...tf_utils import shape_list


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





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

     lowerCAmelCase_			=      vocab_size
     lowerCAmelCase_			=      d_embed
     lowerCAmelCase_			=      d_proj

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

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

     lowerCAmelCase_			=      []
     lowerCAmelCase_			=      []





  def 						UpperCAmelCase_	(	self    ,     _lowerCamelCase ):
     if self.n_clusters > 0:
        lowerCAmelCase_			=      self.add_weight(
            shape=(self.n_clusters, self.d_embed)    ,     initializer='''zeros'''    ,     trainable=_lowerCamelCase    ,     name='''cluster_weight''' )
        lowerCAmelCase_			=      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:
              lowerCAmelCase_			=      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 )
           lowerCAmelCase_			=      self.add_weight(
               shape=(self.vocab_size, self.d_embed)    ,     initializer='''zeros'''    ,     trainable=_lowerCamelCase    ,     name=F'''out_layers_._{i}_._weight'''    ,     )
           lowerCAmelCase_			=      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 ) ):
           lowerCAmelCase_	,lowerCAmelCase_			=      self.cutoff_ends[i], self.cutoff_ends[i + 1]
           lowerCAmelCase_			=      self.d_embed // (self.div_val**i)

           lowerCAmelCase_			=      self.add_weight(
               shape=(d_emb_i, self.d_proj)    ,     initializer='''zeros'''    ,     trainable=_lowerCamelCase    ,     name=F'''out_projs_._{i}''' )
           self.out_projs.append(_lowerCamelCase )
           lowerCAmelCase_			=      self.add_weight(
               shape=(r_idx - l_idx, d_emb_i)    ,     initializer='''zeros'''    ,     trainable=_lowerCamelCase    ,     name=F'''out_layers_._{i}_._weight'''    ,     )
           lowerCAmelCase_			=      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 ):
     lowerCAmelCase_			=      x
     if proj is not None:
        lowerCAmelCase_			=      tf.einsum('''ibd,ed->ibe'''    ,     _lowerCamelCase    ,     _lowerCamelCase )
     return tf.einsum('''ibd,nd->ibn'''    ,     _lowerCamelCase    ,     _lowerCamelCase ) + b





  @staticmethod
  def 						UpperCAmelCase_	(	_lowerCamelCase    ,     _lowerCamelCase ):
     lowerCAmelCase_			=      shape_list(_lowerCamelCase )
     lowerCAmelCase_			=      tf.range(lp_size[0]    ,     dtype=target.dtype )
     lowerCAmelCase_			=      tf.stack([r, target]    ,     1 )
     return tf.gather_nd(_lowerCamelCase    ,     _lowerCamelCase )







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

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

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

              lowerCAmelCase_			=      self._logit(_lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     self.out_projs[0] )
              lowerCAmelCase_			=      tf.nn.log_softmax(_lowerCamelCase )
              out.append(head_logprob[..., : self.cutoffs[0]] )
              if target is not None:
                 lowerCAmelCase_			=      tf.boolean_mask(_lowerCamelCase    ,     _lowerCamelCase )
                 lowerCAmelCase_			=      self._gather_logprob(_lowerCamelCase    ,     _lowerCamelCase )
           else:
              lowerCAmelCase_			=      self._logit(_lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     self.out_projs[i] )
              lowerCAmelCase_			=      tf.nn.log_softmax(_lowerCamelCase )
              lowerCAmelCase_			=      self.cutoffs[0] + i - 1  # No probability for the head cluster
              lowerCAmelCase_			=      head_logprob[..., cluster_prob_idx, None] + tail_logprob
              out.append(_lowerCamelCase )
              if target is not None:
                 lowerCAmelCase_			=      tf.boolean_mask(_lowerCamelCase    ,     _lowerCamelCase )
                 lowerCAmelCase_			=      tf.boolean_mask(_lowerCamelCase    ,     _lowerCamelCase )
                 lowerCAmelCase_			=      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 ) )
        lowerCAmelCase_			=      tf.concat(_lowerCamelCase    ,     axis=-1 )

     if target is not None:
        if return_mean:
           lowerCAmelCase_			=      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






'''simple docstring'''

import os
import tempfile
import unittest
from pathlib import Path

from transformers import AutoConfig, is_torch_available
from transformers.testing_utils import require_torch, torch_device


if is_torch_available():
 from transformers import PyTorchBenchmark, PyTorchBenchmarkArguments


@require_torch
class     __UpperCAmelCase       ( unittest.TestCase		):





  def 						UpperCAmelCase_	(	self    ,     _lowerCamelCase ):
     for model_result in results.values():
        for batch_size, sequence_length in zip(model_result['''bs''']    ,     model_result['''ss'''] ):
           lowerCAmelCase_			=      model_result['''result'''][batch_size][sequence_length]
           self.assertIsNotNone(_lowerCamelCase )





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      '''sshleifer/tiny-gpt2'''
     lowerCAmelCase_			=      PyTorchBenchmarkArguments(
         models=[MODEL_ID]    ,     training=_lowerCamelCase    ,     inference=_lowerCamelCase    ,     sequence_lengths=[8]    ,     batch_sizes=[1]    ,     multi_process=_lowerCamelCase    ,     )
     lowerCAmelCase_			=      PyTorchBenchmark(_lowerCamelCase )
     lowerCAmelCase_			=      benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result )
     self.check_results_dict_not_empty(results.memory_inference_result )





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      '''sgugger/tiny-distilbert-classification'''
     lowerCAmelCase_			=      PyTorchBenchmarkArguments(
         models=[MODEL_ID]    ,     training=_lowerCamelCase    ,     inference=_lowerCamelCase    ,     sequence_lengths=[8]    ,     batch_sizes=[1]    ,     multi_process=_lowerCamelCase    ,     only_pretrain_model=_lowerCamelCase    ,     )
     lowerCAmelCase_			=      PyTorchBenchmark(_lowerCamelCase )
     lowerCAmelCase_			=      benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result )
     self.check_results_dict_not_empty(results.memory_inference_result )





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      '''sshleifer/tiny-gpt2'''
     lowerCAmelCase_			=      PyTorchBenchmarkArguments(
         models=[MODEL_ID]    ,     training=_lowerCamelCase    ,     inference=_lowerCamelCase    ,     torchscript=_lowerCamelCase    ,     sequence_lengths=[8]    ,     batch_sizes=[1]    ,     multi_process=_lowerCamelCase    ,     )
     lowerCAmelCase_			=      PyTorchBenchmark(_lowerCamelCase )
     lowerCAmelCase_			=      benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result )
     self.check_results_dict_not_empty(results.memory_inference_result )





  @unittest.skipIf(torch_device == '''cpu'''    ,     '''Cant do half precision''' )
  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      '''sshleifer/tiny-gpt2'''
     lowerCAmelCase_			=      PyTorchBenchmarkArguments(
         models=[MODEL_ID]    ,     training=_lowerCamelCase    ,     inference=_lowerCamelCase    ,     fpaa=_lowerCamelCase    ,     sequence_lengths=[8]    ,     batch_sizes=[1]    ,     multi_process=_lowerCamelCase    ,     )
     lowerCAmelCase_			=      PyTorchBenchmark(_lowerCamelCase )
     lowerCAmelCase_			=      benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result )
     self.check_results_dict_not_empty(results.memory_inference_result )





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      '''sshleifer/tiny-gpt2'''
     lowerCAmelCase_			=      AutoConfig.from_pretrained(_lowerCamelCase )
     # set architectures equal to `None`
     lowerCAmelCase_			=      None
     lowerCAmelCase_			=      PyTorchBenchmarkArguments(
         models=[MODEL_ID]    ,     training=_lowerCamelCase    ,     inference=_lowerCamelCase    ,     sequence_lengths=[8]    ,     batch_sizes=[1]    ,     multi_process=_lowerCamelCase    ,     )
     lowerCAmelCase_			=      PyTorchBenchmark(_lowerCamelCase    ,     configs=[config] )
     lowerCAmelCase_			=      benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result )
     self.check_results_dict_not_empty(results.memory_inference_result )





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      '''sshleifer/tiny-gpt2'''
     lowerCAmelCase_			=      PyTorchBenchmarkArguments(
         models=[MODEL_ID]    ,     training=_lowerCamelCase    ,     inference=_lowerCamelCase    ,     sequence_lengths=[8]    ,     batch_sizes=[1]    ,     multi_process=_lowerCamelCase    ,     )
     lowerCAmelCase_			=      PyTorchBenchmark(_lowerCamelCase )
     lowerCAmelCase_			=      benchmark.run()
     self.check_results_dict_not_empty(results.time_train_result )
     self.check_results_dict_not_empty(results.memory_train_result )





  @unittest.skipIf(torch_device == '''cpu'''    ,     '''Can\'t do half precision''' )
  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      '''sshleifer/tiny-gpt2'''
     lowerCAmelCase_			=      PyTorchBenchmarkArguments(
         models=[MODEL_ID]    ,     training=_lowerCamelCase    ,     inference=_lowerCamelCase    ,     sequence_lengths=[8]    ,     batch_sizes=[1]    ,     fpaa=_lowerCamelCase    ,     multi_process=_lowerCamelCase    ,     )
     lowerCAmelCase_			=      PyTorchBenchmark(_lowerCamelCase )
     lowerCAmelCase_			=      benchmark.run()
     self.check_results_dict_not_empty(results.time_train_result )
     self.check_results_dict_not_empty(results.memory_train_result )





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      '''sshleifer/tiny-gpt2'''
     lowerCAmelCase_			=      AutoConfig.from_pretrained(_lowerCamelCase )
     lowerCAmelCase_			=      PyTorchBenchmarkArguments(
         models=[MODEL_ID]    ,     training=_lowerCamelCase    ,     inference=_lowerCamelCase    ,     sequence_lengths=[8]    ,     batch_sizes=[1]    ,     multi_process=_lowerCamelCase    ,     )
     lowerCAmelCase_			=      PyTorchBenchmark(_lowerCamelCase    ,     configs=[config] )
     lowerCAmelCase_			=      benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result )
     self.check_results_dict_not_empty(results.memory_inference_result )





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      '''sshleifer/tinier_bart'''
     lowerCAmelCase_			=      AutoConfig.from_pretrained(_lowerCamelCase )
     lowerCAmelCase_			=      PyTorchBenchmarkArguments(
         models=[MODEL_ID]    ,     training=_lowerCamelCase    ,     inference=_lowerCamelCase    ,     sequence_lengths=[8]    ,     batch_sizes=[1]    ,     multi_process=_lowerCamelCase    ,     )
     lowerCAmelCase_			=      PyTorchBenchmark(_lowerCamelCase    ,     configs=[config] )
     lowerCAmelCase_			=      benchmark.run()
     self.check_results_dict_not_empty(results.time_inference_result )
     self.check_results_dict_not_empty(results.memory_inference_result )





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      '''sshleifer/tiny-gpt2'''
     lowerCAmelCase_			=      AutoConfig.from_pretrained(_lowerCamelCase )
     lowerCAmelCase_			=      PyTorchBenchmarkArguments(
         models=[MODEL_ID]    ,     training=_lowerCamelCase    ,     inference=_lowerCamelCase    ,     sequence_lengths=[8]    ,     batch_sizes=[1]    ,     multi_process=_lowerCamelCase    ,     )
     lowerCAmelCase_			=      PyTorchBenchmark(_lowerCamelCase    ,     configs=[config] )
     lowerCAmelCase_			=      benchmark.run()
     self.check_results_dict_not_empty(results.time_train_result )
     self.check_results_dict_not_empty(results.memory_train_result )





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      '''sshleifer/tinier_bart'''
     lowerCAmelCase_			=      AutoConfig.from_pretrained(_lowerCamelCase )
     lowerCAmelCase_			=      PyTorchBenchmarkArguments(
         models=[MODEL_ID]    ,     training=_lowerCamelCase    ,     inference=_lowerCamelCase    ,     sequence_lengths=[8]    ,     batch_sizes=[1]    ,     multi_process=_lowerCamelCase    ,     )
     lowerCAmelCase_			=      PyTorchBenchmark(_lowerCamelCase    ,     configs=[config] )
     lowerCAmelCase_			=      benchmark.run()
     self.check_results_dict_not_empty(results.time_train_result )
     self.check_results_dict_not_empty(results.memory_train_result )





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      '''sshleifer/tiny-gpt2'''
     with tempfile.TemporaryDirectory() as tmp_dir:
        lowerCAmelCase_			=      PyTorchBenchmarkArguments(
            models=[MODEL_ID]    ,     training=_lowerCamelCase    ,     inference=_lowerCamelCase    ,     save_to_csv=_lowerCamelCase    ,     sequence_lengths=[8]    ,     batch_sizes=[1]    ,     inference_time_csv_file=os.path.join(_lowerCamelCase    ,     '''inf_time.csv''' )    ,     train_memory_csv_file=os.path.join(_lowerCamelCase    ,     '''train_mem.csv''' )    ,     inference_memory_csv_file=os.path.join(_lowerCamelCase    ,     '''inf_mem.csv''' )    ,     train_time_csv_file=os.path.join(_lowerCamelCase    ,     '''train_time.csv''' )    ,     env_info_csv_file=os.path.join(_lowerCamelCase    ,     '''env.csv''' )    ,     multi_process=_lowerCamelCase    ,     )
        lowerCAmelCase_			=      PyTorchBenchmark(_lowerCamelCase )
        benchmark.run()
        self.assertTrue(Path(os.path.join(_lowerCamelCase    ,     '''inf_time.csv''' ) ).exists() )
        self.assertTrue(Path(os.path.join(_lowerCamelCase    ,     '''train_time.csv''' ) ).exists() )
        self.assertTrue(Path(os.path.join(_lowerCamelCase    ,     '''inf_mem.csv''' ) ).exists() )
        self.assertTrue(Path(os.path.join(_lowerCamelCase    ,     '''train_mem.csv''' ) ).exists() )
        self.assertTrue(Path(os.path.join(_lowerCamelCase    ,     '''env.csv''' ) ).exists() )







  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      '''sshleifer/tiny-gpt2'''

     def _check_summary_is_not_empty(_lowerCamelCase ):
        self.assertTrue(hasattr(_lowerCamelCase    ,     '''sequential''' ) )
        self.assertTrue(hasattr(_lowerCamelCase    ,     '''cumulative''' ) )
        self.assertTrue(hasattr(_lowerCamelCase    ,     '''current''' ) )
        self.assertTrue(hasattr(_lowerCamelCase    ,     '''total''' ) )

     with tempfile.TemporaryDirectory() as tmp_dir:
        lowerCAmelCase_			=      PyTorchBenchmarkArguments(
            models=[MODEL_ID]    ,     training=_lowerCamelCase    ,     inference=_lowerCamelCase    ,     sequence_lengths=[8]    ,     batch_sizes=[1]    ,     log_filename=os.path.join(_lowerCamelCase    ,     '''log.txt''' )    ,     log_print=_lowerCamelCase    ,     trace_memory_line_by_line=_lowerCamelCase    ,     multi_process=_lowerCamelCase    ,     )
        lowerCAmelCase_			=      PyTorchBenchmark(_lowerCamelCase )
        lowerCAmelCase_			=      benchmark.run()
        _check_summary_is_not_empty(result.inference_summary )
        _check_summary_is_not_empty(result.train_summary )
        self.assertTrue(Path(os.path.join(_lowerCamelCase    ,     '''log.txt''' ) ).exists() )





