




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 		_A(       UpperCamelCase__							:			List[Any]				)    -> Dict:
							'''simple docstring'''


							if isinstance(UpperCamelCase__				,       collections.abc.Iterable				):
														return x
							return (x, x)







@require_tf
class a :
						"""simple docstring"""

						def     UpperCAmelCase_      ( self :     int				,							lowerCamelCase__ :     List[Any]				,							lowerCamelCase__ :     Tuple				)    ->						Optional[Any]:
													"""simple docstring"""


													pass






						def     UpperCAmelCase_      ( self :     List[str]				)    ->						List[Any]:
													"""simple docstring"""


													pass






						def     UpperCAmelCase_      ( self :     Dict				)    ->						Optional[Any]:
													"""simple docstring"""


													pass






						def     UpperCAmelCase_      ( self :     Dict				,							lowerCamelCase__ :     Optional[int]				,							lowerCamelCase__ :     List[Any]				,							lowerCamelCase__ :     Optional[Any]				,							lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     Dict=None				,							**lowerCamelCase__ :     int				)    ->						Dict:
													"""simple docstring"""


													__lowercase      =    VisionTextDualEncoderConfig.from_vision_text_configs(lowerCamelCase__				,							lowerCamelCase__				)

													__lowercase      =    TFVisionTextDualEncoderModel(lowerCamelCase__				)

													__lowercase      =    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 :     List[str]				,							lowerCamelCase__ :     Dict				,							lowerCamelCase__ :     Any				,							lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     str				,							lowerCamelCase__ :     Dict=None				,							**lowerCamelCase__ :     Optional[int]				)    ->						Optional[Any]:
													"""simple docstring"""


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

													__lowercase      =    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 :     Optional[int]				,							lowerCamelCase__ :     Any				,							lowerCamelCase__ :     Union[str, Any]				,							lowerCamelCase__ :     int				,							lowerCamelCase__ :     Optional[Any]				,							lowerCamelCase__ :     Optional[Any]=None				,							**lowerCamelCase__ :     Tuple				)    ->						List[str]:
													"""simple docstring"""


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

													__lowercase      =    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 :     List[str]				,							lowerCamelCase__ :     Optional[Any]				,							lowerCamelCase__ :     Dict				,							lowerCamelCase__ :     str				,							lowerCamelCase__ :     Any				,							lowerCamelCase__ :     str=None				,							**lowerCamelCase__ :     List[str]				)    ->						Optional[int]:
													"""simple docstring"""


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

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

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

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






						def     UpperCAmelCase_      ( self :     Optional[Any]				,							lowerCamelCase__ :     str				,							lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     str				,							lowerCamelCase__ :     Union[str, Any]				,							lowerCamelCase__ :     int=None				,							**lowerCamelCase__ :     Optional[int]				)    ->						Optional[int]:
													"""simple docstring"""


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

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

													__lowercase      =    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)
													__lowercase      =    to_atuple(vision_model.config.image_size				)
													__lowercase      =    to_atuple(vision_model.config.patch_size				)
													__lowercase      =    (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
													__lowercase      =    num_patches + 1
													self.assertEqual(vision_attentions[0].shape[-3:]				,							(vision_config.num_attention_heads, seq_len, seq_len)				)

													__lowercase      =    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 :     List[Any]				,							lowerCamelCase__ :     np.ndarray				,							lowerCamelCase__ :     np.ndarray				,							lowerCamelCase__ :     float				)    ->						List[Any]:
													"""simple docstring"""


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






						def     UpperCAmelCase_      ( self :     Any				)    ->						List[str]:
													"""simple docstring"""


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






						def     UpperCAmelCase_      ( self :     Dict				)    ->						Any:
													"""simple docstring"""


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






						def     UpperCAmelCase_      ( self :     List[Any]				)    ->						Tuple:
													"""simple docstring"""


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






						def     UpperCAmelCase_      ( self :     Tuple				)    ->						int:
													"""simple docstring"""


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






						def     UpperCAmelCase_      ( self :     Union[str, Any]				)    ->						int:
													"""simple docstring"""


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






						@slow
						def     UpperCAmelCase_      ( self :     Union[str, Any]				)    ->						int:
													"""simple docstring"""


													__lowercase     ,     __lowercase      =    self.get_pretrained_model_and_inputs()

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

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

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







