



from typing import List, Optional, Union

import torch
from transformers import (
    XLMRobertaTokenizer,
)

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)
from .text_encoder import MultilingualCLIP


a_      :      Dict    =     logging.get_logger(__name__)  # pylint: disable=invalid-name

a_      :      Union[str, Any]    =     '\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyPriorPipeline.from_pretrained("kandinsky-community/Kandinsky-2-1-prior")\n        >>> pipe_prior.to("cuda")\n\n        >>> prompt = "red cat, 4k photo"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> negative_image_emb = out.negative_image_embeds\n\n        >>> pipe = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1")\n        >>> pipe.to("cuda")\n\n        >>> image = pipe(\n        ...     prompt,\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=negative_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=100,\n        ... ).images\n\n        >>> image[0].save("cat.png")\n        ```\n'






def     __a					(		__UpperCAmelCase							,  __UpperCAmelCase							,  __UpperCAmelCase=8       ):
 a__												= h // scale_factor**2
 if h % scale_factor**2 != 0:
  new_h += 1
 a__												= w // scale_factor**2
 if w % scale_factor**2 != 0:
  new_w += 1
 return new_h * scale_factor, new_w * scale_factor

class    __UpperCamelCase			(       _lowercase				):

 """simple docstring"""



 def __init__(	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				) ->			Union[str, Any]:
  super().__init__()

  self.register_modules(
      text_encoder=SCREAMING_SNAKE_CASE		,				tokenizer=SCREAMING_SNAKE_CASE		,				unet=SCREAMING_SNAKE_CASE		,				scheduler=SCREAMING_SNAKE_CASE		,				movq=SCREAMING_SNAKE_CASE		,				)
  a__												= 2 ** (len(self.movq.config.block_out_channels      ) - 1)

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			int:
  if latents is None:
   a__												= randn_tensor(SCREAMING_SNAKE_CASE		,				generator=SCREAMING_SNAKE_CASE		,				device=SCREAMING_SNAKE_CASE		,				dtype=SCREAMING_SNAKE_CASE      )
  else:
   if latents.shape != shape:
    raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}"      )
   a__												= latents.to(SCREAMING_SNAKE_CASE      )

  a__												= latents * scheduler.init_noise_sigma
  return latents

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE=None		,				) ->			Any:
  a__												= len(SCREAMING_SNAKE_CASE      ) if isinstance(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) else 1
  # get prompt text embeddings
  a__												= self.tokenizer(
      SCREAMING_SNAKE_CASE		,				padding='''max_length'''		,				truncation=SCREAMING_SNAKE_CASE		,				max_length=7_7		,				return_attention_mask=SCREAMING_SNAKE_CASE		,				add_special_tokens=SCREAMING_SNAKE_CASE		,				return_tensors='''pt'''		,				)

  a__												= text_inputs.input_ids
  a__												= self.tokenizer(SCREAMING_SNAKE_CASE		,				padding='''longest'''		,				return_tensors='''pt'''      ).input_ids

  if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ):
   a__												= self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]      )
   logger.warning(
       '''The following part of your input was truncated because CLIP can only handle sequences up to'''
       f" {self.tokenizer.model_max_length} tokens: {removed_text}"      )

  a__												= text_input_ids.to(SCREAMING_SNAKE_CASE      )
  a__												= text_inputs.attention_mask.to(SCREAMING_SNAKE_CASE      )

  a__							,			a__												= self.text_encoder(
      input_ids=SCREAMING_SNAKE_CASE		,				attention_mask=SCREAMING_SNAKE_CASE      )

  a__												= prompt_embeds.repeat_interleave(SCREAMING_SNAKE_CASE		,				dim=0      )
  a__												= text_encoder_hidden_states.repeat_interleave(SCREAMING_SNAKE_CASE		,				dim=0      )
  a__												= text_mask.repeat_interleave(SCREAMING_SNAKE_CASE		,				dim=0      )

  if do_classifier_free_guidance:
   a__												= 42
   if negative_prompt is None:
    a__												= [''''''] * batch_size
   elif type(SCREAMING_SNAKE_CASE      ) is not type(SCREAMING_SNAKE_CASE      ):
    raise TypeError(
        f"`negative_prompt` should be the same type to `prompt`, but got {type(SCREAMING_SNAKE_CASE      )} !="
        f" {type(SCREAMING_SNAKE_CASE      )}."      )
   elif isinstance(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ):
    a__												= [negative_prompt]
   elif batch_size != len(SCREAMING_SNAKE_CASE      ):
    raise ValueError(
        f"`negative_prompt`: {negative_prompt} has batch size {len(SCREAMING_SNAKE_CASE      )}, but `prompt`:"
        f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
        ''' the batch size of `prompt`.'''      )
   else:
    a__												= negative_prompt

   a__												= self.tokenizer(
       SCREAMING_SNAKE_CASE		,				padding='''max_length'''		,				max_length=7_7		,				truncation=SCREAMING_SNAKE_CASE		,				return_attention_mask=SCREAMING_SNAKE_CASE		,				add_special_tokens=SCREAMING_SNAKE_CASE		,				return_tensors='''pt'''		,				)
   a__												= uncond_input.input_ids.to(SCREAMING_SNAKE_CASE      )
   a__												= uncond_input.attention_mask.to(SCREAMING_SNAKE_CASE      )

   a__							,			a__												= self.text_encoder(
       input_ids=SCREAMING_SNAKE_CASE		,				attention_mask=SCREAMING_SNAKE_CASE      )

   # duplicate unconditional embeddings for each generation per prompt, using mps friendly method

   a__												= negative_prompt_embeds.shape[1]
   a__												= negative_prompt_embeds.repeat(1		,				SCREAMING_SNAKE_CASE      )
   a__												= negative_prompt_embeds.view(batch_size * num_images_per_prompt		,				SCREAMING_SNAKE_CASE      )

   a__												= uncond_text_encoder_hidden_states.shape[1]
   a__												= uncond_text_encoder_hidden_states.repeat(1		,				SCREAMING_SNAKE_CASE		,				1      )
   a__												= uncond_text_encoder_hidden_states.view(
       batch_size * num_images_per_prompt		,				SCREAMING_SNAKE_CASE		,				-1      )
   a__												= uncond_text_mask.repeat_interleave(SCREAMING_SNAKE_CASE		,				dim=0      )

   # done duplicates

   # For classifier free guidance, we need to do two forward passes.
   # Here we concatenate the unconditional and text embeddings into a single batch
   # to avoid doing two forward passes
   a__												= torch.cat([negative_prompt_embeds, prompt_embeds]      )
   a__												= torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]      )

   a__												= torch.cat([uncond_text_mask, text_mask]      )

  return prompt_embeds, text_encoder_hidden_states, text_mask

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE=0      ) ->			str:
  if is_accelerate_available():
   from accelerate import cpu_offload
  else:
   raise ImportError('''Please install accelerate via `pip install accelerate`'''      )

  a__												= torch.device(f"cuda:{gpu_id}"      )

  a__												= [
      self.unet,
      self.text_encoder,
      self.movq,
  ]
  for cpu_offloaded_model in models:
   if cpu_offloaded_model is not None:
    cpu_offload(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE=0      ) ->			str:
  if is_accelerate_available() and is_accelerate_version('''>='''		,				'''0.17.0.dev0'''      ):
   from accelerate import cpu_offload_with_hook
  else:
   raise ImportError('''`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'''      )

  a__												= torch.device(f"cuda:{gpu_id}"      )

  if self.device.type != "cpu":
   self.to('''cpu'''		,				silence_dtype_warnings=SCREAMING_SNAKE_CASE      )
   torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

  a__												= None
  for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
   a__							,			a__												= cpu_offload_with_hook(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				prev_module_hook=SCREAMING_SNAKE_CASE      )

  if self.safety_checker is not None:
   a__							,			a__												= cpu_offload_with_hook(self.safety_checker		,				SCREAMING_SNAKE_CASE		,				prev_module_hook=SCREAMING_SNAKE_CASE      )

  # We'll offload the last model manually.
  a__												= hook

 @property
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
 def 						_UpperCAmelCase       (	self      ) ->			List[str]:
  if not hasattr(self.unet		,				'''_hf_hook'''      ):
   return self.device
  for module in self.unet.modules():
   if (
       hasattr(SCREAMING_SNAKE_CASE		,				'''_hf_hook'''      )
       and hasattr(module._hf_hook		,				'''execution_device'''      )
       and module._hf_hook.execution_device is not None
   ):
    return torch.device(module._hf_hook.execution_device      )
  return self.device



 @torch.no_grad()
 @replace_example_docstring(SCREAMING_SNAKE_CASE      )
 def __call__(	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE = None		,				SCREAMING_SNAKE_CASE = 5_1_2		,				SCREAMING_SNAKE_CASE = 5_1_2		,				SCREAMING_SNAKE_CASE = 1_0_0		,				SCREAMING_SNAKE_CASE = 4.0		,				SCREAMING_SNAKE_CASE = 1		,				SCREAMING_SNAKE_CASE = None		,				SCREAMING_SNAKE_CASE = None		,				SCREAMING_SNAKE_CASE = "pil"		,				SCREAMING_SNAKE_CASE = True		,				) ->			Any:

  if isinstance(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ):
   a__												= 1
  elif isinstance(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ):
   a__												= len(SCREAMING_SNAKE_CASE      )
  else:
   raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(SCREAMING_SNAKE_CASE      )}"      )

  a__												= self._execution_device

  a__												= batch_size * num_images_per_prompt
  a__												= guidance_scale > 1.0

  a__							,			a__							,			a__												= self._encode_prompt(
      SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )

  if isinstance(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ):
   a__												= torch.cat(SCREAMING_SNAKE_CASE		,				dim=0      )
  if isinstance(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ):
   a__												= torch.cat(SCREAMING_SNAKE_CASE		,				dim=0      )

  if do_classifier_free_guidance:
   a__												= image_embeds.repeat_interleave(SCREAMING_SNAKE_CASE		,				dim=0      )
   a__												= negative_image_embeds.repeat_interleave(SCREAMING_SNAKE_CASE		,				dim=0      )

  a__												= torch.cat([negative_image_embeds, image_embeds]		,				dim=0      ).to(
      dtype=prompt_embeds.dtype		,				device=SCREAMING_SNAKE_CASE      )

  self.scheduler.set_timesteps(SCREAMING_SNAKE_CASE		,				device=SCREAMING_SNAKE_CASE      )
  a__												= self.scheduler.timesteps

  a__												= self.unet.config.in_channels

  a__							,			a__												= get_new_h_w(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				self.movq_scale_factor      )

  # create initial latent
  a__												= self.prepare_latents(
      (batch_size, num_channels_latents, height, width)		,				text_encoder_hidden_states.dtype		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				self.scheduler		,				)

  for i, t in enumerate(self.progress_bar(SCREAMING_SNAKE_CASE      )      ):
   # expand the latents if we are doing classifier free guidance
   a__												= torch.cat([latents] * 2      ) if do_classifier_free_guidance else latents

   a__												= {'''text_embeds''': prompt_embeds, '''image_embeds''': image_embeds}
   a__												= self.unet(
       sample=SCREAMING_SNAKE_CASE		,				timestep=SCREAMING_SNAKE_CASE		,				encoder_hidden_states=SCREAMING_SNAKE_CASE		,				added_cond_kwargs=SCREAMING_SNAKE_CASE		,				return_dict=SCREAMING_SNAKE_CASE		,				)[0]

   if do_classifier_free_guidance:
    a__							,			a__												= noise_pred.split(latents.shape[1]		,				dim=1      )
    a__							,			a__												= noise_pred.chunk(2      )
    a__							,			a__												= variance_pred.chunk(2      )
    a__												= noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
    a__												= torch.cat([noise_pred, variance_pred_text]		,				dim=1      )

   if not (
       hasattr(self.scheduler.config		,				'''variance_type'''      )
       and self.scheduler.config.variance_type in ["learned", "learned_range"]
   ):
    a__							,			a__												= noise_pred.split(latents.shape[1]		,				dim=1      )

   # compute the previous noisy sample x_t -> x_t-1
   a__												= self.scheduler.step(
       SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				generator=SCREAMING_SNAKE_CASE		,				).prev_sample
  # post-processing
  a__												= self.movq.decode(SCREAMING_SNAKE_CASE		,				force_not_quantize=SCREAMING_SNAKE_CASE      )['''sample''']

  if output_type not in ["pt", "np", "pil"]:
   raise ValueError(f"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}"      )

  if output_type in ["np", "pil"]:
   a__												= image * 0.5 + 0.5
   a__												= image.clamp(0		,				1      )
   a__												= image.cpu().permute(0		,				2		,				3		,				1      ).float().numpy()

  if output_type == "pil":
   a__												= self.numpy_to_pil(SCREAMING_SNAKE_CASE      )

  if not return_dict:
   return (image,)

  return ImagePipelineOutput(images=SCREAMING_SNAKE_CASE      )










