





import inspect
import unittest

from transformers import YolosConfig
from transformers.testing_utils import require_torch, 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
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch
  from torch import nn

  from transformers import YolosForObjectDetection, YolosModel
  from transformers.models.yolos.modeling_yolos import YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import AutoImageProcessor



class 					__snake_case	:


    '''simple docstring'''







    def __init__(       self   ,				A_   ,				A_=13   ,				A_=[30, 30]   ,				A_=2   ,				A_=3   ,				A_=True   ,				A_=True   ,				A_=32   ,				A_=5   ,				A_=4   ,				A_=37   ,				A_="gelu"   ,				A_=0.1   ,				A_=0.1   ,				A_=10   ,				A_=0.02   ,				A_=3   ,				A_=None   ,				A_=8   ,				A_=10   ,				):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			parent
        SCREAMING_SNAKE_CASE__											=			batch_size
        SCREAMING_SNAKE_CASE__											=			image_size
        SCREAMING_SNAKE_CASE__											=			patch_size
        SCREAMING_SNAKE_CASE__											=			num_channels
        SCREAMING_SNAKE_CASE__											=			is_training
        SCREAMING_SNAKE_CASE__											=			use_labels
        SCREAMING_SNAKE_CASE__											=			hidden_size
        SCREAMING_SNAKE_CASE__											=			num_hidden_layers
        SCREAMING_SNAKE_CASE__											=			num_attention_heads
        SCREAMING_SNAKE_CASE__											=			intermediate_size
        SCREAMING_SNAKE_CASE__											=			hidden_act
        SCREAMING_SNAKE_CASE__											=			hidden_dropout_prob
        SCREAMING_SNAKE_CASE__											=			attention_probs_dropout_prob
        SCREAMING_SNAKE_CASE__											=			type_sequence_label_size
        SCREAMING_SNAKE_CASE__											=			initializer_range
        SCREAMING_SNAKE_CASE__											=			num_labels
        SCREAMING_SNAKE_CASE__											=			scope
        SCREAMING_SNAKE_CASE__											=			n_targets
        SCREAMING_SNAKE_CASE__											=			num_detection_tokens
        # we set the expected sequence length (which is used in several tests)
        # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token) + num_detection_tokens
        SCREAMING_SNAKE_CASE__											=			(image_size[1] // patch_size) * (image_size[0] // patch_size)
        SCREAMING_SNAKE_CASE__											=			num_patches + 1 + self.num_detection_tokens



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			floats_tensor([self.batch_size, self.num_channels, self.image_size[0], self.image_size[1]]     )

        SCREAMING_SNAKE_CASE__											=			None
        if self.use_labels:
            # labels is a list of Dict (each Dict being the labels for a given example in the batch)
            SCREAMING_SNAKE_CASE__											=			[]
            for i in range(self.batch_size     ):
                SCREAMING_SNAKE_CASE__											=			{}
                SCREAMING_SNAKE_CASE__											=			torch.randint(
                    high=self.num_labels   ,				size=(self.n_targets,)   ,				device=A_     )
                SCREAMING_SNAKE_CASE__											=			torch.rand(self.n_targets   ,				4   ,				device=A_     )
                labels.append(A_     )

        SCREAMING_SNAKE_CASE__											=			self.get_config()

        return config, pixel_values, labels



    def  lowercase_					(       self     ):





        '''simple docstring'''

        return YolosConfig(
            image_size=self.image_size   ,				patch_size=self.patch_size   ,				num_channels=self.num_channels   ,				hidden_size=self.hidden_size   ,				num_hidden_layers=self.num_hidden_layers   ,				num_attention_heads=self.num_attention_heads   ,				intermediate_size=self.intermediate_size   ,				hidden_act=self.hidden_act   ,				hidden_dropout_prob=self.hidden_dropout_prob   ,				attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,				is_decoder=A_   ,				initializer_range=self.initializer_range   ,				num_detection_tokens=self.num_detection_tokens   ,				num_labels=self.num_labels   ,				)



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





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			YolosModel(config=A_     )
        model.to(A_     )
        model.eval()
        SCREAMING_SNAKE_CASE__											=			model(A_     )
        self.parent.assertEqual(
            result.last_hidden_state.shape   ,				(self.batch_size, self.expected_seq_len, self.hidden_size)     )



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





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			YolosForObjectDetection(A_     )
        model.to(A_     )
        model.eval()

        SCREAMING_SNAKE_CASE__											=			model(pixel_values=A_     )
        SCREAMING_SNAKE_CASE__											=			model(A_     )

        self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.num_detection_tokens, self.num_labels + 1)     )
        self.parent.assertEqual(result.pred_boxes.shape   ,				(self.batch_size, self.num_detection_tokens, 4)     )

        SCREAMING_SNAKE_CASE__											=			model(pixel_values=A_   ,				labels=A_     )

        self.parent.assertEqual(result.loss.shape   ,				()     )
        self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.num_detection_tokens, self.num_labels + 1)     )
        self.parent.assertEqual(result.pred_boxes.shape   ,				(self.batch_size, self.num_detection_tokens, 4)     )




    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.prepare_config_and_inputs()
        SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__											=			config_and_inputs
        SCREAMING_SNAKE_CASE__											=			{'''pixel_values''': pixel_values}
        return config, inputs_dict





