

import unittest

from transformers import XLMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import (
	    XLMForMultipleChoice,
	    XLMForQuestionAnswering,
	    XLMForQuestionAnsweringSimple,
	    XLMForSequenceClassification,
	    XLMForTokenClassification,
	    XLMModel,
	    XLMWithLMHeadModel,
	)
	from transformers.models.xlm.modeling_xlm import XLM_PRETRAINED_MODEL_ARCHIVE_LIST



class        _a		:


					'''simple docstring'''






					def __init__(					self       ,			A__       ,			A__=13       ,			A__=7       ,			A__=True       ,			A__=True       ,			A__=True       ,			A__=True       ,			A__=True       ,			A__=False       ,			A__=False       ,			A__=False       ,			A__=2       ,			A__=99       ,			A__=0       ,			A__=32       ,			A__=5       ,			A__=4       ,			A__=0.1       ,			A__=0.1       ,			A__=512       ,			A__=2       ,			A__=0.0_2       ,			A__=2       ,			A__=4       ,			A__="last"       ,			A__=True       ,			A__=None       ,			A__=0       ,			):
						A__				:       Tuple	  =						parent
						A__				:       int	  =						batch_size
						A__				:       Union[str, Any]	  =						seq_length
						A__				:       List[Any]	  =						is_training
						A__				:       Dict	  =						use_input_lengths
						A__				:       List[str]	  =						use_token_type_ids
						A__				:       List[Any]	  =						use_labels
						A__				:       Union[str, Any]	  =						gelu_activation
						A__				:       Optional[int]	  =						sinusoidal_embeddings
						A__				:       int	  =						causal
						A__				:       Dict	  =						asm
						A__				:       Union[str, Any]	  =						n_langs
						A__				:       Any	  =						vocab_size
						A__				:       str	  =						n_special
						A__				:       int	  =						hidden_size
						A__				:       Tuple	  =						num_hidden_layers
						A__				:       List[str]	  =						num_attention_heads
						A__				:       Union[str, Any]	  =						hidden_dropout_prob
						A__				:       List[Any]	  =						attention_probs_dropout_prob
						A__				:       Tuple	  =						max_position_embeddings
						A__				:       List[str]	  =						type_sequence_label_size
						A__				:       Tuple	  =						initializer_range
						A__				:       Optional[Any]	  =						num_labels
						A__				:       Union[str, Any]	  =						num_choices
						A__				:       Optional[Any]	  =						summary_type
						A__				:       Union[str, Any]	  =						use_proj
						A__				:       str	  =						scope
						A__				:       Tuple	  =						bos_token_id






					def     __A     (					self     ):
						A__				:       Optional[int]	  =						ids_tensor([self.batch_size, self.seq_length]       ,			self.vocab_size     )
						A__				:       List[str]	  =						random_attention_mask([self.batch_size, self.seq_length]     )

						A__				:       int	  =						None
						if self.use_input_lengths:
							A__				:       int	  =						(
							    ids_tensor([self.batch_size]       ,			vocab_size=2     ) + self.seq_length - 2
							)  # small variation of seq_length

						A__				:       int	  =						None
						if self.use_token_type_ids:
							A__				:       List[Any]	  =						ids_tensor([self.batch_size, self.seq_length]       ,			self.n_langs     )

						A__				:       Optional[int]	  =						None
						A__				:       Dict	  =						None
						A__				:       List[Any]	  =						None
						if self.use_labels:
							A__				:       List[str]	  =						ids_tensor([self.batch_size]       ,			self.type_sequence_label_size     )
							A__				:       List[str]	  =						ids_tensor([self.batch_size, self.seq_length]       ,			self.num_labels     )
							A__				:       List[str]	  =						ids_tensor([self.batch_size]       ,			2     ).float()
							A__				:       Optional[int]	  =						ids_tensor([self.batch_size]       ,			self.num_choices     )

						A__				:       Dict	  =						self.get_config()

						return (
						    config,
						    input_ids,
						    token_type_ids,
						    input_lengths,
						    sequence_labels,
						    token_labels,
						    is_impossible_labels,
						    choice_labels,
						    input_mask,
						)






					def     __A     (					self     ):
						return XLMConfig(
						    vocab_size=self.vocab_size       ,			n_special=self.n_special       ,			emb_dim=self.hidden_size       ,			n_layers=self.num_hidden_layers       ,			n_heads=self.num_attention_heads       ,			dropout=self.hidden_dropout_prob       ,			attention_dropout=self.attention_probs_dropout_prob       ,			gelu_activation=self.gelu_activation       ,			sinusoidal_embeddings=self.sinusoidal_embeddings       ,			asm=self.asm       ,			causal=self.causal       ,			n_langs=self.n_langs       ,			max_position_embeddings=self.max_position_embeddings       ,			initializer_range=self.initializer_range       ,			summary_type=self.summary_type       ,			use_proj=self.use_proj       ,			num_labels=self.num_labels       ,			bos_token_id=self.bos_token_id       ,			)






					def     __A     (					self       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			):
						A__				:       List[str]	  =						XLMModel(config=A__     )
						model.to(A__     )
						model.eval()
						A__				:       List[str]	  =						model(A__       ,			lengths=A__       ,			langs=A__     )
						A__				:       int	  =						model(A__       ,			langs=A__     )
						A__				:       List[Any]	  =						model(A__     )
						self.parent.assertEqual(result.last_hidden_state.shape       ,			(self.batch_size, self.seq_length, self.hidden_size)     )






					def     __A     (					self       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			):
						A__				:       List[str]	  =						XLMWithLMHeadModel(A__     )
						model.to(A__     )
						model.eval()

						A__				:       Dict	  =						model(A__       ,			token_type_ids=A__       ,			labels=A__     )
						self.parent.assertEqual(result.loss.shape       ,			()     )
						self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.seq_length, self.vocab_size)     )






					def     __A     (					self       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			):
						A__				:       Tuple	  =						XLMForQuestionAnsweringSimple(A__     )
						model.to(A__     )
						model.eval()

						A__				:       Dict	  =						model(A__     )

						A__				:       int	  =						model(A__       ,			start_positions=A__       ,			end_positions=A__     )
						A__				:       Dict	  =						outputs
						self.parent.assertEqual(result.start_logits.shape       ,			(self.batch_size, self.seq_length)     )
						self.parent.assertEqual(result.end_logits.shape       ,			(self.batch_size, self.seq_length)     )






					def     __A     (					self       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			):
						A__				:       Dict	  =						XLMForQuestionAnswering(A__     )
						model.to(A__     )
						model.eval()

						A__				:       Optional[int]	  =						model(A__     )

						A__				:       Optional[int]	  =						model(
						    A__       ,			start_positions=A__       ,			end_positions=A__       ,			cls_index=A__       ,			is_impossible=A__       ,			p_mask=A__       ,			)

						A__				:       Union[str, Any]	  =						model(
						    A__       ,			start_positions=A__       ,			end_positions=A__       ,			cls_index=A__       ,			is_impossible=A__       ,			)

						((A__)  ,  )				:       str	  =						result_with_labels.to_tuple()

						A__				:       int	  =						model(A__       ,			start_positions=A__       ,			end_positions=A__     )

						((A__)  ,  )				:       Any	  =						result_with_labels.to_tuple()

						self.parent.assertEqual(result_with_labels.loss.shape       ,			()     )
						self.parent.assertEqual(result.start_top_log_probs.shape       ,			(self.batch_size, model.config.start_n_top)     )
						self.parent.assertEqual(result.start_top_index.shape       ,			(self.batch_size, model.config.start_n_top)     )
						self.parent.assertEqual(
						    result.end_top_log_probs.shape       ,			(self.batch_size, model.config.start_n_top * model.config.end_n_top)     )
						self.parent.assertEqual(
						    result.end_top_index.shape       ,			(self.batch_size, model.config.start_n_top * model.config.end_n_top)     )
						self.parent.assertEqual(result.cls_logits.shape       ,			(self.batch_size,)     )






					def     __A     (					self       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			):
						A__				:       int	  =						XLMForSequenceClassification(A__     )
						model.to(A__     )
						model.eval()

						A__				:       Tuple	  =						model(A__     )
						A__				:       Optional[Any]	  =						model(A__       ,			labels=A__     )
						self.parent.assertEqual(result.loss.shape       ,			()     )
						self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.type_sequence_label_size)     )






					def     __A     (					self       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			):
						A__				:       List[str]	  =						self.num_labels
						A__				:       str	  =						XLMForTokenClassification(A__     )
						model.to(A__     )
						model.eval()

						A__				:       List[str]	  =						model(A__       ,			attention_mask=A__       ,			labels=A__     )
						self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.seq_length, self.num_labels)     )






					def     __A     (					self       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			):
						A__				:       Optional[Any]	  =						self.num_choices
						A__				:       Optional[Any]	  =						XLMForMultipleChoice(config=A__     )
						model.to(A__     )
						model.eval()
						A__				:       int	  =						input_ids.unsqueeze(1     ).expand(-1       ,			self.num_choices       ,			-1     ).contiguous()
						A__				:       List[Any]	  =						token_type_ids.unsqueeze(1     ).expand(-1       ,			self.num_choices       ,			-1     ).contiguous()
						A__				:       Tuple	  =						input_mask.unsqueeze(1     ).expand(-1       ,			self.num_choices       ,			-1     ).contiguous()
						A__				:       List[str]	  =						model(
						    A__       ,			attention_mask=A__       ,			token_type_ids=A__       ,			labels=A__       ,			)
						self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.num_choices)     )






					def     __A     (					self     ):
						A__				:       List[Any]	  =						self.prepare_config_and_inputs()
						(
						    (
						    A__
						)  ,  (
						    A__
						)  ,  (
						    A__
						)  ,  (
						    A__
						)  ,  (
						    A__
						)  ,  (
						    A__
						)  ,  (
						    A__
						)  ,  (
						    A__
						)  ,  (
						    A__
						)  ,  
						)				:       List[Any]	  =						config_and_inputs
						A__				:       Union[str, Any]	  =						{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """lengths""": input_lengths}
						return config, inputs_dict



