






import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import (
		    MODEL_FOR_PRETRAINING_MAPPING,
		    MegatronBertForCausalLM,
		    MegatronBertForMaskedLM,
		    MegatronBertForMultipleChoice,
		    MegatronBertForNextSentencePrediction,
		    MegatronBertForPreTraining,
		    MegatronBertForQuestionAnswering,
		    MegatronBertForSequenceClassification,
		    MegatronBertForTokenClassification,
		    MegatronBertModel,
		)



class   lowerCamelCase      :







		'''simple docstring'''



		def __init__(							self				,		_UpperCamelCase				,		_UpperCamelCase=1_3				,		_UpperCamelCase=7				,		_UpperCamelCase=True				,		_UpperCamelCase=True				,		_UpperCamelCase=True				,		_UpperCamelCase=True				,		_UpperCamelCase=9_9				,		_UpperCamelCase=6_4				,		_UpperCamelCase=3_2				,		_UpperCamelCase=5				,		_UpperCamelCase=4				,		_UpperCamelCase=3_7				,		_UpperCamelCase="gelu"				,		_UpperCamelCase=0.1				,		_UpperCamelCase=0.1				,		_UpperCamelCase=5_1_2				,		_UpperCamelCase=1_6				,		_UpperCamelCase=2				,		_UpperCamelCase=0.02				,		_UpperCamelCase=3				,		_UpperCamelCase=4				,		_UpperCamelCase=None				,		)	->							Tuple:
							UpperCAmelCase_    :      Union[str, Any]            =						parent
							UpperCAmelCase_    :      Tuple            =						batch_size
							UpperCAmelCase_    :      Dict            =						seq_length
							UpperCAmelCase_    :      str            =						is_training
							UpperCAmelCase_    :      str            =						use_input_mask
							UpperCAmelCase_    :      Tuple            =						use_token_type_ids
							UpperCAmelCase_    :      Optional[Any]            =						use_labels
							UpperCAmelCase_    :      List[Any]            =						vocab_size
							UpperCAmelCase_    :      Any            =						hidden_size
							UpperCAmelCase_    :      Optional[int]            =						embedding_size
							UpperCAmelCase_    :      Dict            =						num_hidden_layers
							UpperCAmelCase_    :      Optional[Any]            =						num_attention_heads
							UpperCAmelCase_    :      Any            =						intermediate_size
							UpperCAmelCase_    :      List[Any]            =						hidden_act
							UpperCAmelCase_    :      Dict            =						hidden_dropout_prob
							UpperCAmelCase_    :      Optional[int]            =						attention_probs_dropout_prob
							UpperCAmelCase_    :      Tuple            =						max_position_embeddings
							UpperCAmelCase_    :      int            =						type_vocab_size
							UpperCAmelCase_    :      Optional[int]            =						type_sequence_label_size
							UpperCAmelCase_    :      Tuple            =						initializer_range
							UpperCAmelCase_    :      List[Any]            =						num_labels
							UpperCAmelCase_    :      str            =						num_choices
							UpperCAmelCase_    :      Optional[Any]            =						scope



		def 					__UpperCAmelCase			(							self   )	->							Optional[int]:
							UpperCAmelCase_    :      int            =						ids_tensor([self.batch_size, self.seq_length]				,		self.vocab_size   )

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

							UpperCAmelCase_    :      Tuple            =						None
							if self.use_token_type_ids:
												UpperCAmelCase_    :      Optional[Any]            =						ids_tensor([self.batch_size, self.seq_length]				,		self.type_vocab_size   )

							UpperCAmelCase_    :      Any            =						None
							UpperCAmelCase_    :      Any            =						None
							UpperCAmelCase_    :      Any            =						None
							if self.use_labels:
												UpperCAmelCase_    :      Any            =						ids_tensor([self.batch_size]				,		self.type_sequence_label_size   )
												UpperCAmelCase_    :      int            =						ids_tensor([self.batch_size, self.seq_length]				,		self.num_labels   )
												UpperCAmelCase_    :      Any            =						ids_tensor([self.batch_size]				,		self.num_choices   )

							UpperCAmelCase_    :      List[str]            =						self.get_config()

							return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



		def 					__UpperCAmelCase			(							self   )	->							Optional[Any]:
							return MegatronBertConfig(
							    vocab_size=self.vocab_size				,		hidden_size=self.hidden_size				,		num_hidden_layers=self.num_hidden_layers				,		num_attention_heads=self.num_attention_heads				,		intermediate_size=self.intermediate_size				,		embedding_size=self.embedding_size				,		hidden_act=self.hidden_act				,		hidden_dropout_prob=self.hidden_dropout_prob				,		attention_probs_dropout_prob=self.attention_probs_dropout_prob				,		max_position_embeddings=self.max_position_embeddings				,		type_vocab_size=self.type_vocab_size				,		is_decoder=_UpperCamelCase				,		initializer_range=self.initializer_range				,		)



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							Dict:
							UpperCAmelCase_    :      Dict            =						MegatronBertModel(config=_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.eval()
							UpperCAmelCase_    :      Dict            =						model(_UpperCamelCase				,		attention_mask=_UpperCamelCase				,		token_type_ids=_UpperCamelCase   )
							UpperCAmelCase_    :      Optional[int]            =						model(_UpperCamelCase				,		token_type_ids=_UpperCamelCase   )
							UpperCAmelCase_    :      Tuple            =						model(_UpperCamelCase   )

							self.parent.assertEqual(result.last_hidden_state.shape				,		(self.batch_size, self.seq_length, self.hidden_size)   )
							self.parent.assertEqual(result.pooler_output.shape				,		(self.batch_size, self.hidden_size)   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							Optional[Any]:
							UpperCAmelCase_    :      List[Any]            =						MegatronBertForMaskedLM(config=_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.eval()
							UpperCAmelCase_    :      Tuple            =						model(_UpperCamelCase				,		attention_mask=_UpperCamelCase				,		token_type_ids=_UpperCamelCase				,		labels=_UpperCamelCase   )
							self.parent.assertEqual(result.logits.shape				,		(self.batch_size, self.seq_length, self.vocab_size)   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							Optional[int]:
							UpperCAmelCase_    :      Any            =						MegatronBertForCausalLM(config=_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.eval()
							UpperCAmelCase_    :      List[Any]            =						model(_UpperCamelCase				,		attention_mask=_UpperCamelCase				,		token_type_ids=_UpperCamelCase				,		labels=_UpperCamelCase   )
							self.parent.assertEqual(result.logits.shape				,		(self.batch_size, self.seq_length, self.vocab_size)   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							str:
							UpperCAmelCase_    :      Optional[int]            =						MegatronBertForNextSentencePrediction(config=_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.eval()
							UpperCAmelCase_    :      Optional[Any]            =						model(
							    _UpperCamelCase				,		attention_mask=_UpperCamelCase				,		token_type_ids=_UpperCamelCase				,		labels=_UpperCamelCase				,		)
							self.parent.assertEqual(result.logits.shape				,		(self.batch_size, 2)   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							Optional[Any]:
							UpperCAmelCase_    :      Optional[int]            =						MegatronBertForPreTraining(config=_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.eval()
							UpperCAmelCase_    :      List[str]            =						model(
							    _UpperCamelCase				,		attention_mask=_UpperCamelCase				,		token_type_ids=_UpperCamelCase				,		labels=_UpperCamelCase				,		next_sentence_label=_UpperCamelCase				,		)
							self.parent.assertEqual(result.prediction_logits.shape				,		(self.batch_size, self.seq_length, self.vocab_size)   )
							self.parent.assertEqual(result.seq_relationship_logits.shape				,		(self.batch_size, 2)   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							Any:
							UpperCAmelCase_    :      Optional[int]            =						MegatronBertForQuestionAnswering(config=_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.eval()
							UpperCAmelCase_    :      Dict            =						model(
							    _UpperCamelCase				,		attention_mask=_UpperCamelCase				,		token_type_ids=_UpperCamelCase				,		start_positions=_UpperCamelCase				,		end_positions=_UpperCamelCase				,		)
							self.parent.assertEqual(result.start_logits.shape				,		(self.batch_size, self.seq_length)   )
							self.parent.assertEqual(result.end_logits.shape				,		(self.batch_size, self.seq_length)   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							List[Any]:
							UpperCAmelCase_    :      Union[str, Any]            =						self.num_labels
							UpperCAmelCase_    :      List[Any]            =						MegatronBertForSequenceClassification(_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.eval()
							UpperCAmelCase_    :      Any            =						model(_UpperCamelCase				,		attention_mask=_UpperCamelCase				,		token_type_ids=_UpperCamelCase				,		labels=_UpperCamelCase   )
							self.parent.assertEqual(result.logits.shape				,		(self.batch_size, self.num_labels)   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							str:
							UpperCAmelCase_    :      Union[str, Any]            =						self.num_labels
							UpperCAmelCase_    :      Optional[Any]            =						MegatronBertForTokenClassification(config=_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.eval()
							UpperCAmelCase_    :      Optional[int]            =						model(_UpperCamelCase				,		attention_mask=_UpperCamelCase				,		token_type_ids=_UpperCamelCase				,		labels=_UpperCamelCase   )
							self.parent.assertEqual(result.logits.shape				,		(self.batch_size, self.seq_length, self.num_labels)   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							Dict:
							UpperCAmelCase_    :      Optional[int]            =						self.num_choices
							UpperCAmelCase_    :      str            =						MegatronBertForMultipleChoice(config=_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.eval()
							UpperCAmelCase_    :      Optional[Any]            =						input_ids.unsqueeze(1   ).expand(-1				,		self.num_choices				,		-1   ).contiguous()
							UpperCAmelCase_    :      List[str]            =						token_type_ids.unsqueeze(1   ).expand(-1				,		self.num_choices				,		-1   ).contiguous()
							UpperCAmelCase_    :      List[str]            =						input_mask.unsqueeze(1   ).expand(-1				,		self.num_choices				,		-1   ).contiguous()
							UpperCAmelCase_    :      Tuple            =						model(
							    _UpperCamelCase				,		attention_mask=_UpperCamelCase				,		token_type_ids=_UpperCamelCase				,		labels=_UpperCamelCase				,		)
							self.parent.assertEqual(result.logits.shape				,		(self.batch_size, self.num_choices)   )



		def 					__UpperCAmelCase			(							self   )	->							Union[str, Any]:
							UpperCAmelCase_    :      Dict            =						self.prepare_config_and_inputs()
							(
							    (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  
							)    :      Optional[int]            =						config_and_inputs
							UpperCAmelCase_    :      Optional[Any]            =						{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
							return config, inputs_dict





@require_torch
class   lowerCamelCase      (_snake_case ,  _snake_case ,  unittest.TestCase			):







		'''simple docstring'''



		_snake_case    :    Tuple       						=  (
		    (
		        MegatronBertModel,
		        MegatronBertForMaskedLM,
		        MegatronBertForCausalLM,
		        MegatronBertForMultipleChoice,
		        MegatronBertForNextSentencePrediction,
		        MegatronBertForPreTraining,
		        MegatronBertForQuestionAnswering,
		        MegatronBertForSequenceClassification,
		        MegatronBertForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		_snake_case    :    str       						=  (
		    {
		        '''feature-extraction''': MegatronBertModel,
		        '''fill-mask''': MegatronBertForMaskedLM,
		        '''question-answering''': MegatronBertForQuestionAnswering,
		        '''text-classification''': MegatronBertForSequenceClassification,
		        '''text-generation''': MegatronBertForCausalLM,
		        '''token-classification''': MegatronBertForTokenClassification,
		        '''zero-shot''': MegatronBertForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)
		_snake_case    :    Any       						=  True
		# test_resize_embeddings = False
		_snake_case    :    List[str]       						=  False



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase=False   )	->							Any:
							UpperCAmelCase_    :      Dict            =						super()._prepare_for_class(_UpperCamelCase				,		_UpperCamelCase				,		return_labels=_UpperCamelCase   )

							if return_labels:
												if model_class in get_values(_UpperCamelCase   ):
																	UpperCAmelCase_    :      int            =						torch.zeros(
																	    (self.model_tester.batch_size, self.model_tester.seq_length)				,		dtype=torch.long				,		device=_UpperCamelCase   )
																	UpperCAmelCase_    :      int            =						torch.zeros(
																	    self.model_tester.batch_size				,		dtype=torch.long				,		device=_UpperCamelCase   )
							return inputs_dict



		def 					__UpperCAmelCase			(							self   )	->							Tuple:
							UpperCAmelCase_    :      int            =						MegatronBertModelTester(self   )
							UpperCAmelCase_    :      Dict            =						ConfigTester(self				,		config_class=_UpperCamelCase				,		hidden_size=3_7   )



		def 					__UpperCAmelCase			(							self   )	->							Optional[int]:
							self.config_tester.run_common_tests()



		def 					__UpperCAmelCase			(							self   )	->							Any:
							UpperCAmelCase_    :      str            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_model(*_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							str:
							UpperCAmelCase_    :      Optional[Any]            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_masked_lm(*_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							List[str]:
							UpperCAmelCase_    :      Dict            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							Optional[Any]:
							UpperCAmelCase_    :      str            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							Optional[int]:
							UpperCAmelCase_    :      Any            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_pretraining(*_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							str:
							UpperCAmelCase_    :      Optional[int]            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_question_answering(*_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							int:
							UpperCAmelCase_    :      Any            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							Optional[int]:
							UpperCAmelCase_    :      List[Any]            =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_token_classification(*_UpperCamelCase   )





def      lowercase__      (     __snake_case    :							Optional[Any]  ):


					'''simple docstring'''


					return torch.tensor(
					    __snake_case     ,       dtype=torch.long     ,       device=__snake_case     ,       )


__UpperCAmelCase								=						1E-4






@require_torch
@require_sentencepiece
@require_tokenizers
class   lowerCamelCase      (unittest.TestCase			):







		'''simple docstring'''



		@slow
		@unittest.skip('Model is not available.'   )
		def 					__UpperCAmelCase			(							self   )	->							Dict:
							UpperCAmelCase_    :      int            =						'nvidia/megatron-bert-uncased-345m'
							if "MYDIR" in os.environ:
												UpperCAmelCase_    :      List[str]            =						os.path.join(os.environ['MYDIR']				,		_UpperCamelCase   )
							UpperCAmelCase_    :      Tuple            =						MegatronBertModel.from_pretrained(_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.half()
							UpperCAmelCase_    :      Any            =						_long_tensor([[1_0_1, 7_1_1_0, 1_0_0_5, 1_0_5_6, 2_0_2_3, 1_1_3_3_3, 1_7_4_1_3, 1_0_2_9, 1_0_2]]   )
							with torch.no_grad():
												UpperCAmelCase_    :      str            =						model(_UpperCamelCase   )[0]
							UpperCAmelCase_    :      List[Any]            =						torch.Size((1, 9, 1_0_2_4)   )
							self.assertEqual(output.shape				,		_UpperCamelCase   )

							UpperCAmelCase_    :      Optional[int]            =						[-0.60_40, -0.25_17, -0.10_25, 0.34_20, -0.67_58, -0.00_17, -0.10_89, -0.19_90, 0.57_28]
							for ii in range(3   ):
												for jj in range(3   ):
																	UpperCAmelCase_    :      Optional[int]            =						output[0, ii, jj]
																	UpperCAmelCase_    :      Any            =						expected[3 * ii + jj]
																	UpperCAmelCase_    :      List[Any]            =						'ii={} jj={} a={} b={}'.format(_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )
																	self.assertTrue(math.isclose(_UpperCamelCase				,		_UpperCamelCase				,		rel_tol=_UpperCamelCase				,		abs_tol=_UpperCamelCase   )				,		msg=_UpperCamelCase   )














from dataclasses import dataclass
from typing import Dict, Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .attention_processor import AttentionProcessor, AttnProcessor
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, DiagonalGaussianDistribution, Encoder



@dataclass
class   lowerCamelCase      (_snake_case			):







		'''simple docstring'''



		_snake_case    :    "DiagonalGaussianDistribution"





class   lowerCamelCase      (_snake_case ,  _snake_case			):







		'''simple docstring'''



		_snake_case    :    Optional[int]       						=  True



		@register_to_config
		def __init__(							self				,		_UpperCamelCase = 3				,		_UpperCamelCase = 3				,		_UpperCamelCase = ("DownEncoderBlock2D",)				,		_UpperCamelCase = ("UpDecoderBlock2D",)				,		_UpperCamelCase = (6_4,)				,		_UpperCamelCase = 1				,		_UpperCamelCase = "silu"				,		_UpperCamelCase = 4				,		_UpperCamelCase = 3_2				,		_UpperCamelCase = 3_2				,		_UpperCamelCase = 0.1_82_15				,		)	->							List[Any]:
							super().__init__()

							# pass init params to Encoder
							UpperCAmelCase_    :      List[str]            =						Encoder(
							    in_channels=_UpperCamelCase				,		out_channels=_UpperCamelCase				,		down_block_types=_UpperCamelCase				,		block_out_channels=_UpperCamelCase				,		layers_per_block=_UpperCamelCase				,		act_fn=_UpperCamelCase				,		norm_num_groups=_UpperCamelCase				,		double_z=_UpperCamelCase				,		)

							# pass init params to Decoder
							UpperCAmelCase_    :      Dict            =						Decoder(
							    in_channels=_UpperCamelCase				,		out_channels=_UpperCamelCase				,		up_block_types=_UpperCamelCase				,		block_out_channels=_UpperCamelCase				,		layers_per_block=_UpperCamelCase				,		norm_num_groups=_UpperCamelCase				,		act_fn=_UpperCamelCase				,		)

							UpperCAmelCase_    :      Any            =						nn.Convad(2 * latent_channels				,		2 * latent_channels				,		1   )
							UpperCAmelCase_    :      List[Any]            =						nn.Convad(_UpperCamelCase				,		_UpperCamelCase				,		1   )

							UpperCAmelCase_    :      Any            =						False
							UpperCAmelCase_    :      int            =						False

							# only relevant if vae tiling is enabled
							UpperCAmelCase_    :      Optional[int]            =						self.config.sample_size
							UpperCAmelCase_    :      int            =						(
							    self.config.sample_size[0]
							    if isinstance(self.config.sample_size				,		(list, tuple)   )
							    else self.config.sample_size
							)
							UpperCAmelCase_    :      Union[str, Any]            =						int(sample_size / (2 ** (len(self.config.block_out_channels   ) - 1))   )
							UpperCAmelCase_    :      Optional[Any]            =						0.25



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase=False   )	->							List[str]:
							if isinstance(_UpperCamelCase				,		(Encoder, Decoder)   ):
												UpperCAmelCase_    :      Union[str, Any]            =						value



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase = True   )	->							int:
							UpperCAmelCase_    :      Tuple            =						use_tiling



		def 					__UpperCAmelCase			(							self   )	->							Dict:
							self.enable_tiling(_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							Optional[Any]:
							UpperCAmelCase_    :      str            =						True



		def 					__UpperCAmelCase			(							self   )	->							List[Any]:
							UpperCAmelCase_    :      Optional[int]            =						False



		@property
		# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
		def 					__UpperCAmelCase			(							self   )	->							Dict[str, AttentionProcessor]:
							UpperCAmelCase_    :      Optional[int]            =						{}

							def fn_recursive_add_processors(_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   ):
												if hasattr(_UpperCamelCase				,		'set_processor'   ):
																	UpperCAmelCase_    :      Optional[int]            =						module.processor

												for sub_name, child in module.named_children():
																	fn_recursive_add_processors(f"{name}.{sub_name}"				,		_UpperCamelCase				,		_UpperCamelCase   )

												return processors

							for name, module in self.named_children():
												fn_recursive_add_processors(_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )

							return processors



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase   )	->							List[Any]:
							UpperCAmelCase_    :      Union[str, Any]            =						len(self.attn_processors.keys()   )

							if isinstance(_UpperCamelCase				,		_UpperCamelCase   ) and len(_UpperCamelCase   ) != count:
												raise ValueError(
												    f"A dict of processors was passed, but the number of processors {len(_UpperCamelCase   )} does not match the"
												    f" number of attention layers: {count}. Please make sure to pass {count} processor classes."   )

							def fn_recursive_attn_processor(_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   ):
												if hasattr(_UpperCamelCase				,		'set_processor'   ):
																	if not isinstance(_UpperCamelCase				,		_UpperCamelCase   ):
																						module.set_processor(_UpperCamelCase   )
																	else:
																						module.set_processor(processor.pop(f"{name}.processor"   )   )

												for sub_name, child in module.named_children():
																	fn_recursive_attn_processor(f"{name}.{sub_name}"				,		_UpperCamelCase				,		_UpperCamelCase   )

							for name, module in self.named_children():
												fn_recursive_attn_processor(_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							Union[str, Any]:
							self.set_attn_processor(AttnProcessor()   )



		@apply_forward_hook
		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase = True   )	->							AutoencoderKLOutput:
							if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
												return self.tiled_encode(_UpperCamelCase				,		return_dict=_UpperCamelCase   )

							if self.use_slicing and x.shape[0] > 1:
												UpperCAmelCase_    :      Union[str, Any]            =						[self.encoder(_UpperCamelCase   ) for x_slice in x.split(1   )]
												UpperCAmelCase_    :      Tuple            =						torch.cat(_UpperCamelCase   )
							else:
												UpperCAmelCase_    :      List[Any]            =						self.encoder(_UpperCamelCase   )

							UpperCAmelCase_    :      Optional[Any]            =						self.quant_conv(_UpperCamelCase   )
							UpperCAmelCase_    :      Tuple            =						DiagonalGaussianDistribution(_UpperCamelCase   )

							if not return_dict:
												return (posterior,)

							return AutoencoderKLOutput(latent_dist=_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase = True   )	->							Union[DecoderOutput, torch.FloatTensor]:
							if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
												return self.tiled_decode(_UpperCamelCase				,		return_dict=_UpperCamelCase   )

							UpperCAmelCase_    :      str            =						self.post_quant_conv(_UpperCamelCase   )
							UpperCAmelCase_    :      List[str]            =						self.decoder(_UpperCamelCase   )

							if not return_dict:
												return (dec,)

							return DecoderOutput(sample=_UpperCamelCase   )



		@apply_forward_hook
		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase = True   )	->							Union[DecoderOutput, torch.FloatTensor]:
							if self.use_slicing and z.shape[0] > 1:
												UpperCAmelCase_    :      List[str]            =						[self._decode(_UpperCamelCase   ).sample for z_slice in z.split(1   )]
												UpperCAmelCase_    :      Dict            =						torch.cat(_UpperCamelCase   )
							else:
												UpperCAmelCase_    :      Any            =						self._decode(_UpperCamelCase   ).sample

							if not return_dict:
												return (decoded,)

							return DecoderOutput(sample=_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							Any:
							UpperCAmelCase_    :      Tuple            =						min(a.shape[2]				,		b.shape[2]				,		_UpperCamelCase   )
							for y in range(_UpperCamelCase   ):
												UpperCAmelCase_    :      str            =						a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, y, :] * (y / blend_extent)
							return b



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )	->							Dict:
							UpperCAmelCase_    :      Tuple            =						min(a.shape[3]				,		b.shape[3]				,		_UpperCamelCase   )
							for x in range(_UpperCamelCase   ):
												UpperCAmelCase_    :      int            =						a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
							return b



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase = True   )	->							AutoencoderKLOutput:
							UpperCAmelCase_    :      Any            =						int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)   )
							UpperCAmelCase_    :      Tuple            =						int(self.tile_latent_min_size * self.tile_overlap_factor   )
							UpperCAmelCase_    :      Optional[int]            =						self.tile_latent_min_size - blend_extent

							# Split the image into 512x512 tiles and encode them separately.
							UpperCAmelCase_    :      List[str]            =						[]
							for i in range(0				,		x.shape[2]				,		_UpperCamelCase   ):
												UpperCAmelCase_    :      Any            =						[]
												for j in range(0				,		x.shape[3]				,		_UpperCamelCase   ):
																	UpperCAmelCase_    :      Any            =						x[:, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
																	UpperCAmelCase_    :      Dict            =						self.encoder(_UpperCamelCase   )
																	UpperCAmelCase_    :      List[str]            =						self.quant_conv(_UpperCamelCase   )
																	row.append(_UpperCamelCase   )
												rows.append(_UpperCamelCase   )
							UpperCAmelCase_    :      str            =						[]
							for i, row in enumerate(_UpperCamelCase   ):
												UpperCAmelCase_    :      List[Any]            =						[]
												for j, tile in enumerate(_UpperCamelCase   ):
																	# blend the above tile and the left tile
																	# to the current tile and add the current tile to the result row
																	if i > 0:
																						UpperCAmelCase_    :      Dict            =						self.blend_v(rows[i - 1][j]				,		_UpperCamelCase				,		_UpperCamelCase   )
																	if j > 0:
																						UpperCAmelCase_    :      List[str]            =						self.blend_h(row[j - 1]				,		_UpperCamelCase				,		_UpperCamelCase   )
																	result_row.append(tile[:, :, :row_limit, :row_limit]   )
												result_rows.append(torch.cat(_UpperCamelCase				,		dim=3   )   )

							UpperCAmelCase_    :      Union[str, Any]            =						torch.cat(_UpperCamelCase				,		dim=2   )
							UpperCAmelCase_    :      List[Any]            =						DiagonalGaussianDistribution(_UpperCamelCase   )

							if not return_dict:
												return (posterior,)

							return AutoencoderKLOutput(latent_dist=_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase = True   )	->							Union[DecoderOutput, torch.FloatTensor]:
							UpperCAmelCase_    :      str            =						int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)   )
							UpperCAmelCase_    :      Dict            =						int(self.tile_sample_min_size * self.tile_overlap_factor   )
							UpperCAmelCase_    :      Dict            =						self.tile_sample_min_size - blend_extent

							# Split z into overlapping 64x64 tiles and decode them separately.
							# The tiles have an overlap to avoid seams between tiles.
							UpperCAmelCase_    :      Union[str, Any]            =						[]
							for i in range(0				,		z.shape[2]				,		_UpperCamelCase   ):
												UpperCAmelCase_    :      List[str]            =						[]
												for j in range(0				,		z.shape[3]				,		_UpperCamelCase   ):
																	UpperCAmelCase_    :      List[str]            =						z[:, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
																	UpperCAmelCase_    :      Optional[Any]            =						self.post_quant_conv(_UpperCamelCase   )
																	UpperCAmelCase_    :      Tuple            =						self.decoder(_UpperCamelCase   )
																	row.append(_UpperCamelCase   )
												rows.append(_UpperCamelCase   )
							UpperCAmelCase_    :      Optional[Any]            =						[]
							for i, row in enumerate(_UpperCamelCase   ):
												UpperCAmelCase_    :      List[Any]            =						[]
												for j, tile in enumerate(_UpperCamelCase   ):
																	# blend the above tile and the left tile
																	# to the current tile and add the current tile to the result row
																	if i > 0:
																						UpperCAmelCase_    :      Union[str, Any]            =						self.blend_v(rows[i - 1][j]				,		_UpperCamelCase				,		_UpperCamelCase   )
																	if j > 0:
																						UpperCAmelCase_    :      Optional[Any]            =						self.blend_h(row[j - 1]				,		_UpperCamelCase				,		_UpperCamelCase   )
																	result_row.append(tile[:, :, :row_limit, :row_limit]   )
												result_rows.append(torch.cat(_UpperCamelCase				,		dim=3   )   )

							UpperCAmelCase_    :      Dict            =						torch.cat(_UpperCamelCase				,		dim=2   )
							if not return_dict:
												return (dec,)

							return DecoderOutput(sample=_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase = False				,		_UpperCamelCase = True				,		_UpperCamelCase = None				,		)	->							Union[DecoderOutput, torch.FloatTensor]:
							UpperCAmelCase_    :      Optional[Any]            =						sample
							UpperCAmelCase_    :      Union[str, Any]            =						self.encode(_UpperCamelCase   ).latent_dist
							if sample_posterior:
												UpperCAmelCase_    :      str            =						posterior.sample(generator=_UpperCamelCase   )
							else:
												UpperCAmelCase_    :      int            =						posterior.mode()
							UpperCAmelCase_    :      Dict            =						self.decode(_UpperCamelCase   ).sample

							if not return_dict:
												return (dec,)

							return DecoderOutput(sample=_UpperCamelCase   )






