
from abc import ABC, abstractmethod
from typing import List, Optional


class        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self )						->      Optional[Any]:
				"""simple docstring"""
				# test for the above condition
				self.test()



		def __UpperCamelCase				(				self )						->      Dict:
				"""simple docstring"""
				UpperCamelCase						=  0
				UpperCamelCase						=  False
				while not completed:
						if counter == 1:
								self.reset()
						UpperCamelCase						=  self.advance()
						if not self.does_advance(A_ ):
								raise Exception(
								    'Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true.' )

						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase						=  self.update(A_ )
						counter += 1

						if counter > 10_000:
								raise Exception('update() does not fulfill the constraint.' )

				if self.remaining() != 0:
						raise Exception('Custom Constraint is not defined correctly.' )



		@abstractmethod
		def __UpperCamelCase				(				self )						->      Optional[Any]:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )



		@abstractmethod
		def __UpperCamelCase				(				self       ,   A_ )						->      str:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )



		@abstractmethod
		def __UpperCamelCase				(				self       ,   A_ )						->      int:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )



		@abstractmethod
		def __UpperCamelCase				(				self )						->      Any:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )



		@abstractmethod
		def __UpperCamelCase				(				self )						->      str:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )



		@abstractmethod
		def __UpperCamelCase				(				self       ,   A_=False )						->      int:
				"""simple docstring"""
				raise NotImplementedError(
				    F'''{self.__class__} is an abstract class. Only classes inheriting this class can be called.''' )






class        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_ )						->      Any:
				"""simple docstring"""
				super(A_       ,   self ).__init__()

				if not isinstance(A_       ,   A_ ) or len(A_ ) == 0:
						raise ValueError(F'''`token_ids` has to be a non-empty list, but is {token_ids}.''' )
				if any((not isinstance(A_       ,   A_ ) or token_id < 0) for token_id in token_ids ):
						raise ValueError(F'''Each list in `token_ids` has to be a list of positive integers, but is {token_ids}.''' )

				UpperCamelCase						=  token_ids

				UpperCamelCase						=  len(self.token_ids )
				UpperCamelCase						=  -1  # the index of the currently fulfilled step
				UpperCamelCase						=  False



		def __UpperCamelCase				(				self )						->      Optional[Any]:
				"""simple docstring"""
				if self.completed:
						return None
				return self.token_ids[self.fulfilled_idx + 1]



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[int]:
				"""simple docstring"""
				if not isinstance(A_       ,   A_ ):
						raise ValueError(F'''`token_id` has to be an `int`, but is {token_id} of type {type(A_ )}''' )

				if self.completed:
						return False

				return token_id == self.token_ids[self.fulfilled_idx + 1]



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[int]:
				"""simple docstring"""
				if not isinstance(A_       ,   A_ ):
						raise ValueError(F'''`token_id` has to be an `int`, but is {token_id} of type {type(A_ )}''' )

				UpperCamelCase						=  False
				UpperCamelCase						=  False
				UpperCamelCase						=  False

				if self.does_advance(A_ ):
						self.fulfilled_idx += 1
						UpperCamelCase						=  True
						if self.fulfilled_idx == (self.seqlen - 1):
								UpperCamelCase						=  True
						UpperCamelCase						=  completed
				else:
						# failed to make progress.
						UpperCamelCase						=  True
						self.reset()
				return stepped, completed, reset



		def __UpperCamelCase				(				self )						->      Dict:
				"""simple docstring"""
				UpperCamelCase						=  False
				UpperCamelCase						=  0



		def __UpperCamelCase				(				self )						->      int:
				"""simple docstring"""
				return self.seqlen - (self.fulfilled_idx + 1)



		def __UpperCamelCase				(				self       ,   A_=False )						->      Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase						=  PhrasalConstraint(self.token_ids )

				if stateful:
						UpperCamelCase						=  self.seqlen
						UpperCamelCase						=  self.fulfilled_idx
						UpperCamelCase						=  self.completed

				return new_constraint






