


'''simple docstring'''






import unittest
from unittest.util import safe_repr

from transformers import AutoTokenizer, RwkvConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import (
						    RWKV_PRETRAINED_MODEL_ARCHIVE_LIST,
						    RwkvForCausalLM,
						    RwkvModel,
						)
						from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_0
else:
						_A      =     False






class      SCREAMING_SNAKE_CASE_	:




			def __init__(      self					,      lowercase					,      lowercase=1_4					,      lowercase=7					,      lowercase=True					,      lowercase=False					,      lowercase=True					,      lowercase=True					,      lowercase=True					,      lowercase=9_9					,      lowercase=3_2					,      lowercase=5					,      lowercase=3_7					,      lowercase="gelu"					,      lowercase=0.1					,      lowercase=0.1					,      lowercase=5_1_2					,      lowercase=1_6					,      lowercase=2					,      lowercase=3					,      lowercase=4					,      lowercase=None					,      )			->						int:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = parent
								__SCREAMING_SNAKE_CASE					:							int			   = batch_size
								__SCREAMING_SNAKE_CASE					:							str			   = seq_length
								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = is_training
								__SCREAMING_SNAKE_CASE					:							Tuple			   = use_token_type_ids
								__SCREAMING_SNAKE_CASE					:							str			   = use_input_mask
								__SCREAMING_SNAKE_CASE					:							Dict			   = use_labels
								__SCREAMING_SNAKE_CASE					:							Any			   = use_mc_token_ids
								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = vocab_size
								__SCREAMING_SNAKE_CASE					:							Tuple			   = hidden_size
								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = num_hidden_layers
								__SCREAMING_SNAKE_CASE					:							int			   = intermediate_size
								__SCREAMING_SNAKE_CASE					:							int			   = hidden_act
								__SCREAMING_SNAKE_CASE					:							Tuple			   = hidden_dropout_prob
								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = attention_probs_dropout_prob
								__SCREAMING_SNAKE_CASE					:							int			   = max_position_embeddings
								__SCREAMING_SNAKE_CASE					:							int			   = type_vocab_size
								__SCREAMING_SNAKE_CASE					:							Any			   = type_sequence_label_size
								__SCREAMING_SNAKE_CASE					:							Dict			   = num_labels
								__SCREAMING_SNAKE_CASE					:							List[Any]			   = num_choices
								__SCREAMING_SNAKE_CASE					:							Dict			   = scope
								__SCREAMING_SNAKE_CASE					:							Any			   = vocab_size - 1
								__SCREAMING_SNAKE_CASE					:							List[Any]			   = vocab_size - 1
								__SCREAMING_SNAKE_CASE					:							List[Any]			   = vocab_size - 1




			def 			_snake_case   (      self       )			->						Optional[Any]:


								'''simple docstring'''






								return RwkvConfig.from_pretrained('''sgugger/rwkv-4-pile-7b'''       )




			def 			_snake_case   (      self					,      lowercase=False					,      lowercase=False					,      lowercase=False       )			->						Optional[Any]:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = ids_tensor([self.batch_size, self.seq_length]					,      self.vocab_size       )

								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = None
								if self.use_input_mask:
													__SCREAMING_SNAKE_CASE					:							List[Any]			   = random_attention_mask([self.batch_size, self.seq_length]       )

								__SCREAMING_SNAKE_CASE					:							Dict			   = None
								if self.use_token_type_ids:
													__SCREAMING_SNAKE_CASE					:							Optional[int]			   = ids_tensor([self.batch_size, self.seq_length]					,      self.type_vocab_size       )

								__SCREAMING_SNAKE_CASE					:							int			   = None
								if self.use_mc_token_ids:
													__SCREAMING_SNAKE_CASE					:							int			   = ids_tensor([self.batch_size, self.num_choices]					,      self.seq_length       )

								__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = None
								__SCREAMING_SNAKE_CASE					:							Any			   = None
								__SCREAMING_SNAKE_CASE					:							Dict			   = None
								if self.use_labels:
													__SCREAMING_SNAKE_CASE					:							Tuple			   = ids_tensor([self.batch_size]					,      self.type_sequence_label_size       )
													__SCREAMING_SNAKE_CASE					:							List[str]			   = ids_tensor([self.batch_size, self.seq_length]					,      self.num_labels       )
													__SCREAMING_SNAKE_CASE					:							str			   = ids_tensor([self.batch_size]					,      self.num_choices       )

								__SCREAMING_SNAKE_CASE					:							List[str]			   = self.get_config(
								    gradient_checkpointing=lowercase					,      scale_attn_by_inverse_layer_idx=lowercase					,      reorder_and_upcast_attn=lowercase					,      )

								return (
								    config,
								    input_ids,
								    input_mask,
								    None,
								    token_type_ids,
								    mc_token_ids,
								    sequence_labels,
								    token_labels,
								    choice_labels,
								)




			def 			_snake_case   (      self					,      lowercase=False					,      lowercase=False					,      lowercase=False       )			->						int:


								'''simple docstring'''






								return RwkvConfig(
								    vocab_size=self.vocab_size					,      hidden_size=self.hidden_size					,      num_hidden_layers=self.num_hidden_layers					,      intermediate_size=self.intermediate_size					,      activation_function=self.hidden_act					,      resid_pdrop=self.hidden_dropout_prob					,      attn_pdrop=self.attention_probs_dropout_prob					,      n_positions=self.max_position_embeddings					,      type_vocab_size=self.type_vocab_size					,      use_cache=lowercase					,      bos_token_id=self.bos_token_id					,      eos_token_id=self.eos_token_id					,      pad_token_id=self.pad_token_id					,      gradient_checkpointing=lowercase					,      scale_attn_by_inverse_layer_idx=lowercase					,      reorder_and_upcast_attn=lowercase					,      )




			def 			_snake_case   (      self       )			->						int:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Dict			   = self.get_config()
								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = 3_0_0
								return config




			def 			_snake_case   (      self       )			->						Dict:


								'''simple docstring'''






								(
								    (
								    __SCREAMING_SNAKE_CASE
								)   ,   (
								    __SCREAMING_SNAKE_CASE
								)   ,   (
								    __SCREAMING_SNAKE_CASE
								)   ,   (
								    __SCREAMING_SNAKE_CASE
								)   ,   (
								    __SCREAMING_SNAKE_CASE
								)   ,   (
								    __SCREAMING_SNAKE_CASE
								)   ,   (
								    __SCREAMING_SNAKE_CASE
								)   ,   (
								    __SCREAMING_SNAKE_CASE
								)   ,   (
								    __SCREAMING_SNAKE_CASE
								)   ,   
								)					:							int			   = self.prepare_config_and_inputs()

								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
								__SCREAMING_SNAKE_CASE					:							Any			   = ids_tensor([self.batch_size, self.seq_length]					,      vocab_size=2       )

								return (
								    config,
								    input_ids,
								    input_mask,
								    head_mask,
								    token_type_ids,
								    sequence_labels,
								    token_labels,
								    choice_labels,
								    encoder_hidden_states,
								    encoder_attention_mask,
								)




			def 			_snake_case   (      self					,      lowercase					,      lowercase					,      lowercase					,      lowercase					,      lowercase					,      *lowercase       )			->						Dict:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = True
								__SCREAMING_SNAKE_CASE					:							List[str]			   = RwkvModel(config=lowercase       )
								model.to(lowercase       )
								model.eval()

								__SCREAMING_SNAKE_CASE					:							Dict			   = model(lowercase       )

								self.parent.assertEqual(result.last_hidden_state.shape					,      (self.batch_size, self.seq_length, self.hidden_size)       )
								self.parent.assertEqual(len(result.hidden_states       )					,      config.num_hidden_layers + 1       )




			def 			_snake_case   (      self					,      lowercase					,      lowercase					,      lowercase					,      lowercase					,      lowercase					,      *lowercase       )			->						Union[str, Any]:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Dict			   = RwkvForCausalLM(lowercase       )
								model.to(lowercase       )
								model.eval()

								__SCREAMING_SNAKE_CASE					:							List[Any]			   = model(lowercase					,      labels=lowercase       )
								self.parent.assertEqual(result.loss.shape					,      ()       )
								self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.seq_length, self.vocab_size)       )




			def 			_snake_case   (      self					,      lowercase					,      lowercase					,      lowercase					,      lowercase					,      lowercase					,      *lowercase       )			->						List[str]:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							str			   = RwkvModel(config=lowercase       )
								model.to(lowercase       )
								model.eval()

								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = model(lowercase       )
								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = outputs.last_hidden_state

								__SCREAMING_SNAKE_CASE					:							Any			   = model(input_ids[:, :2]       )
								__SCREAMING_SNAKE_CASE					:							Any			   = outputs.last_hidden_state

								# Using the state computed on the first inputs, we will get the same output
								__SCREAMING_SNAKE_CASE					:							int			   = model(input_ids[:, 2:]					,      state=outputs.state       )
								__SCREAMING_SNAKE_CASE					:							int			   = outputs.last_hidden_state

								self.parent.assertTrue(torch.allclose(torch.cat([output_one, output_two]					,      dim=1       )					,      lowercase					,      atol=1e-5       )       )




			def 			_snake_case   (      self					,      lowercase					,      lowercase					,      lowercase					,      lowercase					,      lowercase					,      *lowercase					,      lowercase=False       )			->						Tuple:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Any			   = RwkvForCausalLM(lowercase       )
								model.to(lowercase       )
								if gradient_checkpointing:
													model.gradient_checkpointing_enable()

								__SCREAMING_SNAKE_CASE					:							List[str]			   = model(lowercase					,      labels=lowercase       )
								self.parent.assertEqual(result.loss.shape					,      ()       )
								self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.seq_length, self.vocab_size)       )
								result.loss.backward()




			def 			_snake_case   (      self       )			->						int:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Tuple			   = self.prepare_config_and_inputs()

								(
								    (
								    __SCREAMING_SNAKE_CASE
								)   ,   (
								    __SCREAMING_SNAKE_CASE
								)   ,   (
								    __SCREAMING_SNAKE_CASE
								)   ,   (
								    __SCREAMING_SNAKE_CASE
								)   ,   (
								    __SCREAMING_SNAKE_CASE
								)   ,   (
								    __SCREAMING_SNAKE_CASE
								)   ,   (
								    __SCREAMING_SNAKE_CASE
								)   ,   (
								    __SCREAMING_SNAKE_CASE
								)   ,   (
								    __SCREAMING_SNAKE_CASE
								)   ,   
								)					:							int			   = config_and_inputs

								__SCREAMING_SNAKE_CASE					:							List[str]			   = {'''input_ids''': input_ids}

								return config, inputs_dict