import torch
from torch import nn
class    __UpperCamelCase			(       nn.Module				):

 """simple docstring"""



 def __init__(	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE=1		,				SCREAMING_SNAKE_CASE=False      ) ->			Tuple:
  super().__init__()

  a__												= n_token
  a__												= d_embed
  a__												= d_proj

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

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

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

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

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

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

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

  a__												= keep_order

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

  return logit

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE=None		,				SCREAMING_SNAKE_CASE=False      ) ->			Optional[int]:

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

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

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

    weights.append(SCREAMING_SNAKE_CASE      )
    biases.append(SCREAMING_SNAKE_CASE      )

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

   a__												= self._compute_logit(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )
   a__												= nn.functional.log_softmax(SCREAMING_SNAKE_CASE		,				dim=1      )

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

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

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

     if indices_i.numel() == 0:
      continue

     a__												= labels.index_select(0		,				SCREAMING_SNAKE_CASE      ) - l_idx
     a__												= head_logprob.index_select(0		,				SCREAMING_SNAKE_CASE      )
     a__												= hidden.index_select(0		,				SCREAMING_SNAKE_CASE      )
    else:
     a__												= hidden

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

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

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

  return out



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

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

    weights.append(SCREAMING_SNAKE_CASE      )
    biases.append(SCREAMING_SNAKE_CASE      )

   a__							,			a__							,			a__												= weights[0], biases[0], self.out_projs[0]
   a__												= self._compute_logit(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )

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

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

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

     a__												= self._compute_logit(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )
     a__												= nn.functional.log_softmax(SCREAMING_SNAKE_CASE		,				dim=1      )

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

   return out





