






import inspect
import unittest

from transformers import DecisionTransformerConfig, 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 DecisionTransformerModel
		from transformers.models.decision_transformer.modeling_decision_transformer import (
		    DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
		)



class   lowerCamelCase      :







		'''simple docstring'''



		def __init__(							self				,		_UpperCamelCase				,		_UpperCamelCase=1_3				,		_UpperCamelCase=7				,		_UpperCamelCase=6				,		_UpperCamelCase=1_7				,		_UpperCamelCase=2_3				,		_UpperCamelCase=1_1				,		_UpperCamelCase=True				,		)	->							int:
							UpperCAmelCase_    :      Dict            =						parent
							UpperCAmelCase_    :      str            =						batch_size
							UpperCAmelCase_    :      List[Any]            =						seq_length
							UpperCAmelCase_    :      Any            =						act_dim
							UpperCAmelCase_    :      Dict            =						state_dim
							UpperCAmelCase_    :      List[str]            =						hidden_size
							UpperCAmelCase_    :      List[Any]            =						max_length
							UpperCAmelCase_    :      Optional[Any]            =						is_training



		def 					__UpperCAmelCase			(							self   )	->							Tuple:
							UpperCAmelCase_    :      Union[str, Any]            =						floats_tensor((self.batch_size, self.seq_length, self.state_dim)   )
							UpperCAmelCase_    :      Optional[Any]            =						floats_tensor((self.batch_size, self.seq_length, self.act_dim)   )
							UpperCAmelCase_    :      Tuple            =						floats_tensor((self.batch_size, self.seq_length, 1)   )
							UpperCAmelCase_    :      Optional[Any]            =						floats_tensor((self.batch_size, self.seq_length, 1)   )
							UpperCAmelCase_    :      str            =						ids_tensor((self.batch_size, self.seq_length)				,		vocab_size=1_0_0_0   )
							UpperCAmelCase_    :      Optional[Any]            =						random_attention_mask((self.batch_size, self.seq_length)   )

							UpperCAmelCase_    :      Dict            =						self.get_config()

							return (
							    config,
							    states,
							    actions,
							    rewards,
							    returns_to_go,
							    timesteps,
							    attention_mask,
							)



		def 					__UpperCAmelCase			(							self   )	->							Optional[Any]:
							return DecisionTransformerConfig(
							    batch_size=self.batch_size				,		seq_length=self.seq_length				,		act_dim=self.act_dim				,		state_dim=self.state_dim				,		hidden_size=self.hidden_size				,		max_length=self.max_length				,		)



		def 					__UpperCAmelCase			(							self				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		)	->							List[str]:
							UpperCAmelCase_    :      List[str]            =						DecisionTransformerModel(config=_UpperCamelCase   )
							model.to(_UpperCamelCase   )
							model.eval()
							UpperCAmelCase_    :      Dict            =						model(_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase				,		_UpperCamelCase   )

							self.parent.assertEqual(result.state_preds.shape				,		states.shape   )
							self.parent.assertEqual(result.action_preds.shape				,		actions.shape   )
							self.parent.assertEqual(result.return_preds.shape				,		returns_to_go.shape   )
							self.parent.assertEqual(
							    result.last_hidden_state.shape				,		(self.batch_size, self.seq_length * 3, self.hidden_size)   )  # seq length *3 as there are 3 modelities: states, returns and actions



		def 					__UpperCAmelCase			(							self   )	->							Optional[Any]:
							UpperCAmelCase_    :      List[Any]            =						self.prepare_config_and_inputs()
							(
							    (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  (
							    UpperCAmelCase_
							)       ,  
							)    :      List[str]            =						config_and_inputs
							UpperCAmelCase_    :      Any            =						{
							    'states': states,
							    'actions': actions,
							    'rewards': rewards,
							    'returns_to_go': returns_to_go,
							    'timesteps': timesteps,
							    'attention_mask': attention_mask,
							}
							return config, inputs_dict





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







		'''simple docstring'''



		_snake_case    :    int       						=  (DecisionTransformerModel,) if is_torch_available() else ()
		_snake_case    :    Tuple       						=  ()
		_snake_case    :    Any       						=  {'''feature-extraction''': DecisionTransformerModel} if is_torch_available() else {}

		# Ignoring of a failing test from GenerationTesterMixin, as the model does not use inputs_ids
		_snake_case    :    List[Any]       						=  False

		# Ignoring of a failing tests from ModelTesterMixin, as the model does not implement these features
		_snake_case    :    Union[str, Any]       						=  False
		_snake_case    :    Optional[int]       						=  False
		_snake_case    :    List[str]       						=  False
		_snake_case    :    Optional[int]       						=  False
		_snake_case    :    List[str]       						=  False
		_snake_case    :    int       						=  False
		_snake_case    :    int       						=  False
		_snake_case    :    List[str]       						=  False
		_snake_case    :    List[Any]       						=  False



		def 					__UpperCAmelCase			(							self   )	->							Tuple:
							UpperCAmelCase_    :      str            =						DecisionTransformerModelTester(self   )
							UpperCAmelCase_    :      str            =						ConfigTester(self				,		config_class=_UpperCamelCase				,		hidden_size=3_7   )



		def 					__UpperCAmelCase			(							self   )	->							List[str]:
							self.config_tester.run_common_tests()



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



		@slow
		def 					__UpperCAmelCase			(							self   )	->							List[Any]:
							for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												UpperCAmelCase_    :      int            =						DecisionTransformerModel.from_pretrained(_UpperCamelCase   )
												self.assertIsNotNone(_UpperCamelCase   )



		def 					__UpperCAmelCase			(							self   )	->							Optional[Any]:
							UpperCAmelCase_       ,  UpperCAmelCase_    :      Optional[Any]            =						self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
												UpperCAmelCase_    :      List[Any]            =						model_class(_UpperCamelCase   )
												UpperCAmelCase_    :      Optional[int]            =						inspect.signature(model.forward   )
												# signature.parameters is an OrderedDict => so arg_names order is deterministic
												UpperCAmelCase_    :      Union[str, Any]            =						[*signature.parameters.keys()]

												UpperCAmelCase_    :      List[str]            =						[
												    'states',
												    'actions',
												    'rewards',
												    'returns_to_go',
												    'timesteps',
												    'attention_mask',
												]

												self.assertListEqual(arg_names[: len(_UpperCamelCase   )]				,		_UpperCamelCase   )





@require_torch
class   lowerCamelCase      (unittest.TestCase			):







		'''simple docstring'''



		@slow
		def 					__UpperCAmelCase			(							self   )	->							int:
							UpperCAmelCase_    :      Optional[Any]            =						2  # number of steps of autoregressive prediction we will perform
							UpperCAmelCase_    :      Optional[Any]            =						1_0  # defined by the RL environment, may be normalized
							UpperCAmelCase_    :      Tuple            =						DecisionTransformerModel.from_pretrained('edbeeching/decision-transformer-gym-hopper-expert'   )
							UpperCAmelCase_    :      int            =						model.to(_UpperCamelCase   )
							UpperCAmelCase_    :      int            =						model.config
							torch.manual_seed(0   )
							UpperCAmelCase_    :      str            =						torch.randn(1				,		1				,		config.state_dim   ).to(device=_UpperCamelCase				,		dtype=torch.floataa   )  # env.reset()

							UpperCAmelCase_    :      Tuple            =						torch.tensor(
							    [[0.24_27_93, -0.28_69_30_74, 0.8_74_26_13], [0.67_81_52_74, -0.08_10_10_85, -0.12_95_21_47]]				,		device=_UpperCamelCase   )

							UpperCAmelCase_    :      Tuple            =						torch.tensor(_UpperCamelCase				,		device=_UpperCamelCase				,		dtype=torch.floataa   ).reshape(1				,		1				,		1   )
							UpperCAmelCase_    :      Optional[int]            =						state
							UpperCAmelCase_    :      str            =						torch.zeros(1				,		0				,		config.act_dim				,		device=_UpperCamelCase				,		dtype=torch.floataa   )
							UpperCAmelCase_    :      List[Any]            =						torch.zeros(1				,		0				,		device=_UpperCamelCase				,		dtype=torch.floataa   )
							UpperCAmelCase_    :      Optional[int]            =						torch.tensor(0				,		device=_UpperCamelCase				,		dtype=torch.long   ).reshape(1				,		1   )

							for step in range(_UpperCamelCase   ):
												UpperCAmelCase_    :      Union[str, Any]            =						torch.cat([actions, torch.zeros(1				,		1				,		config.act_dim				,		device=_UpperCamelCase   )]				,		dim=1   )
												UpperCAmelCase_    :      Tuple            =						torch.cat([rewards, torch.zeros(1				,		1				,		device=_UpperCamelCase   )]				,		dim=1   )

												UpperCAmelCase_    :      List[str]            =						torch.ones(1				,		states.shape[1]   ).to(dtype=torch.long				,		device=states.device   )

												with torch.no_grad():
																	UpperCAmelCase_       ,  UpperCAmelCase_       ,  UpperCAmelCase_    :      Optional[int]            =						model(
																	    states=_UpperCamelCase				,		actions=_UpperCamelCase				,		rewards=_UpperCamelCase				,		returns_to_go=_UpperCamelCase				,		timesteps=_UpperCamelCase				,		attention_mask=_UpperCamelCase				,		return_dict=_UpperCamelCase				,		)

												self.assertEqual(action_pred.shape				,		actions.shape   )
												self.assertTrue(torch.allclose(action_pred[0, -1]				,		expected_outputs[step]				,		atol=1E-4   )   )
												UpperCAmelCase_       ,  UpperCAmelCase_       ,  UpperCAmelCase_       ,  UpperCAmelCase_    :      Union[str, Any]            =						(  # env.step(action)
												    torch.randn(1				,		1				,		config.state_dim   ).to(device=_UpperCamelCase				,		dtype=torch.floataa   ),
												    1.0,
												    False,
												    {},
												)

												UpperCAmelCase_    :      List[str]            =						action_pred[0, -1]
												UpperCAmelCase_    :      Tuple            =						torch.cat([states, state]				,		dim=1   )
												UpperCAmelCase_    :      Tuple            =						returns_to_go[0, -1] - reward
												UpperCAmelCase_    :      Optional[Any]            =						torch.cat([returns_to_go, pred_return.reshape(1				,		1				,		1   )]				,		dim=1   )
												UpperCAmelCase_    :      Optional[Any]            =						torch.cat(
												    [timesteps, torch.ones((1, 1)				,		device=_UpperCamelCase				,		dtype=torch.long   ) * (step + 1)]				,		dim=1   )














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   )