@require_torch
class 					__snake_case	(  __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     unittest.TestCase    ):


    '''simple docstring'''







    lowerCamelCase__    :		Tuple				=	(YolosModel, YolosForObjectDetection) if is_torch_available() else ()
    lowerCamelCase__    :		int				=	(
        {"""feature-extraction""": YolosModel, """object-detection""": YolosForObjectDetection} if is_torch_available() else {}
    )

    lowerCamelCase__    :		Optional[Any]				=	False
    lowerCamelCase__    :		Union[str, Any]				=	False
    lowerCamelCase__    :		str				=	False
    lowerCamelCase__    :		Optional[Any]				=	False



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





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			super()._prepare_for_class(A_   ,				A_   ,				return_labels=A_     )

        if return_labels:
            if model_class.__name__ == "YolosForObjectDetection":
                SCREAMING_SNAKE_CASE__											=			[]
                for i in range(self.model_tester.batch_size     ):
                    SCREAMING_SNAKE_CASE__											=			{}
                    SCREAMING_SNAKE_CASE__											=			torch.ones(
                        size=(self.model_tester.n_targets,)   ,				device=A_   ,				dtype=torch.long     )
                    SCREAMING_SNAKE_CASE__											=			torch.ones(
                        self.model_tester.n_targets   ,				4   ,				device=A_   ,				dtype=torch.float     )
                    labels.append(A_     )
                SCREAMING_SNAKE_CASE__											=			labels

        return inputs_dict



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			YolosModelTester(self     )
        SCREAMING_SNAKE_CASE__											=			ConfigTester(self   ,				config_class=A_   ,				has_text_modality=A_   ,				hidden_size=37     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        self.config_tester.run_common_tests()



    def  lowercase_					(       self     ):





        '''simple docstring'''

        pass



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            SCREAMING_SNAKE_CASE__											=			model_class(A_     )
            self.assertIsInstance(model.get_input_embeddings()   ,				(nn.Module)     )
            SCREAMING_SNAKE_CASE__											=			model.get_output_embeddings()
            self.assertTrue(x is None or isinstance(A_   ,				nn.Linear     )     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs_for_common()

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

            SCREAMING_SNAKE_CASE__											=			['''pixel_values''']
            self.assertListEqual(arg_names[:1]   ,				A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs_for_common()
        SCREAMING_SNAKE_CASE__											=			True

        # in YOLOS, the seq_len is different
        SCREAMING_SNAKE_CASE__											=			self.model_tester.expected_seq_len
        for model_class in self.all_model_classes:
            SCREAMING_SNAKE_CASE__											=			True
            SCREAMING_SNAKE_CASE__											=			False
            SCREAMING_SNAKE_CASE__											=			True
            SCREAMING_SNAKE_CASE__											=			model_class(A_     )
            model.to(A_     )
            model.eval()
            with torch.no_grad():
                SCREAMING_SNAKE_CASE__											=			model(**self._prepare_for_class(A_   ,				A_     )     )
            SCREAMING_SNAKE_CASE__											=			outputs.attentions
            self.assertEqual(len(A_     )   ,				self.model_tester.num_hidden_layers     )

            # check that output_attentions also work using config
            del inputs_dict["output_attentions"]
            SCREAMING_SNAKE_CASE__											=			True
            SCREAMING_SNAKE_CASE__											=			model_class(A_     )
            model.to(A_     )
            model.eval()
            with torch.no_grad():
                SCREAMING_SNAKE_CASE__											=			model(**self._prepare_for_class(A_   ,				A_     )     )
            SCREAMING_SNAKE_CASE__											=			outputs.attentions
            self.assertEqual(len(A_     )   ,				self.model_tester.num_hidden_layers     )

            self.assertListEqual(
                list(attentions[0].shape[-3:]     )   ,				[self.model_tester.num_attention_heads, seq_len, seq_len]   ,				)
            SCREAMING_SNAKE_CASE__											=			len(A_     )

            # Check attention is always last and order is fine
            SCREAMING_SNAKE_CASE__											=			True
            SCREAMING_SNAKE_CASE__											=			True
            SCREAMING_SNAKE_CASE__											=			model_class(A_     )
            model.to(A_     )
            model.eval()
            with torch.no_grad():
                SCREAMING_SNAKE_CASE__											=			model(**self._prepare_for_class(A_   ,				A_     )     )

            SCREAMING_SNAKE_CASE__											=			1
            self.assertEqual(out_len + added_hidden_states   ,				len(A_     )     )

            SCREAMING_SNAKE_CASE__											=			outputs.attentions

            self.assertEqual(len(A_     )   ,				self.model_tester.num_hidden_layers     )
            self.assertListEqual(
                list(self_attentions[0].shape[-3:]     )   ,				[self.model_tester.num_attention_heads, seq_len, seq_len]   ,				)



    def  lowercase_					(       self     ):





        '''simple docstring'''

        def check_hidden_states_output(A_   ,				A_   ,				A_     ):
            SCREAMING_SNAKE_CASE__											=			model_class(A_     )
            model.to(A_     )
            model.eval()

            with torch.no_grad():
                SCREAMING_SNAKE_CASE__											=			model(**self._prepare_for_class(A_   ,				A_     )     )

            SCREAMING_SNAKE_CASE__											=			outputs.hidden_states

            SCREAMING_SNAKE_CASE__											=			getattr(
                self.model_tester   ,				'''expected_num_hidden_layers'''   ,				self.model_tester.num_hidden_layers + 1     )
            self.assertEqual(len(A_     )   ,				A_     )

            # YOLOS has a different seq_length
            SCREAMING_SNAKE_CASE__											=			self.model_tester.expected_seq_len

            self.assertListEqual(
                list(hidden_states[0].shape[-2:]     )   ,				[seq_length, self.model_tester.hidden_size]   ,				)

        SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            SCREAMING_SNAKE_CASE__											=			True
            check_hidden_states_output(A_   ,				A_   ,				A_     )

            # check that output_hidden_states also work using config
            del inputs_dict["output_hidden_states"]
            SCREAMING_SNAKE_CASE__											=			True

            check_hidden_states_output(A_   ,				A_   ,				A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_object_detection(*A_     )




    @slow
    def  lowercase_					(       self     ):





        '''simple docstring'''

        for model_name in YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            SCREAMING_SNAKE_CASE__											=			YolosModel.from_pretrained(A_     )
            self.assertIsNotNone(A_     )





def 							__snake_case							(    )				->						Any:
    SCREAMING_SNAKE_CASE__											=			Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''	)
    return image



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


    '''simple docstring'''







    @cached_property
    def  lowercase_					(       self     ):





        '''simple docstring'''

        return AutoImageProcessor.from_pretrained('''hustvl/yolos-small'''     ) if is_vision_available() else None




    @slow
    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			YolosForObjectDetection.from_pretrained('''hustvl/yolos-small'''     ).to(A_     )

        SCREAMING_SNAKE_CASE__											=			self.default_image_processor
        SCREAMING_SNAKE_CASE__											=			prepare_img()
        SCREAMING_SNAKE_CASE__											=			image_processor(images=A_   ,				return_tensors='''pt'''     ).to(A_     )

        # forward pass
        with torch.no_grad():
            SCREAMING_SNAKE_CASE__											=			model(inputs.pixel_values     )

        # verify outputs
        SCREAMING_SNAKE_CASE__											=			torch.Size((1, 1_00, 92)     )
        self.assertEqual(outputs.logits.shape   ,				A_     )

        SCREAMING_SNAKE_CASE__											=			torch.tensor(
            [[-24.0248, -10.3024, -14.8290], [-42.0392, -16.8200, -27.4334], [-27.2743, -11.8154, -18.7148]]   ,				device=A_   ,				)
        SCREAMING_SNAKE_CASE__											=			torch.tensor(
            [[0.2559, 0.5455, 0.4706], [0.2989, 0.7279, 0.1875], [0.7732, 0.4017, 0.4462]]   ,				device=A_     )
        self.assertTrue(torch.allclose(outputs.logits[0, :3, :3]   ,				A_   ,				atol=1E-4     )     )
        self.assertTrue(torch.allclose(outputs.pred_boxes[0, :3, :3]   ,				A_   ,				atol=1E-4     )     )

        # verify postprocessing
        SCREAMING_SNAKE_CASE__											=			image_processor.post_process_object_detection(
            A_   ,				threshold=0.3   ,				target_sizes=[image.size[::-1]]     )[0]
        SCREAMING_SNAKE_CASE__											=			torch.tensor([0.9994, 0.9790, 0.9964, 0.9972, 0.9861]     ).to(A_     )
        SCREAMING_SNAKE_CASE__											=			[75, 75, 17, 63, 17]
        SCREAMING_SNAKE_CASE__											=			torch.tensor([335.0609, 79.3848, 375.4216, 187.2495]     ).to(A_     )

        self.assertEqual(len(results['''scores''']     )   ,				5     )
        self.assertTrue(torch.allclose(results['''scores''']   ,				A_   ,				atol=1E-4     )     )
        self.assertSequenceEqual(results['''labels'''].tolist()   ,				A_     )
        self.assertTrue(torch.allclose(results['''boxes'''][0, :]   ,				A_     )     )







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
