





import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()



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


    '''simple docstring'''







    lowerCamelCase__    :		int				=	StableUnCLIPImgaImgPipeline
    lowerCamelCase__    :		Dict				=	TEXT_GUIDED_IMAGE_VARIATION_PARAMS
    lowerCamelCase__    :		Any				=	TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    lowerCamelCase__    :		List[Any]				=	frozenset(
        []    )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
    lowerCamelCase__    :		int				=	frozenset([]    )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			32
        SCREAMING_SNAKE_CASE__											=			embedder_hidden_size

        # image encoding components

        SCREAMING_SNAKE_CASE__											=			CLIPImageProcessor(crop_size=32   ,				size=32     )

        torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			CLIPVisionModelWithProjection(
            CLIPVisionConfig(
                hidden_size=A_   ,				projection_dim=A_   ,				num_hidden_layers=5   ,				num_attention_heads=4   ,				image_size=32   ,				intermediate_size=37   ,				patch_size=1   ,				)     )

        # regular denoising components

        torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			StableUnCLIPImageNormalizer(embedding_dim=A_     )
        SCREAMING_SNAKE_CASE__											=			DDPMScheduler(beta_schedule='''squaredcos_cap_v2'''     )

        torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )

        torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			CLIPTextModel(
            CLIPTextConfig(
                bos_token_id=0   ,				eos_token_id=2   ,				hidden_size=A_   ,				projection_dim=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   ,				)     )

        torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			UNetaDConditionModel(
            sample_size=32   ,				in_channels=4   ,				out_channels=4   ,				down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''')   ,				up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''')   ,				block_out_channels=(32, 64)   ,				attention_head_dim=(2, 4)   ,				class_embed_type='''projection'''   ,				projection_class_embeddings_input_dim=embedder_projection_dim * 2   ,				cross_attention_dim=A_   ,				layers_per_block=1   ,				upcast_attention=A_   ,				use_linear_projection=A_   ,				)

        torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			DDIMScheduler(
            beta_schedule='''scaled_linear'''   ,				beta_start=0.00085   ,				beta_end=0.012   ,				prediction_type='''v_prediction'''   ,				set_alpha_to_one=A_   ,				steps_offset=1   ,				)

        torch.manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			AutoencoderKL()

        SCREAMING_SNAKE_CASE__											=			{
            # image encoding components
            '''feature_extractor''': feature_extractor,
            '''image_encoder''': image_encoder.eval(),
            # image noising components
            '''image_normalizer''': image_normalizer.eval(),
            '''image_noising_scheduler''': image_noising_scheduler,
            # regular denoising components
            '''tokenizer''': tokenizer,
            '''text_encoder''': text_encoder.eval(),
            '''unet''': unet.eval(),
            '''scheduler''': scheduler,
            '''vae''': vae.eval(),
        }

        return components



    def  lowercase_					(       self   ,				A_   ,				A_=0   ,				A_=True     ):





        '''simple docstring'''

        if str(A_     ).startswith('''mps'''     ):
            SCREAMING_SNAKE_CASE__											=			torch.manual_seed(A_     )
        else:
            SCREAMING_SNAKE_CASE__											=			torch.Generator(device=A_     ).manual_seed(A_     )

        SCREAMING_SNAKE_CASE__											=			floats_tensor((1, 3, 32, 32)   ,				rng=random.Random(A_     )     ).to(A_     )

        if pil_image:
            SCREAMING_SNAKE_CASE__											=			input_image * 0.5 + 0.5
            SCREAMING_SNAKE_CASE__											=			input_image.clamp(0   ,				1     )
            SCREAMING_SNAKE_CASE__											=			input_image.cpu().permute(0   ,				2   ,				3   ,				1     ).float().numpy()
            SCREAMING_SNAKE_CASE__											=			DiffusionPipeline.numpy_to_pil(A_     )[0]

        return {
            "prompt": "An anime racoon running a marathon",
            "image": input_image,
            "generator": generator,
            "num_inference_steps": 2,
            "output_type": "np",
        }



    @skip_mps
    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        SCREAMING_SNAKE_CASE__											=			self.get_dummy_components()
        SCREAMING_SNAKE_CASE__											=			StableUnCLIPImgaImgPipeline(**A_     )
        SCREAMING_SNAKE_CASE__											=			sd_pipe.to(A_     )
        sd_pipe.set_progress_bar_config(disable=A_     )

        SCREAMING_SNAKE_CASE__											=			self.get_dummy_inputs(A_     )
        inputs.update({'''image_embeds''': None}     )
        SCREAMING_SNAKE_CASE__											=			sd_pipe(**A_     ).images
        SCREAMING_SNAKE_CASE__											=			image[0, -3:, -3:, -1]

        assert image.shape == (1, 32, 32, 3)
        SCREAMING_SNAKE_CASE__											=			np.array([0.3872, 0.7224, 0.5601, 0.4741, 0.6872, 0.5814, 0.4636, 0.3867, 0.5078]     )

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



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			torch_device in ['''cpu''', '''mps''']

        self._test_attention_slicing_forward_pass(test_max_difference=A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			torch_device in ['''cpu''', '''mps''']

        self._test_inference_batch_single_identical(test_max_difference=A_     )




    @unittest.skipIf(
        torch_device != '''cuda''' or not is_xformers_available()   ,				reason='''XFormers attention is only available with CUDA and `xformers` installed'''   ,				)
    def  lowercase_					(       self     ):





        '''simple docstring'''

        self._test_xformers_attention_forwardGenerator_pass(test_max_difference=A_     )





@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/stable_unclip/turtle.png'''     )

        SCREAMING_SNAKE_CASE__											=			load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy'''     )

        SCREAMING_SNAKE_CASE__											=			StableUnCLIPImgaImgPipeline.from_pretrained(
            '''fusing/stable-unclip-2-1-l-img2img'''   ,				torch_dtype=torch.floataa     )
        pipe.to(A_     )
        pipe.set_progress_bar_config(disable=A_     )
        # stable unclip will oom when integration tests are run on a V100,
        # so turn on memory savings
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        SCREAMING_SNAKE_CASE__											=			torch.Generator(device='''cpu'''     ).manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			pipe(A_   ,				'''anime turle'''   ,				generator=A_   ,				output_type='''np'''     )

        SCREAMING_SNAKE_CASE__											=			output.images[0]

        assert image.shape == (7_68, 7_68, 3)

        assert_mean_pixel_difference(A_   ,				A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

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

        SCREAMING_SNAKE_CASE__											=			load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy'''     )

        SCREAMING_SNAKE_CASE__											=			StableUnCLIPImgaImgPipeline.from_pretrained(
            '''fusing/stable-unclip-2-1-h-img2img'''   ,				torch_dtype=torch.floataa     )
        pipe.to(A_     )
        pipe.set_progress_bar_config(disable=A_     )
        # stable unclip will oom when integration tests are run on a V100,
        # so turn on memory savings
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        SCREAMING_SNAKE_CASE__											=			torch.Generator(device='''cpu'''     ).manual_seed(0     )
        SCREAMING_SNAKE_CASE__											=			pipe(A_   ,				'''anime turle'''   ,				generator=A_   ,				output_type='''np'''     )

        SCREAMING_SNAKE_CASE__											=			output.images[0]

        assert image.shape == (7_68, 7_68, 3)

        assert_mean_pixel_difference(A_   ,				A_     )




    def  lowercase_					(       self     ):





        '''simple docstring'''

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

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

        SCREAMING_SNAKE_CASE__											=			StableUnCLIPImgaImgPipeline.from_pretrained(
            '''fusing/stable-unclip-2-1-h-img2img'''   ,				torch_dtype=torch.floataa     )
        SCREAMING_SNAKE_CASE__											=			pipe.to(A_     )
        pipe.set_progress_bar_config(disable=A_     )
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        SCREAMING_SNAKE_CASE__											=			pipe(
            A_   ,				'''anime turtle'''   ,				num_inference_steps=2   ,				output_type='''np'''   ,				)

        SCREAMING_SNAKE_CASE__											=			torch.cuda.max_memory_allocated()
        # make sure that less than 7 GB is allocated
        assert mem_bytes < 7 * 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
