





import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()



class 					__snake_case	(  unittest.TestCase    ):


    '''simple docstring'''







    def  lowercase_					(       self     ):





        '''simple docstring'''

        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()



    @property
    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			1
        SCREAMING_SNAKE_CASE__											=			3
        SCREAMING_SNAKE_CASE__											=			(32, 32)

        SCREAMING_SNAKE_CASE__											=			floats_tensor((batch_size, num_channels) + sizes   ,				rng=random.Random(0     )     ).to(A_     )
        return image



    @property
    def  lowercase_					(       self     ):





        '''simple docstring'''

        torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			UNetaDConditionModel(
            block_out_channels=(32, 32, 64)   ,				layers_per_block=2   ,				sample_size=32   ,				in_channels=7   ,				out_channels=4   ,				down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''', '''CrossAttnDownBlock2D''')   ,				up_block_types=('''CrossAttnUpBlock2D''', '''CrossAttnUpBlock2D''', '''UpBlock2D''')   ,				cross_attention_dim=32   ,				attention_head_dim=8   ,				use_linear_projection=A_   ,				only_cross_attention=(True, True, False)   ,				num_class_embeds=1_00   ,				)
        return model



    @property
    def  lowercase_					(       self     ):





        '''simple docstring'''

        torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			AutoencoderKL(
            block_out_channels=[32, 32, 64]   ,				in_channels=3   ,				out_channels=3   ,				down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,				up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,				latent_channels=4   ,				)
        return model



    @property
    def  lowercase_					(       self     ):





        '''simple docstring'''

        torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			CLIPTextConfig(
            bos_token_id=0   ,				eos_token_id=2   ,				hidden_size=32   ,				intermediate_size=37   ,				layer_norm_eps=1E-05   ,				num_attention_heads=4   ,				num_hidden_layers=5   ,				pad_token_id=1   ,				vocab_size=10_00   ,				hidden_act='''gelu'''   ,				projection_dim=5_12   ,				)
        return CLIPTextModel(A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        SCREAMING_SNAKE_CASE__											=			self.dummy_cond_unet_upscale
        SCREAMING_SNAKE_CASE__											=			DDPMScheduler()
        SCREAMING_SNAKE_CASE__											=			DDIMScheduler(prediction_type='''v_prediction'''     )
        SCREAMING_SNAKE_CASE__											=			self.dummy_vae
        SCREAMING_SNAKE_CASE__											=			self.dummy_text_encoder
        SCREAMING_SNAKE_CASE__											=			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )

        SCREAMING_SNAKE_CASE__											=			self.dummy_image.cpu().permute(0   ,				2   ,				3   ,				1     )[0]
        SCREAMING_SNAKE_CASE__											=			Image.fromarray(np.uinta(A_     )     ).convert('''RGB'''     ).resize((64, 64)     )

        # make sure here that pndm scheduler skips prk
        SCREAMING_SNAKE_CASE__											=			StableDiffusionUpscalePipeline(
            unet=A_   ,				low_res_scheduler=A_   ,				scheduler=A_   ,				vae=A_   ,				text_encoder=A_   ,				tokenizer=A_   ,				max_noise_level=3_50   ,				)
        SCREAMING_SNAKE_CASE__											=			sd_pipe.to(A_     )
        sd_pipe.set_progress_bar_config(disable=A_     )

        SCREAMING_SNAKE_CASE__											=			'''A painting of a squirrel eating a burger'''
        SCREAMING_SNAKE_CASE__											=			torch.Generator(device=A_     ).manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			sd_pipe(
            [prompt]   ,				image=A_   ,				generator=A_   ,				guidance_scale=6.0   ,				noise_level=20   ,				num_inference_steps=2   ,				output_type='''np'''   ,				)

        SCREAMING_SNAKE_CASE__											=			output.images

        SCREAMING_SNAKE_CASE__											=			torch.Generator(device=A_     ).manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			sd_pipe(
            [prompt]   ,				image=A_   ,				generator=A_   ,				guidance_scale=6.0   ,				noise_level=20   ,				num_inference_steps=2   ,				output_type='''np'''   ,				return_dict=A_   ,				)[0]

        SCREAMING_SNAKE_CASE__											=			image[0, -3:, -3:, -1]
        SCREAMING_SNAKE_CASE__											=			image_from_tuple[0, -3:, -3:, -1]

        SCREAMING_SNAKE_CASE__											=			low_res_image.size[0] * 4
        assert image.shape == (1, expected_height_width, expected_height_width, 3)
        SCREAMING_SNAKE_CASE__											=			np.array([0.3113, 0.3910, 0.4272, 0.4859, 0.5061, 0.4652, 0.5362, 0.5715, 0.5661]     )

        assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1E-2
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice     ).max() < 1E-2



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        SCREAMING_SNAKE_CASE__											=			self.dummy_cond_unet_upscale
        SCREAMING_SNAKE_CASE__											=			DDPMScheduler()
        SCREAMING_SNAKE_CASE__											=			DDIMScheduler(prediction_type='''v_prediction'''     )
        SCREAMING_SNAKE_CASE__											=			self.dummy_vae
        SCREAMING_SNAKE_CASE__											=			self.dummy_text_encoder
        SCREAMING_SNAKE_CASE__											=			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )

        SCREAMING_SNAKE_CASE__											=			self.dummy_image.cpu().permute(0   ,				2   ,				3   ,				1     )[0]
        SCREAMING_SNAKE_CASE__											=			Image.fromarray(np.uinta(A_     )     ).convert('''RGB'''     ).resize((64, 64)     )

        # make sure here that pndm scheduler skips prk
        SCREAMING_SNAKE_CASE__											=			StableDiffusionUpscalePipeline(
            unet=A_   ,				low_res_scheduler=A_   ,				scheduler=A_   ,				vae=A_   ,				text_encoder=A_   ,				tokenizer=A_   ,				max_noise_level=3_50   ,				)
        SCREAMING_SNAKE_CASE__											=			sd_pipe.to(A_     )
        sd_pipe.set_progress_bar_config(disable=A_     )

        SCREAMING_SNAKE_CASE__											=			'''A painting of a squirrel eating a burger'''
        SCREAMING_SNAKE_CASE__											=			sd_pipe(
            2 * [prompt]   ,				image=2 * [low_res_image]   ,				guidance_scale=6.0   ,				noise_level=20   ,				num_inference_steps=2   ,				output_type='''np'''   ,				)
        SCREAMING_SNAKE_CASE__											=			output.images
        assert image.shape[0] == 2

        SCREAMING_SNAKE_CASE__											=			torch.Generator(device=A_     ).manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			sd_pipe(
            [prompt]   ,				image=A_   ,				generator=A_   ,				num_images_per_prompt=2   ,				guidance_scale=6.0   ,				noise_level=20   ,				num_inference_steps=2   ,				output_type='''np'''   ,				)
        SCREAMING_SNAKE_CASE__											=			output.images
        assert image.shape[0] == 2




    @unittest.skipIf(torch_device != '''cuda'''   ,				'''This test requires a GPU'''     )
    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.dummy_cond_unet_upscale
        SCREAMING_SNAKE_CASE__											=			DDPMScheduler()
        SCREAMING_SNAKE_CASE__											=			DDIMScheduler(prediction_type='''v_prediction'''     )
        SCREAMING_SNAKE_CASE__											=			self.dummy_vae
        SCREAMING_SNAKE_CASE__											=			self.dummy_text_encoder
        SCREAMING_SNAKE_CASE__											=			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )

        SCREAMING_SNAKE_CASE__											=			self.dummy_image.cpu().permute(0   ,				2   ,				3   ,				1     )[0]
        SCREAMING_SNAKE_CASE__											=			Image.fromarray(np.uinta(A_     )     ).convert('''RGB'''     ).resize((64, 64)     )

        # put models in fp16, except vae as it overflows in fp16
        SCREAMING_SNAKE_CASE__											=			unet.half()
        SCREAMING_SNAKE_CASE__											=			text_encoder.half()

        # make sure here that pndm scheduler skips prk
        SCREAMING_SNAKE_CASE__											=			StableDiffusionUpscalePipeline(
            unet=A_   ,				low_res_scheduler=A_   ,				scheduler=A_   ,				vae=A_   ,				text_encoder=A_   ,				tokenizer=A_   ,				max_noise_level=3_50   ,				)
        SCREAMING_SNAKE_CASE__											=			sd_pipe.to(A_     )
        sd_pipe.set_progress_bar_config(disable=A_     )

        SCREAMING_SNAKE_CASE__											=			'''A painting of a squirrel eating a burger'''
        SCREAMING_SNAKE_CASE__											=			torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			sd_pipe(
            [prompt]   ,				image=A_   ,				generator=A_   ,				num_inference_steps=2   ,				output_type='''np'''   ,				).images

        SCREAMING_SNAKE_CASE__											=			low_res_image.size[0] * 4
        assert image.shape == (1, expected_height_width, expected_height_width, 3)





@slow
@require_torch_gpu
class 					__snake_case	(  unittest.TestCase    ):


    '''simple docstring'''







    def  lowercase_					(       self     ):





        '''simple docstring'''

        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/sd2-upscale/low_res_cat.png'''     )
        SCREAMING_SNAKE_CASE__											=			load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'''
            '''/upsampled_cat.npy'''     )

        SCREAMING_SNAKE_CASE__											=			'''stabilityai/stable-diffusion-x4-upscaler'''
        SCREAMING_SNAKE_CASE__											=			StableDiffusionUpscalePipeline.from_pretrained(A_     )
        pipe.to(A_     )
        pipe.set_progress_bar_config(disable=A_     )
        pipe.enable_attention_slicing()

        SCREAMING_SNAKE_CASE__											=			'''a cat sitting on a park bench'''

        SCREAMING_SNAKE_CASE__											=			torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			pipe(
            prompt=A_   ,				image=A_   ,				generator=A_   ,				output_type='''np'''   ,				)
        SCREAMING_SNAKE_CASE__											=			output.images[0]

        assert image.shape == (5_12, 5_12, 3)
        assert np.abs(expected_image - image     ).max() < 1E-3



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/sd2-upscale/low_res_cat.png'''     )
        SCREAMING_SNAKE_CASE__											=			load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'''
            '''/upsampled_cat_fp16.npy'''     )

        SCREAMING_SNAKE_CASE__											=			'''stabilityai/stable-diffusion-x4-upscaler'''
        SCREAMING_SNAKE_CASE__											=			StableDiffusionUpscalePipeline.from_pretrained(
            A_   ,				torch_dtype=torch.floataa   ,				)
        pipe.to(A_     )
        pipe.set_progress_bar_config(disable=A_     )
        pipe.enable_attention_slicing()

        SCREAMING_SNAKE_CASE__											=			'''a cat sitting on a park bench'''

        SCREAMING_SNAKE_CASE__											=			torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			pipe(
            prompt=A_   ,				image=A_   ,				generator=A_   ,				output_type='''np'''   ,				)
        SCREAMING_SNAKE_CASE__											=			output.images[0]

        assert image.shape == (5_12, 5_12, 3)
        assert np.abs(expected_image - image     ).max() < 5E-1




    def  lowercase_					(       self     ):





        '''simple docstring'''

        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        SCREAMING_SNAKE_CASE__											=			load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/sd2-upscale/low_res_cat.png'''     )

        SCREAMING_SNAKE_CASE__											=			'''stabilityai/stable-diffusion-x4-upscaler'''
        SCREAMING_SNAKE_CASE__											=			StableDiffusionUpscalePipeline.from_pretrained(
            A_   ,				torch_dtype=torch.floataa   ,				)
        pipe.to(A_     )
        pipe.set_progress_bar_config(disable=A_     )
        pipe.enable_attention_slicing(1     )
        pipe.enable_sequential_cpu_offload()

        SCREAMING_SNAKE_CASE__											=			'''a cat sitting on a park bench'''

        SCREAMING_SNAKE_CASE__											=			torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			pipe(
            prompt=A_   ,				image=A_   ,				generator=A_   ,				num_inference_steps=5   ,				output_type='''np'''   ,				)

        SCREAMING_SNAKE_CASE__											=			torch.cuda.max_memory_allocated()
        # make sure that less than 2.9 GB is allocated
        assert mem_bytes < 2.9 * 10**9







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
