





'''simple docstring'''







from typing import TYPE_CHECKING

from ...file_utils import _LazyModule, is_tokenizers_available, is_torch_available
from ...utils import OptionalDependencyNotAvailable


snake_case_   :				List[Any]		=     {'configuration_gpt_neox': ['GPT_NEOX_PRETRAINED_CONFIG_ARCHIVE_MAP', 'GPTNeoXConfig']}

try:
					if not is_tokenizers_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					snake_case_   :				List[Any]		=     ['GPTNeoXTokenizerFast']

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					snake_case_   :				Optional[int]		=     [
					    'GPT_NEOX_PRETRAINED_MODEL_ARCHIVE_LIST',
					    'GPTNeoXForCausalLM',
					    'GPTNeoXForQuestionAnswering',
					    'GPTNeoXForSequenceClassification',
					    'GPTNeoXForTokenClassification',
					    'GPTNeoXLayer',
					    'GPTNeoXModel',
					    'GPTNeoXPreTrainedModel',
					]


if TYPE_CHECKING:
					from .configuration_gpt_neox import GPT_NEOX_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTNeoXConfig

					try:
										if not is_tokenizers_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .tokenization_gpt_neox_fast import GPTNeoXTokenizerFast

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_gpt_neox import (
										    GPT_NEOX_PRETRAINED_MODEL_ARCHIVE_LIST,
										    GPTNeoXForCausalLM,
										    GPTNeoXForQuestionAnswering,
										    GPTNeoXForSequenceClassification,
										    GPTNeoXForTokenClassification,
										    GPTNeoXLayer,
										    GPTNeoXModel,
										    GPTNeoXPreTrainedModel,
										)


else:
					import sys

					snake_case_   :				Union[str, Any]		=     _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)







from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block




@dataclass
class        UpperCAmelCase_			( _A     ):


  '''simple docstring'''







  a__           =		42