class        lowercase      :



		def __init__(				self       ,   A_       ,   A_=True )						->      List[Any]:
				"""simple docstring"""
				UpperCamelCase						=  max([len(A_ ) for one in nested_token_ids] )

				UpperCamelCase						=  {}
				for token_ids in nested_token_ids:
						UpperCamelCase						=  root
						for tidx, token_id in enumerate(A_ ):
								if token_id not in level:
										UpperCamelCase						=  {}

								UpperCamelCase						=  level[token_id]

				if no_subsets and self.has_subsets(A_       ,   A_ ):
						raise ValueError(
						    'Each list in `nested_token_ids` can\'t be a complete subset of another list, but is'
						    F''' {nested_token_ids}.''' )

				UpperCamelCase						=  root



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[Any]:
				"""simple docstring"""
				UpperCamelCase						=  self.trie

				for current_token in current_seq:
						UpperCamelCase						=  start[current_token]

				UpperCamelCase						=  list(start.keys() )

				return next_tokens



		def __UpperCamelCase				(				self       ,   A_ )						->      Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase						=  self.next_tokens(A_ )

				return len(A_ ) == 0



		def __UpperCamelCase				(				self       ,   A_ )						->      List[str]:
				"""simple docstring"""
				UpperCamelCase						=  list(root.values() )
				if len(A_ ) == 0:
						return 1
				else:
						return sum([self.count_leaves(A_ ) for nn in next_nodes] )



		def __UpperCamelCase				(				self       ,   A_       ,   A_ )						->      Optional[int]:
				"""simple docstring"""
				UpperCamelCase						=  self.count_leaves(A_ )
				return len(A_ ) != leaf_count






class        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_ )						->      str:
				"""simple docstring"""
				super(A_       ,   self ).__init__()

				if not isinstance(A_       ,   A_ ) or len(A_ ) == 0:
						raise ValueError(F'''`nested_token_ids` has to be a non-empty list, but is {nested_token_ids}.''' )
				if any(not isinstance(A_       ,   A_ ) for token_ids in nested_token_ids ):
						raise ValueError(F'''`nested_token_ids` has to be a list of lists, but is {nested_token_ids}.''' )
				if any(
				    any((not isinstance(A_       ,   A_ ) or token_id < 0) for token_id in token_ids )
				    for token_ids in nested_token_ids ):
						raise ValueError(
						    F'''Each list in `nested_token_ids` has to be a list of positive integers, but is {nested_token_ids}.''' )

				UpperCamelCase						=  DisjunctiveTrie(A_ )
				UpperCamelCase						=  nested_token_ids

				UpperCamelCase						=  self.trie.max_height
				UpperCamelCase						=  []
				UpperCamelCase						=  False



		def __UpperCamelCase				(				self )						->      Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase						=  self.trie.next_tokens(self.current_seq )

				if len(A_ ) == 0:
						return None
				else:
						return token_list



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[Any]:
				"""simple docstring"""
				if not isinstance(A_       ,   A_ ):
						raise ValueError(F'''`token_id` is supposed to be type `int`, but is {token_id} of type {type(A_ )}''' )

				UpperCamelCase						=  self.trie.next_tokens(self.current_seq )

				return token_id in next_tokens



		def __UpperCamelCase				(				self       ,   A_ )						->      Optional[Any]:
				"""simple docstring"""
				if not isinstance(A_       ,   A_ ):
						raise ValueError(F'''`token_id` is supposed to be type `int`, but is {token_id} of type {type(A_ )}''' )

				UpperCamelCase						=  False
				UpperCamelCase						=  False
				UpperCamelCase						=  False

				if self.does_advance(A_ ):
						self.current_seq.append(A_ )
						UpperCamelCase						=  True
				else:
						UpperCamelCase						=  True
						self.reset()

				UpperCamelCase						=  self.trie.reached_leaf(self.current_seq )
				UpperCamelCase						=  completed

				return stepped, completed, reset



		def __UpperCamelCase				(				self )						->      str:
				"""simple docstring"""
				UpperCamelCase						=  False
				UpperCamelCase						=  []



		def __UpperCamelCase				(				self )						->      Optional[int]:
				"""simple docstring"""
				if self.completed:
						# since this can be completed without reaching max height
						return 0
				else:
						return self.seqlen - len(self.current_seq )



		def __UpperCamelCase				(				self       ,   A_=False )						->      int:
				"""simple docstring"""
				UpperCamelCase						=  DisjunctiveConstraint(self.token_ids )

				if stateful:
						UpperCamelCase						=  self.seqlen
						UpperCamelCase						=  self.current_seq
						UpperCamelCase						=  self.completed

				return new_constraint