@require_torch
class        _a		(__magic_name__						,  __magic_name__						,  __magic_name__						,  unittest.TestCase							):


					'''simple docstring'''






					UpperCAmelCase__:  Tuple    				=			(
					    (
					        XLMModel,
					        XLMWithLMHeadModel,
					        XLMForQuestionAnswering,
					        XLMForSequenceClassification,
					        XLMForQuestionAnsweringSimple,
					        XLMForTokenClassification,
					        XLMForMultipleChoice,
					    )
					    if is_torch_available()
					    else ()
					)
					UpperCAmelCase__:  int    				=			(
					    (XLMWithLMHeadModel,) if is_torch_available() else ()
					)  # TODO (PVP): Check other models whether language generation is also applicable
					UpperCAmelCase__:  Optional[Any]    				=			(
					    {
					        '''feature-extraction''': XLMModel,
					        '''fill-mask''': XLMWithLMHeadModel,
					        '''question-answering''': XLMForQuestionAnsweringSimple,
					        '''text-classification''': XLMForSequenceClassification,
					        '''text-generation''': XLMWithLMHeadModel,
					        '''token-classification''': XLMForTokenClassification,
					        '''zero-shot''': XLMForSequenceClassification,
					    }
					    if is_torch_available()
					    else {}
					)






					def     __A     (					self       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__     ):
						if (
						    pipeline_test_casse_name == "QAPipelineTests"
						    and tokenizer_name is not None
						    and not tokenizer_name.endswith("""Fast"""     )
						):
							# `QAPipelineTests` fails for a few models when the slower tokenizer are used.
							# (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
							# TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
							return True

						return False






					def     __A     (					self       ,			A__       ,			A__       ,			A__=False     ):
						A__				:       Dict	  =						super()._prepare_for_class(A__       ,			A__       ,			return_labels=A__     )

						if return_labels:
							if model_class.__name__ == "XLMForQuestionAnswering":
								A__				:       Tuple	  =						torch.zeros(
								    self.model_tester.batch_size       ,			dtype=torch.long       ,			device=A__     )
								A__				:       Optional[int]	  =						torch.zeros(
								    self.model_tester.batch_size       ,			dtype=torch.long       ,			device=A__     )

						return inputs_dict






					def     __A     (					self     ):
						A__				:       str	  =						XLMModelTester(self     )
						A__				:       Tuple	  =						ConfigTester(self       ,			config_class=A__       ,			emb_dim=37     )






					def     __A     (					self     ):
						self.config_tester.run_common_tests()






					def     __A     (					self     ):
						A__				:       int	  =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_model(*A__     )






					def     __A     (					self     ):
						A__				:       Dict	  =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_lm_head(*A__     )






					def     __A     (					self     ):
						A__				:       List[Any]	  =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_simple_qa(*A__     )






					def     __A     (					self     ):
						A__				:       Tuple	  =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_qa(*A__     )






					def     __A     (					self     ):
						A__				:       List[Any]	  =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_sequence_classif(*A__     )






					def     __A     (					self     ):
						A__				:       Optional[int]	  =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_token_classif(*A__     )






					def     __A     (					self     ):
						A__				:       Tuple	  =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_for_multiple_choice(*A__     )






					def     __A     (					self       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__=False       ,			A__=1     ):
						self.assertIsInstance(A__       ,			A__     )
						self.assertListEqual(
						    [isinstance(A__       ,			A__     ) for iter_attentions in attentions]       ,			[True] * len(A__     )     )
						self.assertEqual(len(A__     )       ,			(max_length - min_length) * num_beam_groups     )

						for idx, iter_attentions in enumerate(A__     ):
							# adds PAD dummy token
							A__				:       List[Any]	  =						min_length + idx + 1
							A__				:       Tuple	  =						min_length + idx + 1

							A__				:       List[str]	  =						(
							    batch_size * num_beam_groups,
							    config.num_attention_heads,
							    tgt_len,
							    src_len,
							)
							# check attn size
							self.assertListEqual(
							    [layer_attention.shape for layer_attention in iter_attentions]       ,			[expected_shape] * len(A__     )     )






					def     __A     (					self       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__=False       ,			A__=1     ):
						self.assertIsInstance(A__       ,			A__     )
						self.assertListEqual(
						    [isinstance(A__       ,			A__     ) for iter_hidden_states in hidden_states]       ,			[True] * len(A__     )       ,			)
						self.assertEqual(len(A__     )       ,			(max_length - min_length) * num_beam_groups     )

						for idx, iter_hidden_states in enumerate(A__     ):
							# adds PAD dummy token
							A__				:       Tuple	  =						min_length + idx + 1
							A__				:       int	  =						(batch_size * num_beam_groups, seq_len, config.hidden_size)
							# check hidden size
							self.assertListEqual(
							    [layer_hidden_states.shape for layer_hidden_states in iter_hidden_states]       ,			[expected_shape] * len(A__     )       ,			)
						pass






					@slow
					def     __A     (					self     ):
						for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							A__				:       Union[str, Any]	  =						XLMModel.from_pretrained(A__     )
							self.assertIsNotNone(A__     )