class        UpperCAmelCase_			( nn.Module     ):


  '''simple docstring'''







  def __init__( self			: Tuple			,							UpperCamelCase__			: List[Any]=3			,							UpperCamelCase__			: str=3			,							UpperCamelCase__			: List[Any]=("DownEncoderBlock2D",)			,							UpperCamelCase__			: Optional[Any]=(64,)			,							UpperCamelCase__			: Optional[Any]=2			,							UpperCamelCase__			: Union[str, Any]=32			,							UpperCamelCase__			: Optional[Any]="silu"			,							UpperCamelCase__			: List[str]=True			,							)			->	str:







    """simple docstring"""






    super().__init__()
    __magic_name__				     =       layers_per_block

    __magic_name__				     =       torch.nn.Convad(
        UpperCamelCase__			,							block_out_channels[0]			,							kernel_size=3			,							stride=1			,							padding=1			,							)

    __magic_name__				     =       None
    __magic_name__				     =       nn.ModuleList([]    )

    # down
    __magic_name__				     =       block_out_channels[0]
    for i, down_block_type in enumerate(UpperCamelCase__    ):
      __magic_name__				     =       output_channel
      __magic_name__				     =       block_out_channels[i]
      __magic_name__				     =       i == len(UpperCamelCase__    ) - 1

      __magic_name__				     =       get_down_block(
          UpperCamelCase__			,							num_layers=self.layers_per_block			,							in_channels=UpperCamelCase__			,							out_channels=UpperCamelCase__			,							add_downsample=not is_final_block			,							resnet_eps=1E-6			,							downsample_padding=0			,							resnet_act_fn=UpperCamelCase__			,							resnet_groups=UpperCamelCase__			,							attention_head_dim=UpperCamelCase__			,							temb_channels=UpperCamelCase__			,							)
      self.down_blocks.append(UpperCamelCase__    )

    # mid
    __magic_name__				     =       UNetMidBlockaD(
        in_channels=block_out_channels[-1]			,							resnet_eps=1E-6			,							resnet_act_fn=UpperCamelCase__			,							output_scale_factor=1			,							resnet_time_scale_shift="""default"""			,							attention_head_dim=block_out_channels[-1]			,							resnet_groups=UpperCamelCase__			,							temb_channels=UpperCamelCase__			,							)

    # out
    __magic_name__				     =       nn.GroupNorm(num_channels=block_out_channels[-1]			,							num_groups=UpperCamelCase__			,							eps=1E-6    )
    __magic_name__				     =       nn.SiLU()

    __magic_name__				     =       2 * out_channels if double_z else out_channels
    __magic_name__				     =       nn.Convad(block_out_channels[-1]			,							UpperCamelCase__			,							3			,							padding=1    )

    __magic_name__				     =       False





  def 							_lowercase							( self			: List[str]			,							UpperCamelCase__			: Optional[Any]    )			->	int:







    """simple docstring"""






    __magic_name__				     =       x
    __magic_name__				     =       self.conv_in(UpperCamelCase__    )

    if self.training and self.gradient_checkpointing:

      def create_custom_forward(UpperCamelCase__			: int    ):
        def custom_forward(*UpperCamelCase__			: str    ):
          return module(*UpperCamelCase__    )

        return custom_forward

      # down
      if is_torch_version(""">="""			,							"""1.11.0"""    ):
        for down_block in self.down_blocks:
          __magic_name__				     =       torch.utils.checkpoint.checkpoint(
              create_custom_forward(UpperCamelCase__    )			,							UpperCamelCase__			,							use_reentrant=UpperCamelCase__    )
        # middle
        __magic_name__				     =       torch.utils.checkpoint.checkpoint(
            create_custom_forward(self.mid_block    )			,							UpperCamelCase__			,							use_reentrant=UpperCamelCase__    )
      else:
        for down_block in self.down_blocks:
          __magic_name__				     =       torch.utils.checkpoint.checkpoint(create_custom_forward(UpperCamelCase__    )			,							UpperCamelCase__    )
        # middle
        __magic_name__				     =       torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block    )			,							UpperCamelCase__    )

    else:
      # down
      for down_block in self.down_blocks:
        __magic_name__				     =       down_block(UpperCamelCase__    )

      # middle
      __magic_name__				     =       self.mid_block(UpperCamelCase__    )

    # post-process
    __magic_name__				     =       self.conv_norm_out(UpperCamelCase__    )
    __magic_name__				     =       self.conv_act(UpperCamelCase__    )
    __magic_name__				     =       self.conv_out(UpperCamelCase__    )

    return sample