class        lowercase      :



		def __init__(				self       ,   A_ )						->      Tuple:
				"""simple docstring"""
				UpperCamelCase						=  constraints

				# max # of steps required to fulfill a given constraint
				UpperCamelCase						=  max([c.seqlen for c in constraints] )
				UpperCamelCase						=  len(A_ )
				UpperCamelCase						=  False

				self.init_state()



		def __UpperCamelCase				(				self )						->      List[str]:
				"""simple docstring"""
				UpperCamelCase						=  []
				UpperCamelCase						=  None
				UpperCamelCase						=  [constraint.copy(stateful=A_ ) for constraint in self.constraints]



		def __UpperCamelCase				(				self )						->      Optional[int]:
				"""simple docstring"""
				UpperCamelCase						=  0
				if self.inprogress_constraint:
						# extra points for having a constraint mid-fulfilled
						add += self.max_seqlen - self.inprogress_constraint.remaining()

				return (len(self.complete_constraints ) * self.max_seqlen) + add



		def __UpperCamelCase				(				self )						->      Optional[int]:
				"""simple docstring"""
				UpperCamelCase						=  []
				if self.inprogress_constraint is None:
						for constraint in self.pending_constraints:  # "pending" == "unfulfilled yet"
								UpperCamelCase						=  constraint.advance()
								if isinstance(A_       ,   A_ ):
										token_list.append(A_ )
								elif isinstance(A_       ,   A_ ):
										token_list.extend(A_ )
				else:
						UpperCamelCase						=  self.inprogress_constraint.advance()
						if isinstance(A_       ,   A_ ):
								token_list.append(A_ )
						elif isinstance(A_       ,   A_ ):
								token_list.extend(A_ )

				if len(A_ ) == 0:
						return None
				else:
						return token_list



		def __UpperCamelCase				(				self       ,   A_ )						->      Any:
				"""simple docstring"""
				self.init_state()

				if token_ids is not None:
						for token in token_ids:
								# completes or steps **one** constraint
								UpperCamelCase			,						UpperCamelCase						=  self.add(A_ )

								# the entire list of constraints are fulfilled
								if self.completed:
										break



		def __UpperCamelCase				(				self       ,   A_ )						->      int:
				"""simple docstring"""
				if not isinstance(A_       ,   A_ ):
						raise ValueError(F'''`token_id` should be an `int`, but is `{token_id}`.''' )

				UpperCamelCase			,						UpperCamelCase						=  False, False

				if self.completed:
						UpperCamelCase						=  True
						UpperCamelCase						=  False
						return complete, stepped

				if self.inprogress_constraint is not None:
						# In the middle of fulfilling a constraint. If the `token_id` *does* makes an incremental progress to current
						# job, simply update the state

						UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase						=  self.inprogress_constraint.update(A_ )
						if reset:
								# 1. If the next token breaks the progress, then we must restart.
								#     e.g. constraint = "I love pies" and sequence so far is "I love" but `token_id` == "books".

								#     But that doesn't mean we self.init_state(), since we only reset the state for this particular
								#     constraint, not the full list of constraints.

								self.pending_constraints.append(self.inprogress_constraint.copy(stateful=A_ ) )
								UpperCamelCase						=  None

						if complete:
								# 2. If the next token completes the constraint, move it to completed list, set
								#     inprogress to None. If there are no pending constraints either, then this full list of constraints
								#     is complete.

								self.complete_constraints.append(self.inprogress_constraint )
								UpperCamelCase						=  None

								if len(self.pending_constraints ) == 0:
										# we're done!
										UpperCamelCase						=  True

				else:
						# Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards any of our list
						# of constraints?

						for cidx, pending_constraint in enumerate(self.pending_constraints ):
								if pending_constraint.does_advance(A_ ):
										UpperCamelCase			,						UpperCamelCase			,						UpperCamelCase						=  pending_constraint.update(A_ )

										if not stepped:
												raise Exception(
												    '`constraint.update(token_id)` is not yielding incremental progress, '
												    'even though `constraint.does_advance(token_id)` is true.' )

										if complete:
												self.complete_constraints.append(A_ )
												UpperCamelCase						=  None

										if not complete and stepped:
												UpperCamelCase						=  pending_constraint

										if complete or stepped:
												# If we made any progress at all, then it's at least not a "pending constraint".

												UpperCamelCase						=  (
												    self.pending_constraints[:cidx] + self.pending_constraints[cidx + 1 :]
												)

												if len(self.pending_constraints ) == 0 and self.inprogress_constraint is None:
														# If there's no longer any pending after this and no inprogress either, then we must be
														# complete.

														UpperCamelCase						=  True

												break  # prevent accidentally stepping through multiple constraints with just one token.

				return complete, stepped



		def __UpperCamelCase				(				self       ,   A_=True )						->      Tuple:
				"""simple docstring"""
				UpperCamelCase						=  ConstraintListState(self.constraints )  # we actually never though self.constraints objects
				# throughout this process. So it's at initialization state.

				if stateful:
						UpperCamelCase						=  [
						    constraint.copy(stateful=A_ ) for constraint in self.complete_constraints
						]
						if self.inprogress_constraint is not None:
								UpperCamelCase						=  self.inprogress_constraint.copy(stateful=A_ )
						UpperCamelCase						=  [constraint.copy() for constraint in self.pending_constraints]

				return new_state


