






"""simple docstring"""





import inspect
import unittest

from transformers import MobileViTVaConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation, MobileViTVaModel
							from transformers.models.mobilevitva.modeling_mobilevitva import (
							    MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST,
							    make_divisible,
							)


if is_vision_available():
							from PIL import Image

							from transformers import MobileViTImageProcessor



class       __A		(   SCREAMING_SNAKE_CASE_			):

					def   lowerCamelCase__   (			self  :      Optional[int]					)			->							Dict:
											__magic_name__:	int							=       self.config_class(**self.inputs_dict					)
											self.parent.assertTrue(hasattr(__snake_case     ,		"""width_multiplier"""					)					)



class       __A		:

					def __init__(			self  :      Tuple     ,		__snake_case  :      List[str]     ,		__snake_case  :      int=1_3     ,		__snake_case  :      Optional[int]=6_4     ,		__snake_case  :      Optional[int]=2     ,		__snake_case  :      Tuple=3     ,		__snake_case  :      Any="swish"     ,		__snake_case  :      int=3     ,		__snake_case  :      int=3_2     ,		__snake_case  :      List[str]=0.1     ,		__snake_case  :      int=0.02     ,		__snake_case  :      Union[str, Any]=True     ,		__snake_case  :      List[Any]=True     ,		__snake_case  :      Optional[Any]=1_0     ,		__snake_case  :      Tuple=None     ,		__snake_case  :      Any=0.25     ,		__snake_case  :      List[Any]=0.0     ,		__snake_case  :      Dict=0.0     ,		)			->							Optional[int]:
											__magic_name__:	int							=       parent
											__magic_name__:	Union[str, Any]							=       batch_size
											__magic_name__:	Optional[Any]							=       image_size
											__magic_name__:	Optional[Any]							=       patch_size
											__magic_name__:	Union[str, Any]							=       num_channels
											__magic_name__:	Optional[Any]							=       make_divisible(5_1_2 * width_multiplier     ,		divisor=8					)
											__magic_name__:	List[Any]							=       hidden_act
											__magic_name__:	int							=       conv_kernel_size
											__magic_name__:	Tuple							=       output_stride
											__magic_name__:	Dict							=       classifier_dropout_prob
											__magic_name__:	Optional[Any]							=       use_labels
											__magic_name__:	List[str]							=       is_training
											__magic_name__:	List[str]							=       num_labels
											__magic_name__:	str							=       initializer_range
											__magic_name__:	Dict							=       scope
											__magic_name__:	str							=       width_multiplier
											__magic_name__:	str							=       ffn_dropout
											__magic_name__:	Optional[Any]							=       attn_dropout

					def   lowerCamelCase__   (			self  :      Any					)			->							List[str]:
											__magic_name__:	Dict							=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

											__magic_name__:	List[str]							=       None
											__magic_name__:	Tuple							=       None
											if self.use_labels:
																	__magic_name__:	List[str]							=       ids_tensor([self.batch_size]     ,		self.num_labels					)
																	__magic_name__:	Any							=       ids_tensor([self.batch_size, self.image_size, self.image_size]     ,		self.num_labels					)

											__magic_name__:	Dict							=       self.get_config()

											return config, pixel_values, labels, pixel_labels

					def   lowerCamelCase__   (			self  :      Union[str, Any]					)			->							Any:
											return MobileViTVaConfig(
											    image_size=self.image_size     ,		patch_size=self.patch_size     ,		num_channels=self.num_channels     ,		hidden_act=self.hidden_act     ,		conv_kernel_size=self.conv_kernel_size     ,		output_stride=self.output_stride     ,		classifier_dropout_prob=self.classifier_dropout_prob     ,		initializer_range=self.initializer_range     ,		width_multiplier=self.width_multiplier     ,		ffn_dropout=self.ffn_dropout_prob     ,		attn_dropout=self.attn_dropout_prob     ,		)

					def   lowerCamelCase__   (			self  :      List[str]     ,		__snake_case  :      Dict     ,		__snake_case  :      List[Any]     ,		__snake_case  :      List[str]     ,		__snake_case  :      str					)			->							Optional[Any]:
											__magic_name__:	Optional[int]							=       MobileViTVaModel(config=__snake_case					)
											model.to(__snake_case					)
											model.eval()
											__magic_name__:	Optional[int]							=       model(__snake_case					)
											self.parent.assertEqual(
											    result.last_hidden_state.shape     ,		(
											        self.batch_size,
											        self.last_hidden_size,
											        self.image_size // self.output_stride,
											        self.image_size // self.output_stride,
											    )     ,		)

					def   lowerCamelCase__   (			self  :      Union[str, Any]     ,		__snake_case  :      Optional[int]     ,		__snake_case  :      Optional[Any]     ,		__snake_case  :      Any     ,		__snake_case  :      List[Any]					)			->							Any:
											__magic_name__:	Union[str, Any]							=       self.num_labels
											__magic_name__:	List[str]							=       MobileViTVaForImageClassification(__snake_case					)
											model.to(__snake_case					)
											model.eval()
											__magic_name__:	Optional[int]							=       model(__snake_case     ,		labels=__snake_case					)
											self.parent.assertEqual(result.logits.shape     ,		(self.batch_size, self.num_labels)					)

					def   lowerCamelCase__   (			self  :      List[Any]     ,		__snake_case  :      Dict     ,		__snake_case  :      Optional[Any]     ,		__snake_case  :      str     ,		__snake_case  :      Union[str, Any]					)			->							Dict:
											__magic_name__:	Tuple							=       self.num_labels
											__magic_name__:	Any							=       MobileViTVaForSemanticSegmentation(__snake_case					)
											model.to(__snake_case					)
											model.eval()
											__magic_name__:	List[Any]							=       model(__snake_case					)
											self.parent.assertEqual(
											    result.logits.shape     ,		(
											        self.batch_size,
											        self.num_labels,
											        self.image_size // self.output_stride,
											        self.image_size // self.output_stride,
											    )     ,		)
											__magic_name__:	Any							=       model(__snake_case     ,		labels=__snake_case					)
											self.parent.assertEqual(
											    result.logits.shape     ,		(
											        self.batch_size,
											        self.num_labels,
											        self.image_size // self.output_stride,
											        self.image_size // self.output_stride,
											    )     ,		)






					def   lowerCamelCase__   (			self  :      Any					)			->							Union[str, Any]:
											__magic_name__:	str							=       self.prepare_config_and_inputs()
											__magic_name__,			__magic_name__,			__magic_name__,			__magic_name__:	Optional[int]							=       config_and_inputs
											__magic_name__:	Tuple							=       {"""pixel_values""": pixel_values}
											return config, inputs_dict