@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0					,					reason='''See https://github.com/huggingface/transformers/pull/24204'''		)
@require_torch
class      SCREAMING_SNAKE_CASE_	( snake_case					,					snake_case					,					snake_case					,					unittest.TestCase		):
			__a					:      Optional[Any]					=  (RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
			__a					:      List[Any]					=  (
			    {'''feature-extraction''': RwkvModel, '''text-generation''': RwkvForCausalLM} if is_torch_available() else {}
			)
			# all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
			__a					:      Optional[int]					=  False
			__a					:      Any					=  False
			__a					:      Dict					=  False
			__a					:      Union[str, Any]					=  False
			__a					:      int					=  False  # Rwkv does not support head masking




			def 			_snake_case   (      self       )			->						str:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							str			   = RwkvModelTester(self       )
								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = ConfigTester(
								    self					,      config_class=lowercase					,      n_embd=3_7					,      common_properties=['''hidden_size''', '''num_hidden_layers''']       )




			def 			_snake_case   (      self					,      lowercase					,      lowercase					,      lowercase=None       )			->						str:


								'''simple docstring'''






								if isinstance(lowercase					,      torch.Tensor       ):
													__SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE					:							str			   = member.max().item(), member.min().item()
								elif isinstance(lowercase					,      lowercase       ) or isinstance(lowercase					,      lowercase       ):
													__SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE					:							Optional[int]			   = max(lowercase       ), min(lowercase       )

								if not isinstance(lowercase					,      lowercase       ):
													raise TypeError('''container should be a list or tuple'''       )
								elif len(lowercase       ) != 2:
													raise ValueError('''container should have 2 elements'''       )

								__SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE					:							Optional[int]			   = container

								__SCREAMING_SNAKE_CASE					:							Any			   = (min_value >= expected_min) and (max_value <= expected_max)

								if not is_inside_interval:
													__SCREAMING_SNAKE_CASE					:							Any			   = '''%s not found in %s''' % (safe_repr(lowercase       ), safe_repr(lowercase       ))
													self.fail(self._formatMessage(lowercase					,      lowercase       )       )




			def 			_snake_case   (      self       )			->						Union[str, Any]:


								'''simple docstring'''






								self.config_tester.run_common_tests()




			def 			_snake_case   (      self       )			->						int:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_rwkv_model(*lowercase       )




			def 			_snake_case   (      self       )			->						Any:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_causl_lm(*lowercase       )




			def 			_snake_case   (      self       )			->						Dict:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							List[Any]			   = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_state_equivalency(*lowercase       )




			def 			_snake_case   (      self       )			->						List[Any]:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
													__SCREAMING_SNAKE_CASE					:							int			   = model_class(config=lowercase       )
													for name, param in model.named_parameters():
																		if "time_decay" in name:
																							if param.requires_grad:
																												self.assertTrue(param.data.max().item() == 3.0       )
																												self.assertTrue(param.data.min().item() == -5.0       )
																		elif "time_first" in name:
																							if param.requires_grad:
																												# check if it's a ones like
																												self.assertTrue(torch.allclose(param.data					,      torch.ones_like(param.data       )					,      atol=1e-5					,      rtol=1e-5       )       )
																		elif any(x in name for x in ['''time_mix_key''', '''time_mix_receptance''']       ):
																							if param.requires_grad:
																												self.assertInterval(
																												    param.data					,      [0.0, 1.0]					,      msg=f"""Parameter {name} of model {model_class} seems not properly initialized"""					,      )
																		elif "time_mix_value" in name:
																							if param.requires_grad:
																												self.assertInterval(
																												    param.data					,      [0.0, 1.3]					,      msg=f"""Parameter {name} of model {model_class} seems not properly initialized"""					,      )




			def 			_snake_case   (      self       )			->						Any:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE					:							Optional[Any]			   = self.model_tester.prepare_config_and_inputs_for_common()
								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = True

								__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = getattr(self.model_tester					,      '''seq_length'''					,      lowercase       )

								for model_class in self.all_model_classes:
													__SCREAMING_SNAKE_CASE					:							Optional[int]			   = True
													__SCREAMING_SNAKE_CASE					:							str			   = False
													__SCREAMING_SNAKE_CASE					:							str			   = True
													__SCREAMING_SNAKE_CASE					:							Dict			   = model_class(lowercase       )
													model.to(lowercase       )
													model.eval()

													__SCREAMING_SNAKE_CASE					:							List[Any]			   = self._prepare_for_class(lowercase					,      lowercase       )
													__SCREAMING_SNAKE_CASE					:							List[str]			   = inputs['''input_ids'''].shape[0]
													with torch.no_grad():
																		__SCREAMING_SNAKE_CASE					:							Dict			   = model(**lowercase       )
													__SCREAMING_SNAKE_CASE					:							Dict			   = outputs.attentions
													self.assertEqual(len(lowercase       )					,      self.model_tester.num_hidden_layers       )

													# check that output_attentions also work using config
													del inputs_dict["output_attentions"]
													__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = True
													__SCREAMING_SNAKE_CASE					:							str			   = model_class(lowercase       )
													model.to(lowercase       )
													model.eval()

													__SCREAMING_SNAKE_CASE					:							Dict			   = self._prepare_for_class(lowercase					,      lowercase       )
													__SCREAMING_SNAKE_CASE					:							List[str]			   = inputs['''input_ids'''].shape[0]
													with torch.no_grad():
																		__SCREAMING_SNAKE_CASE					:							Any			   = model(**lowercase       )
													__SCREAMING_SNAKE_CASE					:							Optional[int]			   = outputs.attentions
													self.assertEqual(len(lowercase       )					,      self.model_tester.num_hidden_layers       )

													self.assertListEqual(
													    list(attentions[0].shape[-3:]       )					,      [batch_size, seq_len, config.hidden_size]					,      )
													__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = len(lowercase       )

													# Check attention is always last and order is fine
													__SCREAMING_SNAKE_CASE					:							List[Any]			   = True
													__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = True
													__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = model_class(lowercase       )
													model.to(lowercase       )
													model.eval()

													__SCREAMING_SNAKE_CASE					:							int			   = self._prepare_for_class(lowercase					,      lowercase       )
													__SCREAMING_SNAKE_CASE					:							List[str]			   = inputs['''input_ids'''].shape[0]
													with torch.no_grad():
																		__SCREAMING_SNAKE_CASE					:							str			   = model(**lowercase       )

													__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = 1
													self.assertEqual(out_len + added_hidden_states					,      len(lowercase       )       )

													__SCREAMING_SNAKE_CASE					:							Tuple			   = outputs.attentions

													self.assertEqual(len(lowercase       )					,      self.model_tester.num_hidden_layers       )
													self.assertListEqual(
													    list(self_attentions[0].shape[-3:]       )					,      [batch_size, seq_len, config.hidden_size]					,      )




			@slow
			def 			_snake_case   (      self       )			->						int:


								'''simple docstring'''






								for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													__SCREAMING_SNAKE_CASE					:							Dict			   = RwkvModel.from_pretrained(lowercase       )
													self.assertIsNotNone(lowercase       )






@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0					,					reason='''See https://github.com/huggingface/transformers/pull/24204'''		)
@slow
class      SCREAMING_SNAKE_CASE_	( unittest.TestCase		):




			def 			_snake_case   (      self       )			->						List[str]:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Dict			   = '''RWKV/rwkv-4-169m-pile'''
								__SCREAMING_SNAKE_CASE					:							str			   = AutoTokenizer.from_pretrained(self.model_id       )




			def 			_snake_case   (      self       )			->						Tuple:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							int			   = '''Hello my name is Jasmine and I am a newbie to the'''
								__SCREAMING_SNAKE_CASE					:							str			   = RwkvForCausalLM.from_pretrained(self.model_id       ).to(lowercase       )

								__SCREAMING_SNAKE_CASE					:							List[str]			   = self.tokenizer('''Hello my name is'''					,      return_tensors='''pt'''       ).input_ids.to(lowercase       )
								__SCREAMING_SNAKE_CASE					:							List[str]			   = model.generate(lowercase					,      max_new_tokens=1_0       )
								__SCREAMING_SNAKE_CASE					:							List[Any]			   = self.tokenizer.decode(output[0].tolist()       )

								self.assertEqual(lowercase					,      lowercase       )




			def 			_snake_case   (      self       )			->						Tuple:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = '''Hello my name is Jasmine and I am a newbie to the'''

								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = self.tokenizer('''Hello my name is'''					,      return_tensors='''pt'''       ).input_ids.to(lowercase       )
								__SCREAMING_SNAKE_CASE					:							List[str]			   = RwkvForCausalLM.from_pretrained(self.model_id					,      torch_dtype=torch.bfloataa       ).to(lowercase       )

								__SCREAMING_SNAKE_CASE					:							List[Any]			   = model.generate(lowercase					,      max_new_tokens=1_0       )
								__SCREAMING_SNAKE_CASE					:							List[Any]			   = self.tokenizer.decode(output[0].tolist()       )

								self.assertEqual(lowercase					,      lowercase       )