class        UpperCAmelCase_			( nn.Module     ):


  '''simple docstring'''







  def __init__( self			: List[str]			,							UpperCamelCase__			: int=3			,							UpperCamelCase__			: Dict=3			,							UpperCamelCase__			: List[Any]=("UpDecoderBlock2D",)			,							UpperCamelCase__			: List[Any]=(64,)			,							UpperCamelCase__			: Dict=2			,							UpperCamelCase__			: int=32			,							UpperCamelCase__			: Optional[int]="silu"			,							UpperCamelCase__			: Tuple="group"			,							)			->	Dict:







    """simple docstring"""






    super().__init__()
    __magic_name__				     =       layers_per_block

    __magic_name__				     =       nn.Convad(
        UpperCamelCase__			,							block_out_channels[-1]			,							kernel_size=3			,							stride=1			,							padding=1			,							)

    __magic_name__				     =       None
    __magic_name__				     =       nn.ModuleList([]    )

    __magic_name__				     =       in_channels if norm_type == """spatial""" else None

    # mid
    __magic_name__				     =       UNetMidBlockaD(
        in_channels=block_out_channels[-1]			,							resnet_eps=1E-6			,							resnet_act_fn=UpperCamelCase__			,							output_scale_factor=1			,							resnet_time_scale_shift="""default""" if norm_type == """group""" else norm_type			,							attention_head_dim=block_out_channels[-1]			,							resnet_groups=UpperCamelCase__			,							temb_channels=UpperCamelCase__			,							)

    # up
    __magic_name__				     =       list(reversed(UpperCamelCase__    )    )
    __magic_name__				     =       reversed_block_out_channels[0]
    for i, up_block_type in enumerate(UpperCamelCase__    ):
      __magic_name__				     =       output_channel
      __magic_name__				     =       reversed_block_out_channels[i]

      __magic_name__				     =       i == len(UpperCamelCase__    ) - 1

      __magic_name__				     =       get_up_block(
          UpperCamelCase__			,							num_layers=self.layers_per_block + 1			,							in_channels=UpperCamelCase__			,							out_channels=UpperCamelCase__			,							prev_output_channel=UpperCamelCase__			,							add_upsample=not is_final_block			,							resnet_eps=1E-6			,							resnet_act_fn=UpperCamelCase__			,							resnet_groups=UpperCamelCase__			,							attention_head_dim=UpperCamelCase__			,							temb_channels=UpperCamelCase__			,							resnet_time_scale_shift=UpperCamelCase__			,							)
      self.up_blocks.append(UpperCamelCase__    )
      __magic_name__				     =       output_channel

    # out
    if norm_type == "spatial":
      __magic_name__				     =       SpatialNorm(block_out_channels[0]			,							UpperCamelCase__    )
    else:
      __magic_name__				     =       nn.GroupNorm(num_channels=block_out_channels[0]			,							num_groups=UpperCamelCase__			,							eps=1E-6    )
    __magic_name__				     =       nn.SiLU()
    __magic_name__				     =       nn.Convad(block_out_channels[0]			,							UpperCamelCase__			,							3			,							padding=1    )

    __magic_name__				     =       False





  def 							_lowercase							( self			: str			,							UpperCamelCase__			: Union[str, Any]			,							UpperCamelCase__			: Tuple=None    )			->	Tuple:







    """simple docstring"""






    __magic_name__				     =       z
    __magic_name__				     =       self.conv_in(UpperCamelCase__    )

    __magic_name__				     =       next(iter(self.up_blocks.parameters()    )    ).dtype
    if self.training and self.gradient_checkpointing:

      def create_custom_forward(UpperCamelCase__			: Optional[int]    ):
        def custom_forward(*UpperCamelCase__			: int    ):
          return module(*UpperCamelCase__    )

        return custom_forward

      if is_torch_version(""">="""			,							"""1.11.0"""    ):
        # middle
        __magic_name__				     =       torch.utils.checkpoint.checkpoint(
            create_custom_forward(self.mid_block    )			,							UpperCamelCase__			,							UpperCamelCase__			,							use_reentrant=UpperCamelCase__    )
        __magic_name__				     =       sample.to(UpperCamelCase__    )

        # up
        for up_block in self.up_blocks:
          __magic_name__				     =       torch.utils.checkpoint.checkpoint(
              create_custom_forward(UpperCamelCase__    )			,							UpperCamelCase__			,							UpperCamelCase__			,							use_reentrant=UpperCamelCase__    )
      else:
        # middle
        __magic_name__				     =       torch.utils.checkpoint.checkpoint(
            create_custom_forward(self.mid_block    )			,							UpperCamelCase__			,							UpperCamelCase__    )
        __magic_name__				     =       sample.to(UpperCamelCase__    )

        # up
        for up_block in self.up_blocks:
          __magic_name__				     =       torch.utils.checkpoint.checkpoint(create_custom_forward(UpperCamelCase__    )			,							UpperCamelCase__			,							UpperCamelCase__    )
    else:
      # middle
      __magic_name__				     =       self.mid_block(UpperCamelCase__			,							UpperCamelCase__    )
      __magic_name__				     =       sample.to(UpperCamelCase__    )

      # up
      for up_block in self.up_blocks:
        __magic_name__				     =       up_block(UpperCamelCase__			,							UpperCamelCase__    )

        # post-process
    if latent_embeds is None:
      __magic_name__				     =       self.conv_norm_out(UpperCamelCase__    )
    else:
      __magic_name__				     =       self.conv_norm_out(UpperCamelCase__			,							UpperCamelCase__    )
    __magic_name__				     =       self.conv_act(UpperCamelCase__    )
    __magic_name__				     =       self.conv_out(UpperCamelCase__    )

    return sample