@require_tf
class a (     __SCREAMING_SNAKE_CASE       ,						unittest.TestCase  ):
						"""simple docstring"""

						def     UpperCAmelCase_      ( self :     Optional[int]				)    ->						Union[str, Any]:
													"""simple docstring"""


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

													return model, inputs






						def     UpperCAmelCase_      ( self :     List[str]				,							lowerCamelCase__ :     Any				,							lowerCamelCase__ :     List[str]				)    ->						Optional[int]:
													"""simple docstring"""


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






						def     UpperCAmelCase_      ( self :     Any				)    ->						List[str]:
													"""simple docstring"""


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

													__lowercase     ,     __lowercase     ,     __lowercase      =    vision_config_and_inputs

													(
													    (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,          
													)      =    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 a (     __SCREAMING_SNAKE_CASE       ,						unittest.TestCase  ):
						"""simple docstring"""

						def     UpperCAmelCase_      ( self :     Any				)    ->						Tuple:
													"""simple docstring"""


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

													return model, inputs






						def     UpperCAmelCase_      ( self :     int				,							lowerCamelCase__ :     Any				,							lowerCamelCase__ :     List[str]				,							lowerCamelCase__ :     Optional[int]				,							lowerCamelCase__ :     List[str]				,							lowerCamelCase__ :     Optional[int]=None				,							**lowerCamelCase__ :     Union[str, Any]				)    ->						Tuple:
													"""simple docstring"""


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

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

													__lowercase      =    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)
													__lowercase      =    to_atuple(vision_model.config.image_size				)
													__lowercase      =    to_atuple(vision_model.config.patch_size				)
													__lowercase      =    (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
													__lowercase      =    num_patches + 2
													self.assertEqual(vision_attentions[0].shape[-3:]				,							(vision_config.num_attention_heads, seq_len, seq_len)				)

													__lowercase      =    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 :     List[str]				,							lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     Dict				)    ->						Optional[Any]:
													"""simple docstring"""


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






						def     UpperCAmelCase_      ( self :     int				)    ->						Dict:
													"""simple docstring"""


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

													__lowercase     ,     __lowercase     ,     __lowercase      =    vision_config_and_inputs

													(
													    (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,          
													)      =    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 a (     __SCREAMING_SNAKE_CASE       ,						unittest.TestCase  ):
						"""simple docstring"""

						def     UpperCAmelCase_      ( self :     List[str]				)    ->						Optional[int]:
													"""simple docstring"""


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

													return model, inputs






						def     UpperCAmelCase_      ( self :     Tuple				,							lowerCamelCase__ :     Union[str, Any]				,							lowerCamelCase__ :     List[str]				)    ->						List[str]:
													"""simple docstring"""


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






						def     UpperCAmelCase_      ( self :     Dict				)    ->						Tuple:
													"""simple docstring"""


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

													__lowercase     ,     __lowercase      =    vision_config_and_inputs

													(
													    (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,     (
													    __lowercase     
													)     ,          
													)      =    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 a (     unittest.TestCase  ):
						"""simple docstring"""

						@slow
						def     UpperCAmelCase_      ( self :     Union[str, Any]				)    ->						Dict:
													"""simple docstring"""


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

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

													__lowercase      =    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])				,							)

													__lowercase      =    np.array([[1.2_2_8_4_7_2_7, 0.3_1_0_4_1_2_2]]				)

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









import tensorflow as tf

from ...tf_utils import shape_list







class a (     tf.keras.layers.Layer  ):
						"""simple docstring"""

						def __init__( self :     int				,							lowerCamelCase__ :     Union[str, Any]				,							lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     List[Any]				,							lowerCamelCase__ :     Optional[Any]				,							lowerCamelCase__ :     Tuple=1				,							lowerCamelCase__ :     Optional[Any]=False				,							**lowerCamelCase__ :     Optional[int]				)    ->						str:
													"""simple docstring"""


													super().__init__(**lowerCamelCase__				)

													__lowercase      =    vocab_size
													__lowercase      =    d_embed
													__lowercase      =    d_proj

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

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

													__lowercase      =    []
													__lowercase      =    []






						def     UpperCAmelCase_      ( self :     Optional[int]				,							lowerCamelCase__ :     Tuple				)    ->						Tuple:
													"""simple docstring"""


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

																											__lowercase      =    self.add_weight(
																											    shape=(d_emb_i, self.d_proj)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name=f'out_projs_._{i}'				)
																											self.out_projs.append(lowerCamelCase__				)
																											__lowercase      =    self.add_weight(
																											    shape=(r_idx - l_idx, d_emb_i)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name=f'out_layers_._{i}_._weight'				,							)
																											__lowercase      =    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__ :     List[str]				,							lowerCamelCase__ :     int				,							lowerCamelCase__ :     List[Any]				,							lowerCamelCase__ :     List[str]=None				)    ->						Tuple:
													"""simple docstring"""


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






						@staticmethod
						def     UpperCAmelCase_      ( lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     List[str]				)    ->						str:
													"""simple docstring"""


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






						def     UpperCAmelCase_      ( self :     Optional[int]				,							lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     str				,							lowerCamelCase__ :     Optional[int]=True				,							lowerCamelCase__ :     Optional[int]=False				)    ->						List[str]:
													"""simple docstring"""


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

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

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

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

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