'''simple docstring'''






from dataclasses import dataclass
from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..models.embeddings import ImagePositionalEmbeddings
from ..utils import BaseOutput, deprecate
from .attention import BasicTransformerBlock
from .embeddings import PatchEmbed
from .modeling_utils import ModelMixin






@dataclass
class      SCREAMING_SNAKE_CASE_	( snake_case		):
			__a					:      torch.FloatTensor






class      SCREAMING_SNAKE_CASE_	( snake_case					,					snake_case		):




			@register_to_config
			def __init__(      self					,      lowercase = 1_6					,      lowercase = 8_8					,      lowercase = None					,      lowercase = None					,      lowercase = 1					,      lowercase = 0.0					,      lowercase = 3_2					,      lowercase = None					,      lowercase = False					,      lowercase = None					,      lowercase = None					,      lowercase = None					,      lowercase = "geglu"					,      lowercase = None					,      lowercase = False					,      lowercase = False					,      lowercase = False					,      lowercase = "layer_norm"					,      lowercase = True					,      )			->						Optional[int]:


								'''simple docstring'''






								super().__init__()
								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = use_linear_projection
								__SCREAMING_SNAKE_CASE					:							Dict			   = num_attention_heads
								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = attention_head_dim
								__SCREAMING_SNAKE_CASE					:							Any			   = num_attention_heads * attention_head_dim

								# 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
								# Define whether input is continuous or discrete depending on configuration
								__SCREAMING_SNAKE_CASE					:							List[Any]			   = (in_channels is not None) and (patch_size is None)
								__SCREAMING_SNAKE_CASE					:							List[str]			   = num_vector_embeds is not None
								__SCREAMING_SNAKE_CASE					:							Any			   = in_channels is not None and patch_size is not None

								if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
													__SCREAMING_SNAKE_CASE					:							Any			   = (
													    f"""The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"""
													    ''' incorrectly set to `\'layer_norm\'`.Make sure to set `norm_type` to `\'ada_norm\'` in the config.'''
													    ''' Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect'''
													    ''' results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it'''
													    ''' would be very nice if you could open a Pull request for the `transformer/config.json` file'''
													)
													deprecate('''norm_type!=num_embeds_ada_norm'''					,      '''1.0.0'''					,      lowercase					,      standard_warn=lowercase       )
													__SCREAMING_SNAKE_CASE					:							str			   = '''ada_norm'''

								if self.is_input_continuous and self.is_input_vectorized:
													raise ValueError(
													    f"""Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"""
													    ''' sure that either `in_channels` or `num_vector_embeds` is None.'''       )
								elif self.is_input_vectorized and self.is_input_patches:
													raise ValueError(
													    f"""Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make"""
													    ''' sure that either `num_vector_embeds` or `num_patches` is None.'''       )
								elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
													raise ValueError(
													    f"""Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:"""
													    f""" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None."""       )

								# 2. Define input layers
								if self.is_input_continuous:
													__SCREAMING_SNAKE_CASE					:							List[Any]			   = in_channels

													__SCREAMING_SNAKE_CASE					:							Optional[int]			   = torch.nn.GroupNorm(num_groups=lowercase					,      num_channels=lowercase					,      eps=1e-6					,      affine=lowercase       )
													if use_linear_projection:
																		__SCREAMING_SNAKE_CASE					:							List[Any]			   = nn.Linear(lowercase					,      lowercase       )
													else:
																		__SCREAMING_SNAKE_CASE					:							Any			   = nn.Convad(lowercase					,      lowercase					,      kernel_size=1					,      stride=1					,      padding=0       )
								elif self.is_input_vectorized:
													assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
													assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"

													__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = sample_size
													__SCREAMING_SNAKE_CASE					:							str			   = sample_size
													__SCREAMING_SNAKE_CASE					:							int			   = num_vector_embeds
													__SCREAMING_SNAKE_CASE					:							str			   = self.height * self.width

													__SCREAMING_SNAKE_CASE					:							int			   = ImagePositionalEmbeddings(
													    num_embed=lowercase					,      embed_dim=lowercase					,      height=self.height					,      width=self.width       )
								elif self.is_input_patches:
													assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"

													__SCREAMING_SNAKE_CASE					:							Any			   = sample_size
													__SCREAMING_SNAKE_CASE					:							Any			   = sample_size

													__SCREAMING_SNAKE_CASE					:							int			   = patch_size
													__SCREAMING_SNAKE_CASE					:							Any			   = PatchEmbed(
													    height=lowercase					,      width=lowercase					,      patch_size=lowercase					,      in_channels=lowercase					,      embed_dim=lowercase					,      )

								# 3. Define transformers blocks
								__SCREAMING_SNAKE_CASE					:							Tuple			   = nn.ModuleList(
								    [
								        BasicTransformerBlock(
								            lowercase					,      lowercase					,      lowercase					,      dropout=lowercase					,      cross_attention_dim=lowercase					,      activation_fn=lowercase					,      num_embeds_ada_norm=lowercase					,      attention_bias=lowercase					,      only_cross_attention=lowercase					,      upcast_attention=lowercase					,      norm_type=lowercase					,      norm_elementwise_affine=lowercase					,      )
								        for d in range(lowercase       )
								    ]       )

								# 4. Define output layers
								__SCREAMING_SNAKE_CASE					:							List[str]			   = in_channels if out_channels is None else out_channels
								if self.is_input_continuous:
													# TODO: should use out_channels for continuous projections
													if use_linear_projection:
																		__SCREAMING_SNAKE_CASE					:							str			   = nn.Linear(lowercase					,      lowercase       )
													else:
																		__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = nn.Convad(lowercase					,      lowercase					,      kernel_size=1					,      stride=1					,      padding=0       )
								elif self.is_input_vectorized:
													__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = nn.LayerNorm(lowercase       )
													__SCREAMING_SNAKE_CASE					:							int			   = nn.Linear(lowercase					,      self.num_vector_embeds - 1       )
								elif self.is_input_patches:
													__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = nn.LayerNorm(lowercase					,      elementwise_affine=lowercase					,      eps=1e-6       )
													__SCREAMING_SNAKE_CASE					:							Any			   = nn.Linear(lowercase					,      2 * inner_dim       )
													__SCREAMING_SNAKE_CASE					:							Any			   = nn.Linear(lowercase					,      patch_size * patch_size * self.out_channels       )




			def 			_snake_case   (      self					,      lowercase					,      lowercase = None					,      lowercase = None					,      lowercase = None					,      lowercase = None					,      lowercase = None					,      lowercase = None					,      lowercase = True					,      )			->						Optional[Any]:


								'''simple docstring'''






								if attention_mask is not None and attention_mask.ndim == 2:
													# assume that mask is expressed as:
													#   (1 = keep,      0 = discard)
													# convert mask into a bias that can be added to attention scores:
													#       (keep = +0,     discard = -10000.0)
													__SCREAMING_SNAKE_CASE					:							Optional[int]			   = (1 - attention_mask.to(hidden_states.dtype       )) * -1_0_0_0_0.0
													__SCREAMING_SNAKE_CASE					:							Dict			   = attention_mask.unsqueeze(1       )

								# convert encoder_attention_mask to a bias the same way we do for attention_mask
								if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
													__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = (1 - encoder_attention_mask.to(hidden_states.dtype       )) * -1_0_0_0_0.0
													__SCREAMING_SNAKE_CASE					:							List[Any]			   = encoder_attention_mask.unsqueeze(1       )

								# 1. Input
								if self.is_input_continuous:
													__SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE					:							Optional[Any]			   = hidden_states.shape
													__SCREAMING_SNAKE_CASE					:							Dict			   = hidden_states

													__SCREAMING_SNAKE_CASE					:							str			   = self.norm(lowercase       )
													if not self.use_linear_projection:
																		__SCREAMING_SNAKE_CASE					:							int			   = self.proj_in(lowercase       )
																		__SCREAMING_SNAKE_CASE					:							List[str]			   = hidden_states.shape[1]
																		__SCREAMING_SNAKE_CASE					:							List[str]			   = hidden_states.permute(0					,      2					,      3					,      1       ).reshape(lowercase					,      height * width					,      lowercase       )
													else:
																		__SCREAMING_SNAKE_CASE					:							Tuple			   = hidden_states.shape[1]
																		__SCREAMING_SNAKE_CASE					:							Optional[int]			   = hidden_states.permute(0					,      2					,      3					,      1       ).reshape(lowercase					,      height * width					,      lowercase       )
																		__SCREAMING_SNAKE_CASE					:							Optional[int]			   = self.proj_in(lowercase       )
								elif self.is_input_vectorized:
													__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = self.latent_image_embedding(lowercase       )
								elif self.is_input_patches:
													__SCREAMING_SNAKE_CASE					:							int			   = self.pos_embed(lowercase       )

								# 2. Blocks
								for block in self.transformer_blocks:
													__SCREAMING_SNAKE_CASE					:							Tuple			   = block(
													    lowercase					,      attention_mask=lowercase					,      encoder_hidden_states=lowercase					,      encoder_attention_mask=lowercase					,      timestep=lowercase					,      cross_attention_kwargs=lowercase					,      class_labels=lowercase					,      )

								# 3. Output
								if self.is_input_continuous:
													if not self.use_linear_projection:
																		__SCREAMING_SNAKE_CASE					:							Tuple			   = hidden_states.reshape(lowercase					,      lowercase					,      lowercase					,      lowercase       ).permute(0					,      3					,      1					,      2       ).contiguous()
																		__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = self.proj_out(lowercase       )
													else:
																		__SCREAMING_SNAKE_CASE					:							str			   = self.proj_out(lowercase       )
																		__SCREAMING_SNAKE_CASE					:							str			   = hidden_states.reshape(lowercase					,      lowercase					,      lowercase					,      lowercase       ).permute(0					,      3					,      1					,      2       ).contiguous()

													__SCREAMING_SNAKE_CASE					:							str			   = hidden_states + residual
								elif self.is_input_vectorized:
													__SCREAMING_SNAKE_CASE					:							Dict			   = self.norm_out(lowercase       )
													__SCREAMING_SNAKE_CASE					:							str			   = self.out(lowercase       )
													# (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
													__SCREAMING_SNAKE_CASE					:							Optional[int]			   = logits.permute(0					,      2					,      1       )

													# log(p(x_0))
													__SCREAMING_SNAKE_CASE					:							int			   = F.log_softmax(logits.double()					,      dim=1       ).float()
								elif self.is_input_patches:
													# TODO: cleanup!
													__SCREAMING_SNAKE_CASE					:							Optional[int]			   = self.transformer_blocks[0].norma.emb(
													    lowercase					,      lowercase					,      hidden_dtype=hidden_states.dtype       )
													__SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE					:							Optional[int]			   = self.proj_out_a(F.silu(lowercase       )       ).chunk(2					,      dim=1       )
													__SCREAMING_SNAKE_CASE					:							int			   = self.norm_out(lowercase       ) * (1 + scale[:, None]) + shift[:, None]
													__SCREAMING_SNAKE_CASE					:							int			   = self.proj_out_a(lowercase       )

													# unpatchify
													__SCREAMING_SNAKE_CASE					:							List[str]			   = int(hidden_states.shape[1] ** 0.5       )
													__SCREAMING_SNAKE_CASE					:							int			   = hidden_states.reshape(
													    shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)       )
													__SCREAMING_SNAKE_CASE					:							int			   = torch.einsum('''nhwpqc->nchpwq'''					,      lowercase       )
													__SCREAMING_SNAKE_CASE					:							Any			   = hidden_states.reshape(
													    shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)       )

								if not return_dict:
													return (output,)

								return TransformeraDModelOutput(sample=lowercase       )