from typing import Callable, List, Optional, Tuple, Union

import torch
from transformers import CLIPTextModel, CLIPTokenizer

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin, TransformeraDModel, VQModel
from ...schedulers import VQDiffusionScheduler
from ...utils import logging
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput


_UpperCAmelCase      :				str          = logging.get_logger(__name__)  # pylint: disable=invalid-name


class        lowercase      (	_SCREAMING_SNAKE_CASE						,     _SCREAMING_SNAKE_CASE     ):



		@register_to_config
		def __init__(				self       ,   A_       ,   A_ = None       ,   A_ = None )						->      Any:
				"""simple docstring"""
				super().__init__()

				UpperCamelCase						=  learnable

				if self.learnable:
						assert hidden_size is not None, "learnable=True requires `hidden_size` to be set"
						assert length is not None, "learnable=True requires `length` to be set"

						UpperCamelCase						=  torch.zeros(A_       ,   A_ )
				else:
						UpperCamelCase						=  None

				UpperCamelCase						=  torch.nn.Parameter(A_ )






class        lowercase      (	_SCREAMING_SNAKE_CASE     ):
		__lowercase : VQModel
		__lowercase : CLIPTextModel
		__lowercase : CLIPTokenizer
		__lowercase : TransformeraDModel
		__lowercase : LearnedClassifierFreeSamplingEmbeddings
		__lowercase : VQDiffusionScheduler



		def __init__(				self       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   A_       ,   )						->      Optional[Any]:
				"""simple docstring"""
				super().__init__()

				self.register_modules(
				    vqvae=A_       ,   transformer=A_       ,   text_encoder=A_       ,   tokenizer=A_       ,   scheduler=A_       ,   learned_classifier_free_sampling_embeddings=A_       ,   )



		def __UpperCamelCase				(				self       ,   A_       ,   A_       ,   A_ )						->      Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase						=  len(A_ ) if isinstance(A_       ,   A_ ) else 1

				# get prompt text embeddings
				UpperCamelCase						=  self.tokenizer(
				    A_       ,   padding='max_length'       ,   max_length=self.tokenizer.model_max_length       ,   return_tensors='pt'       ,   )
				UpperCamelCase						=  text_inputs.input_ids

				if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
						UpperCamelCase						=  self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :] )
						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}''' )
						UpperCamelCase						=  text_input_ids[:, : self.tokenizer.model_max_length]
				UpperCamelCase						=  self.text_encoder(text_input_ids.to(self.device ) )[0]

				# NOTE: This additional step of normalizing the text embeddings is from VQ-Diffusion.
				# While CLIP does normalize the pooled output of the text transformer when combining
				# the image and text embeddings, CLIP does not directly normalize the last hidden state.
				#
				# CLIP normalizing the pooled output.
				# https://github.com/huggingface/transformers/blob/d92e22d1f28324f513f3080e5c47c071a3916721/src/transformers/models/clip/modeling_clip.py#L1052-L1053
				UpperCamelCase						=  prompt_embeds / prompt_embeds.norm(dim=-1       ,   keepdim=A_ )

				# duplicate text embeddings for each generation per prompt
				UpperCamelCase						=  prompt_embeds.repeat_interleave(A_       ,   dim=0 )

				if do_classifier_free_guidance:
						if self.learned_classifier_free_sampling_embeddings.learnable:
								UpperCamelCase						=  self.learned_classifier_free_sampling_embeddings.embeddings
								UpperCamelCase						=  negative_prompt_embeds.unsqueeze(0 ).repeat(A_       ,   1       ,   1 )
						else:
								UpperCamelCase						=  [''] * batch_size

								UpperCamelCase						=  text_input_ids.shape[-1]
								UpperCamelCase						=  self.tokenizer(
								    A_       ,   padding='max_length'       ,   max_length=A_       ,   truncation=A_       ,   return_tensors='pt'       ,   )
								UpperCamelCase						=  self.text_encoder(uncond_input.input_ids.to(self.device ) )[0]
								# See comment for normalizing text embeddings
								UpperCamelCase						=  negative_prompt_embeds / negative_prompt_embeds.norm(dim=-1       ,   keepdim=A_ )

						# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
						UpperCamelCase						=  negative_prompt_embeds.shape[1]
						UpperCamelCase						=  negative_prompt_embeds.repeat(1       ,   A_       ,   1 )
						UpperCamelCase						=  negative_prompt_embeds.view(batch_size * num_images_per_prompt       ,   A_       ,   -1 )

						# 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
						UpperCamelCase						=  torch.cat([negative_prompt_embeds, prompt_embeds] )

				return prompt_embeds



		@torch.no_grad()
		def __call__(				self       ,   A_       ,   A_ = 100       ,   A_ = 5.0       ,   A_ = 1.0       ,   A_ = 1       ,   A_ = None       ,   A_ = None       ,   A_ = "pil"       ,   A_ = True       ,   A_ = None       ,   A_ = 1       ,   )						->      Union[ImagePipelineOutput, Tuple]:
				"""simple docstring"""
				if isinstance(A_       ,   A_ ):
						UpperCamelCase						=  1
				elif isinstance(A_       ,   A_ ):
						UpperCamelCase						=  len(A_ )
				else:
						raise ValueError(F'''`prompt` has to be of type `str` or `list` but is {type(A_ )}''' )

				UpperCamelCase						=  batch_size * num_images_per_prompt

				UpperCamelCase						=  guidance_scale > 1.0

				UpperCamelCase						=  self._encode_prompt(A_       ,   A_       ,   A_ )

				if (callback_steps is None) or (
				    callback_steps is not None and (not isinstance(A_       ,   A_ ) or callback_steps <= 0)
				):
						raise ValueError(
						    F'''`callback_steps` has to be a positive integer but is {callback_steps} of type'''
						    F''' {type(A_ )}.''' )

				# get the initial completely masked latents unless the user supplied it

				UpperCamelCase						=  (batch_size, self.transformer.num_latent_pixels)
				if latents is None:
						UpperCamelCase						=  self.transformer.num_vector_embeds - 1
						UpperCamelCase						=  torch.full(A_       ,   A_ ).to(self.device )
				else:
						if latents.shape != latents_shape:
								raise ValueError(F'''Unexpected latents shape, got {latents.shape}, expected {latents_shape}''' )
						if (latents < 0).any() or (latents >= self.transformer.num_vector_embeds).any():
								raise ValueError(
								    'Unexpected latents value(s). All latents be valid embedding indices i.e. in the range 0,'
								    F''' {self.transformer.num_vector_embeds - 1} (inclusive).''' )
						UpperCamelCase						=  latents.to(self.device )

				# set timesteps
				self.scheduler.set_timesteps(A_       ,   device=self.device )

				UpperCamelCase						=  self.scheduler.timesteps.to(self.device )

				UpperCamelCase						=  latents

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

						# predict the un-noised image
						# model_output == `log_p_x_0`
						UpperCamelCase						=  self.transformer(A_       ,   encoder_hidden_states=A_       ,   timestep=A_ ).sample

						if do_classifier_free_guidance:
								UpperCamelCase			,						UpperCamelCase						=  model_output.chunk(2 )
								UpperCamelCase						=  model_output_uncond + guidance_scale * (model_output_text - model_output_uncond)
								model_output -= torch.logsumexp(A_       ,   dim=1       ,   keepdim=A_ )

						UpperCamelCase						=  self.truncate(A_       ,   A_ )

						# remove `log(0)`'s (`-inf`s)
						UpperCamelCase						=  model_output.clamp(-70 )

						# compute the previous noisy sample x_t -> x_t-1
						UpperCamelCase						=  self.scheduler.step(A_       ,   timestep=A_       ,   sample=A_       ,   generator=A_ ).prev_sample

						# call the callback, if provided
						if callback is not None and i % callback_steps == 0:
								callback(A_       ,   A_       ,   A_ )

				UpperCamelCase						=  self.vqvae.config.vq_embed_dim
				UpperCamelCase						=  (batch_size, self.transformer.height, self.transformer.width, embedding_channels)
				UpperCamelCase						=  self.vqvae.quantize.get_codebook_entry(A_       ,   shape=A_ )
				UpperCamelCase						=  self.vqvae.decode(A_       ,   force_not_quantize=A_ ).sample

				UpperCamelCase						=  (image / 2 + 0.5).clamp(0       ,   1 )
				UpperCamelCase						=  image.cpu().permute(0       ,   2       ,   3       ,   1 ).numpy()

				if output_type == "pil":
						UpperCamelCase						=  self.numpy_to_pil(A_ )

				if not return_dict:
						return (image,)

				return ImagePipelineOutput(images=A_ )



		def __UpperCamelCase				(				self       ,   A_       ,   A_ )						->      torch.FloatTensor:
				"""simple docstring"""
				UpperCamelCase			,						UpperCamelCase						=  torch.sort(A_       ,   1       ,   descending=A_ )
				UpperCamelCase						=  torch.exp(A_ )
				UpperCamelCase						=  sorted_p_x_0.cumsum(dim=1 ) < truncation_rate

				# Ensure that at least the largest probability is not zeroed out
				UpperCamelCase						=  torch.full_like(keep_mask[:, 0:1, :]       ,   A_ )
				UpperCamelCase						=  torch.cat((all_true, keep_mask)       ,   dim=1 )
				UpperCamelCase						=  keep_mask[:, :-1, :]

				UpperCamelCase						=  keep_mask.gather(1       ,   indices.argsort(1 ) )

				UpperCamelCase						=  log_p_x_0.clone()

				UpperCamelCase						=  -torch.inf  # -inf = log(0)

				return rv