class        UpperCAmelCase_			( nn.Module     ):


  '''simple docstring'''







  def __init__( self			: str			,							UpperCamelCase__			: Optional[Any]			,							UpperCamelCase__			: Optional[Any]			,							UpperCamelCase__			: List[str]			,							UpperCamelCase__			: Union[str, Any]=None			,							UpperCamelCase__			: Dict="random"			,							UpperCamelCase__			: List[Any]=False			,							UpperCamelCase__			: Dict=True    )			->	Optional[Any]:







    """simple docstring"""






    super().__init__()
    __magic_name__				     =       n_e
    __magic_name__				     =       vq_embed_dim
    __magic_name__				     =       beta
    __magic_name__				     =       legacy

    __magic_name__				     =       nn.Embedding(self.n_e			,							self.vq_embed_dim    )
    self.embedding.weight.data.uniform_(-1.0 / self.n_e			,							1.0 / self.n_e    )

    __magic_name__				     =       remap
    if self.remap is not None:
      self.register_buffer("""used"""			,							torch.tensor(np.load(self.remap    )    )    )
      __magic_name__				     =       self.used.shape[0]
      __magic_name__				     =       unknown_index  # "random" or "extra" or integer
      if self.unknown_index == "extra":
        __magic_name__				     =       self.re_embed
        __magic_name__				     =       self.re_embed + 1
      print(
          F'''Remapping {self.n_e} indices to {self.re_embed} indices. '''
          F'''Using {self.unknown_index} for unknown indices.'''    )
    else:
      __magic_name__				     =       n_e

    __magic_name__				     =       sane_index_shape
  def 							_lowercase							( self			: Union[str, Any]			,							UpperCamelCase__			: Dict    )			->	Union[str, Any]:







    """simple docstring"""






    __magic_name__				     =       inds.shape
    assert len(UpperCamelCase__    ) > 1
    __magic_name__				     =       inds.reshape(ishape[0]			,							-1    )
    __magic_name__				     =       self.used.to(UpperCamelCase__    )
    __magic_name__				     =       (inds[:, :, None] == used[None, None, ...]).long()
    __magic_name__				     =       match.argmax(-1    )
    __magic_name__				     =       match.sum(2    ) < 1
    if self.unknown_index == "random":
      __magic_name__				     =       torch.randint(0			,							self.re_embed			,							size=new[unknown].shape    ).to(device=new.device    )
    else:
      __magic_name__				     =       self.unknown_index
    return new.reshape(UpperCamelCase__    )
  def 							_lowercase							( self			: Optional[Any]			,							UpperCamelCase__			: str    )			->	Tuple:







    """simple docstring"""






    __magic_name__				     =       inds.shape
    assert len(UpperCamelCase__    ) > 1
    __magic_name__				     =       inds.reshape(ishape[0]			,							-1    )
    __magic_name__				     =       self.used.to(UpperCamelCase__    )
    if self.re_embed > self.used.shape[0]:  # extra token
      __magic_name__				     =       0  # simply set to zero
    __magic_name__				     =       torch.gather(used[None, :][inds.shape[0] * [0], :]			,							1			,							UpperCamelCase__    )
    return back.reshape(UpperCamelCase__    )
  def 							_lowercase							( self			: List[str]			,							UpperCamelCase__			: List[str]    )			->	List[str]:







    """simple docstring"""






    __magic_name__				     =       z.permute(0			,							2			,							3			,							1    ).contiguous()
    __magic_name__				     =       z.view(-1			,							self.vq_embed_dim    )

    # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
    __magic_name__				     =       torch.argmin(torch.cdist(UpperCamelCase__			,							self.embedding.weight    )			,							dim=1    )

    __magic_name__				     =       self.embedding(UpperCamelCase__    ).view(z.shape    )
    __magic_name__				     =       None
    __magic_name__				     =       None

    # compute loss for embedding
    if not self.legacy:
      __magic_name__				     =       self.beta * torch.mean((z_q.detach() - z) ** 2    ) + torch.mean((z_q - z.detach()) ** 2    )
    else:
      __magic_name__				     =       torch.mean((z_q.detach() - z) ** 2    ) + self.beta * torch.mean((z_q - z.detach()) ** 2    )

    # preserve gradients
    __magic_name__				     =       z + (z_q - z).detach()

    # reshape back to match original input shape
    __magic_name__				     =       z_q.permute(0			,							3			,							1			,							2    ).contiguous()

    if self.remap is not None:
      __magic_name__				     =       min_encoding_indices.reshape(z.shape[0]			,							-1    )  # add batch axis
      __magic_name__				     =       self.remap_to_used(UpperCamelCase__    )
      __magic_name__				     =       min_encoding_indices.reshape(-1			,							1    )  # flatten

    if self.sane_index_shape:
      __magic_name__				     =       min_encoding_indices.reshape(z_q.shape[0]			,							z_q.shape[2]			,							z_q.shape[3]    )

    return z_q, loss, (perplexity, min_encodings, min_encoding_indices)





  def 							_lowercase							( self			: Union[str, Any]			,							UpperCamelCase__			: List[str]			,							UpperCamelCase__			: Union[str, Any]    )			->	int:







    """simple docstring"""






    if self.remap is not None:
      __magic_name__				     =       indices.reshape(shape[0]			,							-1    )  # add batch axis
      __magic_name__				     =       self.unmap_to_all(UpperCamelCase__    )
      __magic_name__				     =       indices.reshape(-1    )  # flatten again

    # get quantized latent vectors
    __magic_name__				     =       self.embedding(UpperCamelCase__    )

    if shape is not None:
      __magic_name__				     =       z_q.view(UpperCamelCase__    )
      # reshape back to match original input shape
      __magic_name__				     =       z_q.permute(0			,							3			,							1			,							2    ).contiguous()

    return z_q