@require_torch
class       __A		(   SCREAMING_SNAKE_CASE_     ,SCREAMING_SNAKE_CASE_     ,unittest.TestCase			):
					UpperCAmelCase__								=				(
					    (MobileViTVaModel, MobileViTVaForImageClassification, MobileViTVaForSemanticSegmentation)
					    if is_torch_available()
					    else ()
					)

					UpperCAmelCase__								=				(
					    {
					        "feature-extraction": MobileViTVaModel,
					        "image-classification": MobileViTVaForImageClassification,
					        "image-segmentation": MobileViTVaForSemanticSegmentation,
					    }
					    if is_torch_available()
					    else {}
					)

					UpperCAmelCase__								=				False
					UpperCAmelCase__								=				False
					UpperCAmelCase__								=				False
					UpperCAmelCase__								=				False

					def   lowerCamelCase__   (			self  :      List[Any]					)			->							Any:
											__magic_name__:	List[str]							=       MobileViTVaModelTester(self					)
											__magic_name__:	str							=       MobileViTVaConfigTester(self     ,		config_class=__snake_case     ,		has_text_modality=__snake_case					)

					def   lowerCamelCase__   (			self  :      int					)			->							Dict:
											self.config_tester.run_common_tests()

					@unittest.skip(reason="""MobileViTV2 does not use inputs_embeds"""					)
					def   lowerCamelCase__   (			self  :      Optional[Any]					)			->							List[str]:
											pass

					@unittest.skip(reason="""MobileViTV2 does not support input and output embeddings"""					)
					def   lowerCamelCase__   (			self  :      List[str]					)			->							Dict:
											pass

					@unittest.skip(reason="""MobileViTV2 does not output attentions"""					)
					def   lowerCamelCase__   (			self  :      Optional[Any]					)			->							List[Any]:
											pass

					@require_torch_multi_gpu
					@unittest.skip(reason="""Got `CUDA error: misaligned address` for tests after this one being run."""					)
					def   lowerCamelCase__   (			self  :      str					)			->							str:
											pass

					@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""					)
					def   lowerCamelCase__   (			self  :      str					)			->							Dict:
											pass

					def   lowerCamelCase__   (			self  :      Dict					)			->							List[str]:
											__magic_name__,			__magic_name__:	Dict							=       self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	__magic_name__:	str							=       model_class(__snake_case					)
																	__magic_name__:	Dict							=       inspect.signature(model.forward					)
																	# signature.parameters is an OrderedDict => so arg_names order is deterministic
																	__magic_name__:	int							=       [*signature.parameters.keys()]

																	__magic_name__:	Dict							=       ["""pixel_values"""]
																	self.assertListEqual(arg_names[:1]     ,		__snake_case					)

					def   lowerCamelCase__   (			self  :      str					)			->							Optional[Any]:
											__magic_name__:	Any							=       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__snake_case					)

					def   lowerCamelCase__   (			self  :      Tuple					)			->							Optional[Any]:
											def check_hidden_states_output(__snake_case  :      Union[str, Any]     ,		__snake_case  :      str     ,		__snake_case  :      List[Any]					):
																	__magic_name__:	Union[str, Any]							=       model_class(__snake_case					)
																	model.to(__snake_case					)
																	model.eval()

																	with torch.no_grad():
																							__magic_name__:	str							=       model(**self._prepare_for_class(__snake_case     ,		__snake_case					)					)

																	__magic_name__:	List[Any]							=       outputs.hidden_states

																	__magic_name__:	List[Any]							=       5
																	self.assertEqual(len(__snake_case					)     ,		__snake_case					)

																	# MobileViTV2's feature maps are of shape (batch_size, num_channels, height, width)
																	# with the width and height being successively divided by 2.
																	__magic_name__:	int							=       2
																	for i in range(len(__snake_case					)					):
																							self.assertListEqual(
																							    list(hidden_states[i].shape[-2:]					)     ,		[self.model_tester.image_size // divisor, self.model_tester.image_size // divisor]     ,		)
																							divisor *= 2

																	self.assertEqual(self.model_tester.output_stride     ,		divisor // 2					)

											__magic_name__,			__magic_name__:	Dict							=       self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	__magic_name__:	Dict							=       True
																	check_hidden_states_output(__snake_case     ,		__snake_case     ,		__snake_case					)

																	# check that output_hidden_states also work using config
																	del inputs_dict["output_hidden_states"]
																	__magic_name__:	Optional[int]							=       True

																	check_hidden_states_output(__snake_case     ,		__snake_case     ,		__snake_case					)

					def   lowerCamelCase__   (			self  :      Tuple					)			->							Union[str, Any]:
											__magic_name__:	Any							=       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*__snake_case					)

					def   lowerCamelCase__   (			self  :      Tuple					)			->							str:
											__magic_name__:	Dict							=       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_semantic_segmentation(*__snake_case					)






					@slow
					def   lowerCamelCase__   (			self  :      Union[str, Any]					)			->							Optional[int]:
											for model_name in MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	__magic_name__:	Union[str, Any]							=       MobileViTVaModel.from_pretrained(__snake_case					)
																	self.assertIsNotNone(__snake_case					)



def   a    (     )    ->							Any:
						__magic_name__:	Any							=       Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""  )
						return image



@require_torch
@require_vision
class       __A		(   unittest.TestCase			):

					@cached_property
					def   lowerCamelCase__   (			self  :      Dict					)			->							str:
											return (
											    MobileViTImageProcessor.from_pretrained("""apple/mobilevitv2-1.0-imagenet1k-256"""					)
											    if is_vision_available()
											    else None
											)

					@slow
					def   lowerCamelCase__   (			self  :      Dict					)			->							int:
											__magic_name__:	Any							=       MobileViTVaForImageClassification.from_pretrained("""apple/mobilevitv2-1.0-imagenet1k-256"""					).to(
											    __snake_case					)

											__magic_name__:	Optional[Any]							=       self.default_image_processor
											__magic_name__:	Tuple							=       prepare_img()
											__magic_name__:	str							=       image_processor(images=__snake_case     ,		return_tensors="""pt"""					).to(__snake_case					)

											# forward pass
											with torch.no_grad():
																	__magic_name__:	Any							=       model(**__snake_case					)

											# verify the logits
											__magic_name__:	List[Any]							=       torch.Size((1, 1_0_0_0)					)
											self.assertEqual(outputs.logits.shape     ,		__snake_case					)

											__magic_name__:	Optional[int]							=       torch.tensor([-1.63_36E00, -7.32_04E-02, -5.18_83E-01]					).to(__snake_case					)

											self.assertTrue(torch.allclose(outputs.logits[0, :3]     ,		__snake_case     ,		atol=1E-4					)					)

					@slow
					def   lowerCamelCase__   (			self  :      List[Any]					)			->							Optional[Any]:
											__magic_name__:	Any							=       MobileViTVaForSemanticSegmentation.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""					)
											__magic_name__:	List[Any]							=       model.to(__snake_case					)

											__magic_name__:	str							=       MobileViTImageProcessor.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""					)

											__magic_name__:	str							=       prepare_img()
											__magic_name__:	str							=       image_processor(images=__snake_case     ,		return_tensors="""pt"""					).to(__snake_case					)

											# forward pass
											with torch.no_grad():
																	__magic_name__:	Any							=       model(**__snake_case					)
											__magic_name__:	Optional[Any]							=       outputs.logits

											# verify the logits
											__magic_name__:	List[str]							=       torch.Size((1, 2_1, 3_2, 3_2)					)
											self.assertEqual(logits.shape     ,		__snake_case					)

											__magic_name__:	Optional[Any]							=       torch.tensor(
											    [
											        [[7.0863, 7.1525, 6.8201], [6.6931, 6.8770, 6.8933], [6.2978, 7.0366, 6.9636]],
											        [[-3.7134, -3.6712, -3.6675], [-3.5825, -3.3549, -3.4777], [-3.3435, -3.3979, -3.2857]],
											        [[-2.9329, -2.8003, -2.7369], [-3.0564, -2.4780, -2.0207], [-2.6889, -1.9298, -1.7640]],
											    ]     ,		device=__snake_case     ,		)

											self.assertTrue(torch.allclose(logits[0, :3, :3, :3]     ,		__snake_case     ,		atol=1E-4					)					)






					@slow
					def   lowerCamelCase__   (			self  :      int					)			->							Optional[Any]:
											__magic_name__:	List[str]							=       MobileViTVaForSemanticSegmentation.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""					)
											__magic_name__:	Union[str, Any]							=       model.to(__snake_case					)

											__magic_name__:	List[Any]							=       MobileViTImageProcessor.from_pretrained("""shehan97/mobilevitv2-1.0-voc-deeplabv3"""					)

											__magic_name__:	str							=       prepare_img()
											__magic_name__:	List[str]							=       image_processor(images=__snake_case     ,		return_tensors="""pt"""					).to(__snake_case					)

											# forward pass
											with torch.no_grad():
																	__magic_name__:	Tuple							=       model(**__snake_case					)

											__magic_name__:	Union[str, Any]							=       outputs.logits.detach().cpu()

											__magic_name__:	Any							=       image_processor.post_process_semantic_segmentation(outputs=__snake_case     ,		target_sizes=[(5_0, 6_0)]					)
											__magic_name__:	Tuple							=       torch.Size((5_0, 6_0)					)
											self.assertEqual(segmentation[0].shape     ,		__snake_case					)

											__magic_name__:	str							=       image_processor.post_process_semantic_segmentation(outputs=__snake_case					)
											__magic_name__:	Optional[int]							=       torch.Size((3_2, 3_2)					)
											self.assertEqual(segmentation[0].shape     ,		__snake_case					)










