





"""simple docstring"""





import math
import unittest

from transformers import BioGptConfig, 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, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import (
           BioGptForCausalLM,
           BioGptForSequenceClassification,
           BioGptForTokenClassification,
           BioGptModel,
           BioGptTokenizer,
       )
       from transformers.models.biogpt.modeling_biogpt import BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST

class 		lowercase					:





      def __init__(       self      ,						lowercase      ,						lowercase=13      ,						lowercase=7      ,						lowercase=True      ,						lowercase=True      ,						lowercase=False      ,						lowercase=True      ,						lowercase=99      ,						lowercase=32      ,						lowercase=5      ,						lowercase=4      ,						lowercase=37      ,						lowercase="gelu"      ,						lowercase=0.1      ,						lowercase=0.1      ,						lowercase=512      ,						lowercase=16      ,						lowercase=2      ,						lowercase=0.02      ,						lowercase=3      ,						lowercase=4      ,						lowercase=None      ,						)		->     Tuple:
            lowerCAmelCase  						=       parent
            lowerCAmelCase  						=       batch_size
            lowerCAmelCase  						=       seq_length
            lowerCAmelCase  						=       is_training
            lowerCAmelCase  						=       use_input_mask
            lowerCAmelCase  						=       use_token_type_ids
            lowerCAmelCase  						=       use_labels
            lowerCAmelCase  						=       vocab_size
            lowerCAmelCase  						=       hidden_size
            lowerCAmelCase  						=       num_hidden_layers
            lowerCAmelCase  						=       num_attention_heads
            lowerCAmelCase  						=       intermediate_size
            lowerCAmelCase  						=       hidden_act
            lowerCAmelCase  						=       hidden_dropout_prob
            lowerCAmelCase  						=       attention_probs_dropout_prob
            lowerCAmelCase  						=       max_position_embeddings
            lowerCAmelCase  						=       type_vocab_size
            lowerCAmelCase  						=       type_sequence_label_size
            lowerCAmelCase  						=       initializer_range
            lowerCAmelCase  						=       num_labels
            lowerCAmelCase  						=       num_choices
            lowerCAmelCase  						=       scope





      def     _snake_case			(       self )		->     int:
            lowerCAmelCase  						=       ids_tensor([self.batch_size, self.seq_length]      ,						self.vocab_size )

            lowerCAmelCase  						=       None
            if self.use_input_mask:
                  lowerCAmelCase  						=       random_attention_mask([self.batch_size, self.seq_length] )

            lowerCAmelCase  						=       None
            if self.use_token_type_ids:
                  lowerCAmelCase  						=       ids_tensor([self.batch_size, self.seq_length]      ,						self.type_vocab_size )

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

            lowerCAmelCase  						=       self.get_config()

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





      def     _snake_case			(       self )		->     Dict:
            return BioGptConfig(
                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      ,						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=lowercase      ,						initializer_range=self.initializer_range      ,						)





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     str:
            lowerCAmelCase  						=       BioGptModel(config=lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase )
            lowerCAmelCase  						=       model(lowercase )
            self.parent.assertEqual(result.last_hidden_state.shape      ,						(self.batch_size, self.seq_length, self.hidden_size) )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						)		->     Any:
            lowerCAmelCase  						=       BioGptForCausalLM(config=lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase      ,						token_type_ids=lowercase      ,						labels=lowercase )
            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 )		->     Optional[Any]:
            lowerCAmelCase  						=       BioGptModel(config=lowercase )
            model.to(lowercase )
            model.eval()

            # create attention mask
            lowerCAmelCase  						=       torch.ones(input_ids.shape      ,						dtype=torch.long      ,						device=lowercase )
            lowerCAmelCase  						=       self.seq_length // 2
            lowerCAmelCase  						=       0

            # first forward pass
            lowerCAmelCase						,		lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase ).to_tuple()

            # create hypothetical next token and extent to next_input_ids
            lowerCAmelCase  						=       ids_tensor((self.batch_size, 1)      ,						config.vocab_size )

            # change a random masked slice from input_ids
            lowerCAmelCase  						=       ids_tensor((1,)      ,						lowercase ).item() + 1
            lowerCAmelCase  						=       ids_tensor((self.batch_size, 1)      ,						config.vocab_size ).squeeze(-1 )
            lowerCAmelCase  						=       random_other_next_tokens

            # append to next input_ids and attn_mask
            lowerCAmelCase  						=       torch.cat([input_ids, next_tokens]      ,						dim=-1 )
            lowerCAmelCase  						=       torch.cat(
                [attn_mask, torch.ones((attn_mask.shape[0], 1)      ,						dtype=torch.long      ,						device=lowercase )]      ,						dim=1      ,						)

            # get two different outputs
            lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase )["""last_hidden_state"""]
            lowerCAmelCase  						=       model(lowercase      ,						past_key_values=lowercase      ,						attention_mask=lowercase )["""last_hidden_state"""]

            # select random slice
            lowerCAmelCase  						=       ids_tensor((1,)      ,						output_from_past.shape[-1] ).item()
            lowerCAmelCase  						=       output_from_no_past[:, -1, random_slice_idx].detach()
            lowerCAmelCase  						=       output_from_past[:, 0, random_slice_idx].detach()

            # test that outputs are equal for slice
            self.parent.assertTrue(torch.allclose(lowercase      ,						lowercase      ,						atol=1e-3 ) )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						*lowercase )		->     Optional[Any]:
            lowerCAmelCase  						=       BioGptModel(config=lowercase ).to(lowercase ).eval()

            lowerCAmelCase  						=       torch.ones(input_ids.shape      ,						dtype=torch.long      ,						device=lowercase )

            # first forward pass
            lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase      ,						use_cache=lowercase )

            lowerCAmelCase						,		lowerCAmelCase  						=       outputs.to_tuple()

            # create hypothetical multiple next token and extent to next_input_ids
            lowerCAmelCase  						=       ids_tensor((self.batch_size, 3)      ,						config.vocab_size )
            lowerCAmelCase  						=       ids_tensor((self.batch_size, 3)      ,						2 )

            # append to next input_ids and
            lowerCAmelCase  						=       torch.cat([input_ids, next_tokens]      ,						dim=-1 )
            lowerCAmelCase  						=       torch.cat([attention_mask, next_attn_mask]      ,						dim=-1 )

            lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase )["""last_hidden_state"""]
            lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase      ,						past_key_values=lowercase )[
                """last_hidden_state"""
            ]

            # select random slice
            lowerCAmelCase  						=       ids_tensor((1,)      ,						output_from_past.shape[-1] ).item()
            lowerCAmelCase  						=       output_from_no_past[:, -3:, random_slice_idx].detach()
            lowerCAmelCase  						=       output_from_past[:, :, random_slice_idx].detach()

            self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1] )

            # test that outputs are equal for slice
            self.parent.assertTrue(torch.allclose(lowercase      ,						lowercase      ,						atol=1e-3 ) )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						*lowercase      ,						lowercase=False )		->     Optional[Any]:
            lowerCAmelCase  						=       BioGptForCausalLM(lowercase )
            model.to(lowercase )
            if gradient_checkpointing:
                  model.gradient_checkpointing_enable()

            lowerCAmelCase  						=       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      ,						lowercase      ,						*lowercase )		->     List[str]:
            lowerCAmelCase  						=       BioGptModel(lowercase )
            lowerCAmelCase  						=       model.config.initializer_range / math.sqrt(2 * model.config.num_hidden_layers )
            for key in model.state_dict().keys():
                  if "c_proj" in key and "weight" in key:
                        self.parent.assertLessEqual(abs(torch.std(model.state_dict()[key] ) - model_std )      ,						0.001 )
                        self.parent.assertLessEqual(abs(torch.mean(model.state_dict()[key] ) - 0.0 )      ,						0.01 )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						*lowercase )		->     Optional[int]:
            lowerCAmelCase  						=       self.num_labels
            lowerCAmelCase  						=       BioGptForTokenClassification(lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase      ,						token_type_ids=lowercase )
            self.parent.assertEqual(result.logits.shape      ,						(self.batch_size, self.seq_length, self.num_labels) )





      def     _snake_case			(       self )		->     Dict:
            lowerCAmelCase  						=       self.prepare_config_and_inputs()
            (
                (
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		
            )  						=       config_and_inputs
            lowerCAmelCase  						=       {"""input_ids""": input_ids, """attention_mask""": input_mask}
            return config, inputs_dict







@require_torch
class 		lowercase					(					_UpperCAmelCase     ,  _UpperCAmelCase     ,  _UpperCAmelCase     ,  unittest.TestCase					):
      _SCREAMING_SNAKE_CASE					=      (
          (BioGptModel, BioGptForCausalLM, BioGptForSequenceClassification, BioGptForTokenClassification)
          if is_torch_available()
          else ()
      )
      _SCREAMING_SNAKE_CASE					=      (BioGptForCausalLM,) if is_torch_available() else ()
      _SCREAMING_SNAKE_CASE					=      (
          {
              'feature-extraction': BioGptModel,
              'text-classification': BioGptForSequenceClassification,
              'text-generation': BioGptForCausalLM,
              'token-classification': BioGptForTokenClassification,
              'zero-shot': BioGptForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      _SCREAMING_SNAKE_CASE					=      False





      def     _snake_case			(       self )		->     List[Any]:
            lowerCAmelCase  						=       BioGptModelTester(self )
            lowerCAmelCase  						=       ConfigTester(self      ,						config_class=lowercase      ,						hidden_size=37 )





      def     _snake_case			(       self )		->     Any:
            self.config_tester.run_common_tests()





      def     _snake_case			(       self )		->     Union[str, Any]:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*lowercase )





      def     _snake_case			(       self )		->     Dict:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            for type in ["absolute", "relative_key", "relative_key_query"]:
                  lowerCAmelCase  						=       type
                  self.model_tester.create_and_check_model(*lowercase )





      def     _snake_case			(       self )		->     Tuple:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_biogpt_model_attention_mask_past(*lowercase )





      def     _snake_case			(       self )		->     Tuple:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_forward_and_backwards(*lowercase      ,						gradient_checkpointing=lowercase )





      def     _snake_case			(       self )		->     Optional[int]:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_biogpt_model_past_large_inputs(*lowercase )





      def     _snake_case			(       self )		->     Union[str, Any]:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_biogpt_weight_initialization(*lowercase )





      def     _snake_case			(       self )		->     Union[str, Any]:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_biogpt_for_token_classification(*lowercase )





      @slow
      def     _snake_case			(       self )		->     List[str]:
            lowerCAmelCase  						=       BioGptForCausalLM.from_pretrained("""microsoft/biogpt""" )
            model.to(lowercase )
            lowerCAmelCase  						=       BioGptTokenizer.from_pretrained("""microsoft/biogpt""" )

            lowerCAmelCase  						=       """left"""

            # Define PAD Token = EOS Token = 50256
            lowerCAmelCase  						=       tokenizer.eos_token
            lowerCAmelCase  						=       model.config.eos_token_id

            # use different length sentences to test batching
            lowerCAmelCase  						=       [
                """Hello, my dog is a little""",
                """Today, I""",
            ]

            lowerCAmelCase  						=       tokenizer(lowercase      ,						return_tensors="""pt"""      ,						padding=lowercase )
            lowerCAmelCase  						=       inputs["""input_ids"""].to(lowercase )

            lowerCAmelCase  						=       model.generate(
                input_ids=lowercase      ,						attention_mask=inputs["""attention_mask"""].to(lowercase )      ,						)

            lowerCAmelCase  						=       tokenizer(sentences[0]      ,						return_tensors="""pt""" ).input_ids.to(lowercase )
            lowerCAmelCase  						=       model.generate(input_ids=lowercase )

            lowerCAmelCase  						=       inputs_non_padded.shape[-1] - inputs["""attention_mask"""][-1].long().sum().cpu().item()
            lowerCAmelCase  						=       tokenizer(sentences[1]      ,						return_tensors="""pt""" ).input_ids.to(lowercase )
            lowerCAmelCase  						=       model.generate(input_ids=lowercase      ,						max_length=model.config.max_length - num_paddings )

            lowerCAmelCase  						=       tokenizer.batch_decode(lowercase      ,						skip_special_tokens=lowercase )
            lowerCAmelCase  						=       tokenizer.decode(output_non_padded[0]      ,						skip_special_tokens=lowercase )
            lowerCAmelCase  						=       tokenizer.decode(output_padded[0]      ,						skip_special_tokens=lowercase )

            lowerCAmelCase  						=       [
                """Hello, my dog is a little bit bigger than a little bit.""",
                """Today, I have a good idea of how to use the information""",
            ]
            self.assertListEqual(lowercase      ,						lowercase )
            self.assertListEqual(lowercase      ,						[non_padded_sentence, padded_sentence] )





      @slow
      def     _snake_case			(       self )		->     Optional[int]:
            for model_name in BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  lowerCAmelCase  						=       BioGptModel.from_pretrained(lowercase )
                  self.assertIsNotNone(lowercase )





      def     _snake_case			(       self )		->     Tuple:
            lowerCAmelCase						,		lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs_for_common()
            lowerCAmelCase  						=       3
            lowerCAmelCase  						=       input_dict["""input_ids"""]
            lowerCAmelCase  						=       input_ids.ne(1 ).to(lowercase )
            lowerCAmelCase  						=       ids_tensor([self.model_tester.batch_size]      ,						self.model_tester.type_sequence_label_size )
            lowerCAmelCase  						=       BioGptForSequenceClassification(lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase      ,						labels=lowercase )
            self.assertEqual(result.logits.shape      ,						(self.model_tester.batch_size, self.model_tester.num_labels) )





      def     _snake_case			(       self )		->     Optional[Any]:
            lowerCAmelCase						,		lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs_for_common()
            lowerCAmelCase  						=       3
            lowerCAmelCase  						=       """multi_label_classification"""
            lowerCAmelCase  						=       input_dict["""input_ids"""]
            lowerCAmelCase  						=       input_ids.ne(1 ).to(lowercase )
            lowerCAmelCase  						=       ids_tensor(
                [self.model_tester.batch_size, config.num_labels]      ,						self.model_tester.type_sequence_label_size ).to(torch.float )
            lowerCAmelCase  						=       BioGptForSequenceClassification(lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase      ,						labels=lowercase )
            self.assertEqual(result.logits.shape      ,						(self.model_tester.batch_size, self.model_tester.num_labels) )







@require_torch
class 		lowercase					(					unittest.TestCase					):





      @slow
      def     _snake_case			(       self )		->     Optional[Any]:
            lowerCAmelCase  						=       BioGptForCausalLM.from_pretrained("""microsoft/biogpt""" )
            lowerCAmelCase  						=       torch.tensor([[2, 4_805, 9, 656, 21]] )
            lowerCAmelCase  						=       model(lowercase )[0]

            lowerCAmelCase  						=       42_384

            lowerCAmelCase  						=       torch.Size((1, 5, vocab_size) )
            self.assertEqual(output.shape      ,						lowercase )

            lowerCAmelCase  						=       torch.tensor(
                [[[-9.5_236, -9.8_918, 10.4_557], [-11.0_469, -9.6_423, 8.1_022], [-8.8_664, -7.8_826, 5.5_325]]] )

            self.assertTrue(torch.allclose(output[:, :3, :3]      ,						lowercase      ,						atol=1e-4 ) )





      @slow
      def     _snake_case			(       self )		->     List[Any]:
            lowerCAmelCase  						=       BioGptTokenizer.from_pretrained("""microsoft/biogpt""" )
            lowerCAmelCase  						=       BioGptForCausalLM.from_pretrained("""microsoft/biogpt""" )
            model.to(lowercase )

            torch.manual_seed(0 )
            lowerCAmelCase  						=       tokenizer("""COVID-19 is"""      ,						return_tensors="""pt""" ).to(lowercase )
            lowerCAmelCase  						=       model.generate(
                **lowercase      ,						min_length=100      ,						max_length=1_024      ,						num_beams=5      ,						early_stopping=lowercase      ,						)
            lowerCAmelCase  						=       tokenizer.decode(output_ids[0]      ,						skip_special_tokens=lowercase )

            lowerCAmelCase  						=       (
                """COVID-19 is a global pandemic caused by severe acute respiratory syndrome coronavirus 2 (SARS-CoV-2), the"""
                """ causative agent of coronavirus disease 2019 (COVID-19), which has spread to more than 200 countries and"""
                """ territories, including the United States (US), Canada, Australia, New Zealand, the United Kingdom (UK),"""
                """ and the United States of America (USA), as of March 11, 2020, with more than 800,000 confirmed cases and"""
                """ more than 800,000 deaths."""
            )
            self.assertEqual(lowercase      ,						lowercase )













"""simple docstring"""





from dataclasses import dataclass
from typing import 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 .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, Encoder, VectorQuantizer

@dataclass
class 		lowercase					(					_UpperCAmelCase					):
      _SCREAMING_SNAKE_CASE					=      42







class 		lowercase					(					_UpperCAmelCase     ,  _UpperCAmelCase					):





      @register_to_config
      def __init__(       self      ,						lowercase = 3      ,						lowercase = 3      ,						lowercase = ("DownEncoderBlock2D",)      ,						lowercase = ("UpDecoderBlock2D",)      ,						lowercase = (64,)      ,						lowercase = 1      ,						lowercase = "silu"      ,						lowercase = 3      ,						lowercase = 32      ,						lowercase = 256      ,						lowercase = 32      ,						lowercase = None      ,						lowercase = 0.18_215      ,						lowercase = "group"      ,						)		->     Union[str, Any]:
            super().__init__()

            # pass init params to Encoder
            lowerCAmelCase  						=       Encoder(
                in_channels=lowercase      ,						out_channels=lowercase      ,						down_block_types=lowercase      ,						block_out_channels=lowercase      ,						layers_per_block=lowercase      ,						act_fn=lowercase      ,						norm_num_groups=lowercase      ,						double_z=lowercase      ,						)

            lowerCAmelCase  						=       vq_embed_dim if vq_embed_dim is not None else latent_channels

            lowerCAmelCase  						=       nn.Convad(lowercase      ,						lowercase      ,						1 )
            lowerCAmelCase  						=       VectorQuantizer(lowercase      ,						lowercase      ,						beta=0.25      ,						remap=lowercase      ,						sane_index_shape=lowercase )
            lowerCAmelCase  						=       nn.Convad(lowercase      ,						lowercase      ,						1 )

            # pass init params to Decoder
            lowerCAmelCase  						=       Decoder(
                in_channels=lowercase      ,						out_channels=lowercase      ,						up_block_types=lowercase      ,						block_out_channels=lowercase      ,						layers_per_block=lowercase      ,						act_fn=lowercase      ,						norm_num_groups=lowercase      ,						norm_type=lowercase      ,						)





      @apply_forward_hook
      def     _snake_case			(       self      ,						lowercase      ,						lowercase = True )		->     VQEncoderOutput:
            lowerCAmelCase  						=       self.encoder(lowercase )
            lowerCAmelCase  						=       self.quant_conv(lowercase )

            if not return_dict:
                  return (h,)

            return VQEncoderOutput(latents=lowercase )





      @apply_forward_hook
      def     _snake_case			(       self      ,						lowercase      ,						lowercase = False      ,						lowercase = True )		->     Union[DecoderOutput, torch.FloatTensor]:
            # also go through quantization layer
            if not force_not_quantize:
                  lowerCAmelCase						,		lowerCAmelCase						,		lowerCAmelCase  						=       self.quantize(lowercase )
            else:
                  lowerCAmelCase  						=       h
            lowerCAmelCase  						=       self.post_quant_conv(lowercase )
            lowerCAmelCase  						=       self.decoder(lowercase      ,						quant if self.config.norm_type == """spatial""" else None )

            if not return_dict:
                  return (dec,)

            return DecoderOutput(sample=lowercase )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase = True )		->     Union[DecoderOutput, torch.FloatTensor]:
            lowerCAmelCase  						=       sample
            lowerCAmelCase  						=       self.encode(lowercase ).latents
            lowerCAmelCase  						=       self.decode(lowercase ).sample

            if not return_dict:
                  return (dec,)

            return DecoderOutput(sample=lowercase )