@require_torch
class        _a		(unittest.TestCase							):


					'''simple docstring'''






					@slow
					def     __A     (					self     ):
						A__				:       Tuple	  =						XLMWithLMHeadModel.from_pretrained("""xlm-mlm-en-2048"""     )
						model.to(A__     )
						A__				:       Tuple	  =						torch.tensor([[14, 447]]       ,			dtype=torch.long       ,			device=A__     )  # the president
						A__				:       Optional[Any]	  =						[
						    14,
						    447,
						    14,
						    447,
						    14,
						    447,
						    14,
						    447,
						    14,
						    447,
						    14,
						    447,
						    14,
						    447,
						    14,
						    447,
						    14,
						    447,
						    14,
						    447,
						]  # the president the president the president the president the president the president the president the president the president the president
						# TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
						A__				:       Tuple	  =						model.generate(A__       ,			do_sample=A__     )
						self.assertListEqual(output_ids[0].cpu().numpy().tolist()       ,			A__     )






import tensorflow as tf

from ...tf_utils import shape_list



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


					'''simple docstring'''






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

						A__				:       Tuple	  =						vocab_size
						A__				:       List[str]	  =						d_embed
						A__				:       int	  =						d_proj

						A__				:       Tuple	  =						cutoffs + [vocab_size]
						A__				:       Any	  =						[0] + self.cutoffs
						A__				:       str	  =						div_val

						A__				:       int	  =						self.cutoffs[0]
						A__				:       List[str]	  =						len(self.cutoffs     ) - 1
						A__				:       Optional[int]	  =						self.shortlist_size + self.n_clusters
						A__				:       List[Any]	  =						keep_order

						A__				:       Optional[int]	  =						[]
						A__				:       Optional[Any]	  =						[]






					def     __A     (					self       ,			A__     ):
						if self.n_clusters > 0:
							A__				:       List[str]	  =						self.add_weight(
							    shape=(self.n_clusters, self.d_embed)       ,			initializer="""zeros"""       ,			trainable=A__       ,			name="""cluster_weight"""     )
							A__				:       Union[str, Any]	  =						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:
									A__				:       List[str]	  =						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__     )
								A__				:       int	  =						self.add_weight(
								    shape=(self.vocab_size, self.d_embed)       ,			initializer="""zeros"""       ,			trainable=A__       ,			name=F"""out_layers_._{i}_._weight"""       ,			)
								A__				:       str	  =						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     )     ):
								A__  ,  A__				:       Optional[int]	  =						self.cutoff_ends[i], self.cutoff_ends[i + 1]
								A__				:       Optional[int]	  =						self.d_embed // (self.div_val**i)

								A__				:       Optional[Any]	  =						self.add_weight(
								    shape=(d_emb_i, self.d_proj)       ,			initializer="""zeros"""       ,			trainable=A__       ,			name=F"""out_projs_._{i}"""     )
								self.out_projs.append(A__     )
								A__				:       List[str]	  =						self.add_weight(
								    shape=(r_idx - l_idx, d_emb_i)       ,			initializer="""zeros"""       ,			trainable=A__       ,			name=F"""out_layers_._{i}_._weight"""       ,			)
								A__				:       Optional[int]	  =						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     ):
						A__				:       Tuple	  =						x
						if proj is not None:
							A__				:       Tuple	  =						tf.einsum("""ibd,ed->ibe"""       ,			A__       ,			A__     )
						return tf.einsum("""ibd,nd->ibn"""       ,			A__       ,			A__     ) + b






					@staticmethod
					def     __A     (					A__       ,			A__     ):
						A__				:       Optional[Any]	  =						shape_list(A__     )
						A__				:       Optional[Any]	  =						tf.range(lp_size[0]       ,			dtype=target.dtype     )
						A__				:       List[str]	  =						tf.stack([r, target]       ,			1     )
						return tf.gather_nd(A__       ,			A__     )






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

								if self.div_val == 1:
									A__				:       Any	  =						self.out_layers[0][0][l_idx:r_idx]
									A__				:       List[Any]	  =						self.out_layers[0][1][l_idx:r_idx]
								else:
									A__				:       Optional[int]	  =						self.out_layers[i][0]
									A__				:       Tuple	  =						self.out_layers[i][1]

								if i == 0:
									A__				:       Union[str, Any]	  =						tf.concat([cur_W, self.cluster_weight]       ,			0     )
									A__				:       Optional[int]	  =						tf.concat([cur_b, self.cluster_bias]       ,			0     )

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

						if target is not None:
							if return_mean:
								A__				:       List[Any]	  =						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