import torch
from torch import nn



class 					__snake_case	(  nn.Module    ):


    '''simple docstring'''







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





        '''simple docstring'''

        super().__init__()

        SCREAMING_SNAKE_CASE__											=			n_token
        SCREAMING_SNAKE_CASE__											=			d_embed
        SCREAMING_SNAKE_CASE__											=			d_proj

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

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

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

        SCREAMING_SNAKE_CASE__											=			nn.ModuleList()
        SCREAMING_SNAKE_CASE__											=			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(A_   ,				A_     )     )     )
                else:
                    self.out_projs.append(A_     )

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

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

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

        SCREAMING_SNAKE_CASE__											=			keep_order



    def  lowercase_					(       self   ,				A_   ,				A_   ,				A_   ,				A_     ):





        '''simple docstring'''

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

        return logit



    def  lowercase_					(       self   ,				A_   ,				A_=None   ,				A_=False     ):





        '''simple docstring'''

        if labels is not None:
            # Shift so that tokens < n predict n
            SCREAMING_SNAKE_CASE__											=			hidden[..., :-1, :].contiguous()
            SCREAMING_SNAKE_CASE__											=			labels[..., 1:].contiguous()
            SCREAMING_SNAKE_CASE__											=			hidden.view(-1   ,				hidden.size(-1     )     )
            SCREAMING_SNAKE_CASE__											=			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:
            SCREAMING_SNAKE_CASE__											=			hidden.view(-1   ,				hidden.size(-1     )     )

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

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

                weights.append(A_     )
                biases.append(A_     )

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

            SCREAMING_SNAKE_CASE__											=			self._compute_logit(A_   ,				A_   ,				A_   ,				A_     )
            SCREAMING_SNAKE_CASE__											=			nn.functional.log_softmax(A_   ,				dim=1     )

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

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

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

                    if indices_i.numel() == 0:
                        continue

                    SCREAMING_SNAKE_CASE__											=			labels.index_select(0   ,				A_     ) - l_idx
                    SCREAMING_SNAKE_CASE__											=			head_logprob.index_select(0   ,				A_     )
                    SCREAMING_SNAKE_CASE__											=			hidden.index_select(0   ,				A_     )
                else:
                    SCREAMING_SNAKE_CASE__											=			hidden

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

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

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

        return out




    def  lowercase_					(       self   ,				A_     ):





        '''simple docstring'''

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

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

                weights.append(A_     )
                biases.append(A_     )

            SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__											=			weights[0], biases[0], self.out_projs[0]
            SCREAMING_SNAKE_CASE__											=			self._compute_logit(A_   ,				A_   ,				A_   ,				A_     )

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

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

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

                    SCREAMING_SNAKE_CASE__											=			self._compute_logit(A_   ,				A_   ,				A_   ,				A_     )
                    SCREAMING_SNAKE_CASE__											=			nn.functional.log_softmax(A_   ,				dim=1     )

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

            return out