class        UpperCAmelCase_			( _A     ):


  '''simple docstring'''







  def __init__( self			: Dict			,							UpperCamelCase__			: Optional[Any]			,							UpperCamelCase__			: Tuple=False    )			->	Optional[int]:







    """simple docstring"""






    __magic_name__				     =       parameters
    __magic_name__     ,    __magic_name__				     =       torch.chunk(UpperCamelCase__			,							2			,							dim=1    )
    __magic_name__				     =       torch.clamp(self.logvar			,							-30.0			,							20.0    )
    __magic_name__				     =       deterministic
    __magic_name__				     =       torch.exp(0.5 * self.logvar    )
    __magic_name__				     =       torch.exp(self.logvar    )
    if self.deterministic:
      __magic_name__				     =       __magic_name__				     =       torch.zeros_like(
          self.mean			,							device=self.parameters.device			,							dtype=self.parameters.dtype    )
  def 							_lowercase							( self			: Tuple			,							UpperCamelCase__			: Optional[torch.Generator] = None    )			->	torch.FloatTensor:







    """simple docstring"""






    __magic_name__				     =       randn_tensor(
        self.mean.shape			,							generator=UpperCamelCase__			,							device=self.parameters.device			,							dtype=self.parameters.dtype    )
    __magic_name__				     =       self.mean + self.std * sample
    return x
  def 							_lowercase							( self			: Dict			,							UpperCamelCase__			: Optional[int]=None    )			->	Any:







    """simple docstring"""






    if self.deterministic:
      return torch.Tensor([0.0]    )
    else:
      if other is None:
        return 0.5 * torch.sum(torch.pow(self.mean			,							2    ) + self.var - 1.0 - self.logvar			,							dim=[1, 2, 3]    )
      else:
        return 0.5 * torch.sum(
            torch.pow(self.mean - other.mean			,							2    ) / other.var
            + self.var / other.var
            - 1.0
            - self.logvar
            + other.logvar			,							dim=[1, 2, 3]			,							)
  def 							_lowercase							( self			: Optional[int]			,							UpperCamelCase__			: List[Any]			,							UpperCamelCase__			: Dict=[1, 2, 3]    )			->	Optional[int]:







    """simple docstring"""






    if self.deterministic:
      return torch.Tensor([0.0]    )
    __magic_name__				     =       np.log(2.0 * np.pi    )
    return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean			,							2    ) / self.var			,							dim=UpperCamelCase__    )





  def 							_lowercase							( self			: Optional[int]    )			->	Optional[int]:







    """simple docstring"""






    return self.mean

