
'''simple docstring'''


from __future__ import annotations

import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import is_tf_available, is_vision_available

from ...test_modeling_tf_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_tf_bert import TFBertModelTester
from ..clip.test_modeling_tf_clip import TFCLIPVisionModelTester
from ..deit.test_modeling_tf_deit import TFDeiTModelTester
from ..roberta.test_modeling_tf_roberta import TFRobertaModelTester
from ..vit.test_modeling_tf_vit import TFViTModelTester


if is_tf_available():
   from transformers import (
       TFBertModel,
       TFCLIPVisionModel,
       TFDeiTModel,
       TFRobertaModel,
       TFVisionTextDualEncoderModel,
       TFViTModel,
       VisionTextDualEncoderConfig,
   )

if is_vision_available():
   from PIL import Image

   from transformers import VisionTextDualEncoderProcessor



def  __snake_case     (    _UpperCAmelCase					:       List[str]):
   if isinstance(_UpperCAmelCase,      collections.abc.Iterable):
      return x
   return (x, x)


@require_tf
class 						lowercase__	:






   '''simple docstring'''







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__   ):

      '''simple docstring'''







      pass







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      pass







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      pass







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=None ,   **lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  VisionTextDualEncoderConfig.from_vision_text_configs(lowerCamelCase__ ,   lowerCamelCase__   )

      UpperCamelCase       =  TFVisionTextDualEncoderModel(lowerCamelCase__   )

      UpperCamelCase       =  model(input_ids=lowerCamelCase__ ,   pixel_values=lowerCamelCase__ ,   attention_mask=lowerCamelCase__   )

      self.assertEqual(output['''text_embeds'''].shape ,   (input_ids.shape[0], config.projection_dim)   )
      self.assertEqual(output['''image_embeds'''].shape ,   (pixel_values.shape[0], config.projection_dim)   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=None ,   **lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase     , UpperCamelCase       =  self.get_vision_text_model(lowerCamelCase__ ,   lowerCamelCase__   )
      UpperCamelCase       =  TFVisionTextDualEncoderModel(vision_model=lowerCamelCase__ ,   text_model=lowerCamelCase__   )

      UpperCamelCase       =  model(input_ids=lowerCamelCase__ ,   pixel_values=lowerCamelCase__ ,   attention_mask=lowerCamelCase__   )

      self.assertEqual(output['''text_embeds'''].shape ,   (input_ids.shape[0], model.config.projection_dim)   )
      self.assertEqual(output['''image_embeds'''].shape ,   (pixel_values.shape[0], model.config.projection_dim)   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=None ,   **lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase     , UpperCamelCase       =  self.get_vision_text_model(lowerCamelCase__ ,   lowerCamelCase__   )
      UpperCamelCase       =  {'''vision_model''': vision_model, '''text_model''': text_model}
      UpperCamelCase       =  TFVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCamelCase__   )

      UpperCamelCase       =  model(input_ids=lowerCamelCase__ ,   pixel_values=lowerCamelCase__ ,   attention_mask=lowerCamelCase__   )

      self.assertEqual(output['''text_embeds'''].shape ,   (input_ids.shape[0], model.config.projection_dim)   )
      self.assertEqual(output['''image_embeds'''].shape ,   (pixel_values.shape[0], model.config.projection_dim)   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=None ,   **lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase     , UpperCamelCase       =  self.get_vision_text_model(lowerCamelCase__ ,   lowerCamelCase__   )
      UpperCamelCase       =  TFVisionTextDualEncoderModel(vision_model=lowerCamelCase__ ,   text_model=lowerCamelCase__   )

      UpperCamelCase       =  model(input_ids=lowerCamelCase__ ,   pixel_values=lowerCamelCase__ ,   attention_mask=lowerCamelCase__   )
      UpperCamelCase       =  output[0].numpy()

      with tempfile.TemporaryDirectory() as tmpdirname:
         model.save_pretrained(lowerCamelCase__   )
         UpperCamelCase       =  TFVisionTextDualEncoderModel.from_pretrained(lowerCamelCase__   )

         UpperCamelCase       =  model(input_ids=lowerCamelCase__ ,   pixel_values=lowerCamelCase__ ,   attention_mask=lowerCamelCase__   )
         UpperCamelCase       =  after_output[0].numpy()
         UpperCamelCase       =  np.amax(np.abs(out_a - out_a   )   )
         self.assertLessEqual(lowerCamelCase__ ,   1e-5   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=None ,   **lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase     , UpperCamelCase       =  self.get_vision_text_model(lowerCamelCase__ ,   lowerCamelCase__   )
      UpperCamelCase       =  TFVisionTextDualEncoderModel(vision_model=lowerCamelCase__ ,   text_model=lowerCamelCase__   )

      UpperCamelCase       =  model(
          input_ids=lowerCamelCase__ ,   pixel_values=lowerCamelCase__ ,   attention_mask=lowerCamelCase__ ,   output_attentions=lowerCamelCase__   )

      UpperCamelCase       =  output.vision_model_output.attentions
      self.assertEqual(len(lowerCamelCase__   ) ,   vision_config.num_hidden_layers   )

      # in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
      UpperCamelCase       =  to_atuple(vision_model.config.image_size   )
      UpperCamelCase       =  to_atuple(vision_model.config.patch_size   )
      UpperCamelCase       =  (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
      UpperCamelCase       =  num_patches + 1
      self.assertEqual(vision_attentions[0].shape[-3:] ,   (vision_config.num_attention_heads, seq_len, seq_len)   )

      UpperCamelCase       =  output.text_model_output.attentions
      self.assertEqual(len(lowerCamelCase__   ) ,   text_config.num_hidden_layers   )

      self.assertEqual(
          text_attentions[0].shape[-3:] ,   (text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1]) ,   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  np.abs((a - b)   ).max()
      self.assertLessEqual(lowerCamelCase__ ,   lowerCamelCase__ ,   f'Difference between torch and flax is {diff} (>= {tol}).'   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.prepare_config_and_inputs()
      self.check_vision_text_dual_encoder_model(**lowerCamelCase__   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.prepare_config_and_inputs()
      self.check_model_from_pretrained_configs(**lowerCamelCase__   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.prepare_config_and_inputs()
      self.check_vision_text_dual_encoder_from_pretrained(**lowerCamelCase__   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.prepare_config_and_inputs()
      self.check_save_load(**lowerCamelCase__   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.prepare_config_and_inputs()
      self.check_vision_text_output_attention(**lowerCamelCase__   )







   @slow
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase     , UpperCamelCase       =  self.get_pretrained_model_and_inputs()

      UpperCamelCase       =  model_a(**lowerCamelCase__   )
      UpperCamelCase       =  outputs[0].numpy()

      with tempfile.TemporaryDirectory() as tmp_dirname:
         model_a.save_pretrained(lowerCamelCase__   )
         UpperCamelCase       =  TFVisionTextDualEncoderModel.from_pretrained(lowerCamelCase__   )

         UpperCamelCase       =  model_a(**lowerCamelCase__   )
         UpperCamelCase       =  after_outputs[0].numpy()
         UpperCamelCase       =  np.amax(np.abs(out_a - out_a   )   )
         self.assertLessEqual(lowerCamelCase__ ,   1e-5   )




@require_tf
class 						lowercase__	(  snake_case_,					unittest.TestCase   ):






   '''simple docstring'''







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  TFVisionTextDualEncoderModel.from_vision_text_pretrained(
          '''hf-internal-testing/tiny-random-vit''' ,   '''hf-internal-testing/tiny-random-bert'''   )
      UpperCamelCase       =  1_3
      UpperCamelCase       =  floats_tensor(
          [
              batch_size,
              model.vision_model.config.num_channels,
              model.vision_model.config.image_size,
              model.vision_model.config.image_size,
          ]   )
      UpperCamelCase       =  ids_tensor([batch_size, 4] ,   model.text_model.config.vocab_size   )
      UpperCamelCase       =  random_attention_mask([batch_size, 4]   )
      UpperCamelCase       =  {'''pixel_values''': pixel_values, '''input_ids''': input_ids, '''attention_mask''': attention_mask}

      return model, inputs







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  TFViTModel(lowerCamelCase__ ,   name='''vision_model'''   )
      UpperCamelCase       =  TFBertModel(lowerCamelCase__ ,   name='''text_model'''   )
      return vision_model, text_model







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  TFViTModelTester(self   )
      UpperCamelCase       =  TFBertModelTester(self   )
      UpperCamelCase       =  vit_model_tester.prepare_config_and_inputs()
      UpperCamelCase       =  bert_model_tester.prepare_config_and_inputs()

      UpperCamelCase     , UpperCamelCase     , UpperCamelCase       =  vision_config_and_inputs

      (
          (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , 
      )       =  text_config_and_inputs

      return {
          "text_config": text_config,
          "vision_config": vision_config,
          "pixel_values": pixel_values,
          "attention_mask": input_mask,
          "input_ids": input_ids,
          "text_token_type_ids": token_type_ids,
          "text_sequence_labels": sequence_labels,
          "text_token_labels": token_labels,
          "text_choice_labels": choice_labels,
      }




@require_tf
class 						lowercase__	(  snake_case_,					unittest.TestCase   ):






   '''simple docstring'''







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  TFVisionTextDualEncoderModel.from_vision_text_pretrained(
          '''Rocketknight1/tiny-random-deit-tf''' ,   '''hf-internal-testing/tiny-random-roberta'''   )
      UpperCamelCase       =  1_3
      UpperCamelCase       =  floats_tensor(
          [
              batch_size,
              model.vision_model.config.num_channels,
              model.vision_model.config.image_size,
              model.vision_model.config.image_size,
          ]   )
      UpperCamelCase       =  ids_tensor([batch_size, 4] ,   model.text_model.config.vocab_size   )
      UpperCamelCase       =  random_attention_mask([batch_size, 4]   )
      UpperCamelCase       =  {'''pixel_values''': pixel_values, '''input_ids''': input_ids, '''attention_mask''': attention_mask}

      return model, inputs







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=None ,   **lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase     , UpperCamelCase       =  self.get_vision_text_model(lowerCamelCase__ ,   lowerCamelCase__   )
      UpperCamelCase       =  TFVisionTextDualEncoderModel(vision_model=lowerCamelCase__ ,   text_model=lowerCamelCase__   )

      UpperCamelCase       =  model(
          input_ids=lowerCamelCase__ ,   pixel_values=lowerCamelCase__ ,   attention_mask=lowerCamelCase__ ,   output_attentions=lowerCamelCase__   )

      UpperCamelCase       =  output.vision_model_output.attentions
      self.assertEqual(len(lowerCamelCase__   ) ,   vision_config.num_hidden_layers   )

      # in DEiT, the seq_len equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
      UpperCamelCase       =  to_atuple(vision_model.config.image_size   )
      UpperCamelCase       =  to_atuple(vision_model.config.patch_size   )
      UpperCamelCase       =  (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
      UpperCamelCase       =  num_patches + 2
      self.assertEqual(vision_attentions[0].shape[-3:] ,   (vision_config.num_attention_heads, seq_len, seq_len)   )

      UpperCamelCase       =  output.text_model_output.attentions
      self.assertEqual(len(lowerCamelCase__   ) ,   text_config.num_hidden_layers   )

      self.assertEqual(
          text_attentions[0].shape[-3:] ,   (text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1]) ,   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  TFDeiTModel(lowerCamelCase__ ,   name='''vision_model'''   )
      UpperCamelCase       =  TFRobertaModel(lowerCamelCase__ ,   name='''text_model'''   )
      return vision_model, text_model







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  TFDeiTModelTester(self   )
      UpperCamelCase       =  TFRobertaModelTester(self   )
      UpperCamelCase       =  vit_model_tester.prepare_config_and_inputs()
      UpperCamelCase       =  bert_model_tester.prepare_config_and_inputs()

      UpperCamelCase     , UpperCamelCase     , UpperCamelCase       =  vision_config_and_inputs

      (
          (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , 
      )       =  text_config_and_inputs

      return {
          "text_config": text_config,
          "vision_config": vision_config,
          "pixel_values": pixel_values,
          "attention_mask": input_mask,
          "input_ids": input_ids,
          "text_token_type_ids": token_type_ids,
          "text_sequence_labels": sequence_labels,
          "text_token_labels": token_labels,
          "text_choice_labels": choice_labels,
      }




@require_tf
class 						lowercase__	(  snake_case_,					unittest.TestCase   ):






   '''simple docstring'''







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  TFVisionTextDualEncoderModel.from_vision_text_pretrained(
          '''Rocketknight1/tiny-random-clip-tf''' ,   '''hf-internal-testing/tiny-random-bert'''   )
      UpperCamelCase       =  1_3
      UpperCamelCase       =  floats_tensor(
          [
              batch_size,
              model.vision_model.config.num_channels,
              model.vision_model.config.image_size,
              model.vision_model.config.image_size,
          ]   )
      UpperCamelCase       =  ids_tensor([batch_size, 4] ,   model.text_model.config.vocab_size   )
      UpperCamelCase       =  random_attention_mask([batch_size, 4]   )
      UpperCamelCase       =  {'''pixel_values''': pixel_values, '''input_ids''': input_ids, '''attention_mask''': attention_mask}

      return model, inputs







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  TFCLIPVisionModel(lowerCamelCase__ ,   name='''vision_model'''   )
      UpperCamelCase       =  TFBertModel(lowerCamelCase__ ,   name='''text_model'''   )
      return vision_model, text_model







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  TFCLIPVisionModelTester(self   )
      UpperCamelCase       =  TFBertModelTester(self   )
      UpperCamelCase       =  clip_model_tester.prepare_config_and_inputs()
      UpperCamelCase       =  bert_model_tester.prepare_config_and_inputs()

      UpperCamelCase     , UpperCamelCase       =  vision_config_and_inputs

      (
          (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , (
          UpperCamelCase
      )     , 
      )       =  text_config_and_inputs

      return {
          "text_config": text_config,
          "vision_config": vision_config,
          "pixel_values": pixel_values,
          "attention_mask": input_mask,
          "input_ids": input_ids,
          "text_token_type_ids": token_type_ids,
          "text_sequence_labels": sequence_labels,
          "text_token_labels": token_labels,
          "text_choice_labels": choice_labels,
      }




@require_vision
@require_tf
class 						lowercase__	(  unittest.TestCase   ):






   '''simple docstring'''







   @slow
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  TFVisionTextDualEncoderModel.from_pretrained(
          '''clip-italian/clip-italian''' ,   logit_scale_init_value=1.0 ,   from_pt=lowerCamelCase__   )
      UpperCamelCase       =  VisionTextDualEncoderProcessor.from_pretrained('''clip-italian/clip-italian'''   )

      UpperCamelCase       =  Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''   )
      UpperCamelCase       =  processor(
          text=['''una foto di un gatto''', '''una foto di un cane'''] ,   images=lowerCamelCase__ ,   padding=lowerCamelCase__ ,   return_tensors='''np'''   )

      UpperCamelCase       =  model(**lowerCamelCase__   )

      # verify the logits
      self.assertEqual(outputs.logits_per_image.shape ,   (inputs.pixel_values.shape[0], inputs.input_ids.shape[0])   )
      self.assertEqual(
          outputs.logits_per_text.shape ,   (inputs.input_ids.shape[0], inputs.pixel_values.shape[0]) ,   )

      UpperCamelCase       =  np.array([[1.2284727, 0.3104122]]   )

      self.assertTrue(np.allclose(outputs.logits_per_image.numpy() ,   lowerCamelCase__ ,   atol=1e-3   )   )







'''simple docstring'''


import torch
from torch import nn


class 						lowercase__	(  nn.Module   ):






   '''simple docstring'''







   def __init__(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=1 ,   lowerCamelCase__=False   ):

      '''simple docstring'''







      super().__init__()

      UpperCamelCase       =  n_token
      UpperCamelCase       =  d_embed
      UpperCamelCase       =  d_proj

      UpperCamelCase       =  cutoffs + [n_token]
      UpperCamelCase       =  [0] + self.cutoffs
      UpperCamelCase       =  div_val

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

      if self.n_clusters > 0:
         UpperCamelCase       =  nn.Parameter(torch.zeros(self.n_clusters ,   self.d_embed   )   )
         UpperCamelCase       =  nn.Parameter(torch.zeros(self.n_clusters   )   )

      UpperCamelCase       =  nn.ModuleList()
      UpperCamelCase       =  nn.ParameterList()

      if div_val == 1:
         for i in range(len(self.cutoffs   )   ):
            if d_proj != d_embed:
               self.out_projs.append(nn.Parameter(torch.FloatTensor(lowerCamelCase__ ,   lowerCamelCase__   )   )   )
            else:
               self.out_projs.append(lowerCamelCase__   )

         self.out_layers.append(nn.Linear(lowerCamelCase__ ,   lowerCamelCase__   )   )
      else:
         for i in range(len(self.cutoffs   )   ):
            UpperCamelCase     , UpperCamelCase       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
            UpperCamelCase       =  d_embed // (div_val**i)

            self.out_projs.append(nn.Parameter(torch.FloatTensor(lowerCamelCase__ ,   lowerCamelCase__   )   )   )

            self.out_layers.append(nn.Linear(lowerCamelCase__ ,   r_idx - l_idx   )   )

      UpperCamelCase       =  keep_order







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   ):

      '''simple docstring'''







      if proj is None:
         UpperCamelCase       =  nn.functional.linear(lowerCamelCase__ ,   lowerCamelCase__ ,   bias=lowerCamelCase__   )
      else:
         # if CUDA_MAJOR <= 9 and CUDA_MINOR <= 1:
         UpperCamelCase       =  nn.functional.linear(lowerCamelCase__ ,   proj.t().contiguous()   )
         UpperCamelCase       =  nn.functional.linear(lowerCamelCase__ ,   lowerCamelCase__ ,   bias=lowerCamelCase__   )
         # else:
         #     logit = torch.einsum('bd,de,ev->bv', (hidden, proj, weight.t()))
         #     if bias is not None:
         #         logit = logit + bias

      return logit







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__=None ,   lowerCamelCase__=False   ):

      '''simple docstring'''







      if labels is not None:
         # Shift so that tokens < n predict n
         UpperCamelCase       =  hidden[..., :-1, :].contiguous()
         UpperCamelCase       =  labels[..., 1:].contiguous()
         UpperCamelCase       =  hidden.view(-1 ,   hidden.size(-1   )   )
         UpperCamelCase       =  labels.view(-1   )
         if hidden.size(0   ) != labels.size(0   ):
            raise RuntimeError('''Input and labels should have the same size in the batch dimension.'''   )
      else:
         UpperCamelCase       =  hidden.view(-1 ,   hidden.size(-1   )   )

      if self.n_clusters == 0:
         UpperCamelCase       =  self._compute_logit(lowerCamelCase__ ,   self.out_layers[0].weight ,   self.out_layers[0].bias ,   self.out_projs[0]   )
         if labels is not None:
            UpperCamelCase       =  labels != -1_0_0
            UpperCamelCase       =  torch.zeros_like(lowerCamelCase__ ,   dtype=hidden.dtype ,   device=hidden.device   )
            UpperCamelCase       =  (
                -nn.functional.log_softmax(lowerCamelCase__ ,   dim=-1   )[mask].gather(1 ,   labels[mask].unsqueeze(1   )   ).squeeze(1   )
            )
         else:
            UpperCamelCase       =  nn.functional.log_softmax(lowerCamelCase__ ,   dim=-1   )
      else:
         # construct weights and biases
         UpperCamelCase     , UpperCamelCase       =  [], []
         for i in range(len(self.cutoffs   )   ):
            if self.div_val == 1:
               UpperCamelCase     , UpperCamelCase       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
               UpperCamelCase       =  self.out_layers[0].weight[l_idx:r_idx]
               UpperCamelCase       =  self.out_layers[0].bias[l_idx:r_idx]
            else:
               UpperCamelCase       =  self.out_layers[i].weight
               UpperCamelCase       =  self.out_layers[i].bias

            if i == 0:
               UpperCamelCase       =  torch.cat([weight_i, self.cluster_weight] ,   dim=0   )
               UpperCamelCase       =  torch.cat([bias_i, self.cluster_bias] ,   dim=0   )

            weights.append(lowerCamelCase__   )
            biases.append(lowerCamelCase__   )

         UpperCamelCase     , UpperCamelCase     , UpperCamelCase       =  weights[0], biases[0], self.out_projs[0]

         UpperCamelCase       =  self._compute_logit(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   )
         UpperCamelCase       =  nn.functional.log_softmax(lowerCamelCase__ ,   dim=1   )

         if labels is None:
            UpperCamelCase       =  hidden.new_empty((head_logit.size(0   ), self.n_token)   )
         else:
            UpperCamelCase       =  torch.zeros_like(lowerCamelCase__ ,   dtype=hidden.dtype ,   device=hidden.device   )

         UpperCamelCase       =  0
         UpperCamelCase       =  [0] + self.cutoffs
         for i in range(len(lowerCamelCase__   ) - 1   ):
            UpperCamelCase     , UpperCamelCase       =  cutoff_values[i], cutoff_values[i + 1]

            if labels is not None:
               UpperCamelCase       =  (labels >= l_idx) & (labels < r_idx)
               UpperCamelCase       =  mask_i.nonzero().squeeze()

               if indices_i.numel() == 0:
                  continue

               UpperCamelCase       =  labels.index_select(0 ,   lowerCamelCase__   ) - l_idx
               UpperCamelCase       =  head_logprob.index_select(0 ,   lowerCamelCase__   )
               UpperCamelCase       =  hidden.index_select(0 ,   lowerCamelCase__   )
            else:
               UpperCamelCase       =  hidden

            if i == 0:
               if labels is not None:
                  UpperCamelCase       =  head_logprob_i.gather(1 ,   target_i[:, None]   ).squeeze(1   )
               else:
                  UpperCamelCase       =  head_logprob[:, : self.cutoffs[0]]
            else:
               UpperCamelCase     , UpperCamelCase     , UpperCamelCase       =  weights[i], biases[i], self.out_projs[i]

               UpperCamelCase       =  self._compute_logit(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   )
               UpperCamelCase       =  nn.functional.log_softmax(lowerCamelCase__ ,   dim=1   )
               UpperCamelCase       =  self.cutoffs[0] + i - 1  # No probability for the head cluster
               if labels is not None:
                  UpperCamelCase       =  head_logprob_i[:, cluster_prob_idx] + tail_logprob_i.gather(
                      1 ,   target_i[:, None]   ).squeeze(1   )
               else:
                  UpperCamelCase       =  head_logprob[:, cluster_prob_idx, None] + tail_logprob_i
                  UpperCamelCase       =  logprob_i

            if labels is not None:
               if (hasattr(self ,   '''keep_order'''   ) and self.keep_order) or keep_order:
                  out.index_copy_(0 ,   lowerCamelCase__ ,   -logprob_i   )
               else:
                  out[offset : offset + logprob_i.size(0   )].copy_(-logprob_i   )
               offset += logprob_i.size(0   )

      return out







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      if self.n_clusters == 0:
         UpperCamelCase       =  self._compute_logit(lowerCamelCase__ ,   self.out_layers[0].weight ,   self.out_layers[0].bias ,   self.out_projs[0]   )
         return nn.functional.log_softmax(lowerCamelCase__ ,   dim=-1   )
      else:
         # construct weights and biases
         UpperCamelCase     , UpperCamelCase       =  [], []
         for i in range(len(self.cutoffs   )   ):
            if self.div_val == 1:
               UpperCamelCase     , UpperCamelCase       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
               UpperCamelCase       =  self.out_layers[0].weight[l_idx:r_idx]
               UpperCamelCase       =  self.out_layers[0].bias[l_idx:r_idx]
            else:
               UpperCamelCase       =  self.out_layers[i].weight
               UpperCamelCase       =  self.out_layers[i].bias

            if i == 0:
               UpperCamelCase       =  torch.cat([weight_i, self.cluster_weight] ,   dim=0   )
               UpperCamelCase       =  torch.cat([bias_i, self.cluster_bias] ,   dim=0   )

            weights.append(lowerCamelCase__   )
            biases.append(lowerCamelCase__   )

         UpperCamelCase     , UpperCamelCase     , UpperCamelCase       =  weights[0], biases[0], self.out_projs[0]
         UpperCamelCase       =  self._compute_logit(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   )

         UpperCamelCase       =  hidden.new_empty((head_logit.size(0   ), self.n_token)   )
         UpperCamelCase       =  nn.functional.log_softmax(lowerCamelCase__ ,   dim=1   )

         UpperCamelCase       =  [0] + self.cutoffs
         for i in range(len(lowerCamelCase__   ) - 1   ):
            UpperCamelCase     , UpperCamelCase       =  cutoff_values[i], cutoff_values[i + 1]

            if i == 0:
               UpperCamelCase       =  head_logprob[:, : self.cutoffs[0]]
            else:
               UpperCamelCase     , UpperCamelCase     , UpperCamelCase       =  weights[i], biases[i], self.out_projs[i]

               UpperCamelCase       =  self._compute_logit(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__   )
               UpperCamelCase       =  nn.functional.log_softmax(lowerCamelCase__ ,   dim=1   )

               UpperCamelCase       =  head_logprob[:, -i] + tail_logprob_i
               UpperCamelCase       =  logprob_i

         return out





