
import unittest

from transformers import XLMConfig, 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 (
          XLMForMultipleChoice,
          XLMForQuestionAnswering,
          XLMForQuestionAnsweringSimple,
          XLMForSequenceClassification,
          XLMForTokenClassification,
          XLMModel,
          XLMWithLMHeadModel,
      )
      from transformers.models.xlm.modeling_xlm import XLM_PRETRAINED_MODEL_ARCHIVE_LIST
class 			UpperCamelCase      :
     def __init__(     self  ,     snake_case__  ,     snake_case__=13  ,     snake_case__=7  ,     snake_case__=True  ,     snake_case__=True  ,     snake_case__=True  ,     snake_case__=True  ,     snake_case__=True  ,     snake_case__=False  ,     snake_case__=False  ,     snake_case__=False  ,     snake_case__=2  ,     snake_case__=99  ,     snake_case__=0  ,     snake_case__=32  ,     snake_case__=5  ,     snake_case__=4  ,     snake_case__=0.1  ,     snake_case__=0.1  ,     snake_case__=512  ,     snake_case__=2  ,     snake_case__=0.02  ,     snake_case__=2  ,     snake_case__=4  ,     snake_case__="last"  ,     snake_case__=True  ,     snake_case__=None  ,     snake_case__=0  ,     ):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				List[Any]							    = parent
       _SCREAMING_SNAKE_CASE   :				Any							    = batch_size
       _SCREAMING_SNAKE_CASE   :				Any							    = seq_length
       _SCREAMING_SNAKE_CASE   :				str							    = is_training
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = use_input_lengths
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = use_token_type_ids
       _SCREAMING_SNAKE_CASE   :				Dict							    = use_labels
       _SCREAMING_SNAKE_CASE   :				List[str]							    = gelu_activation
       _SCREAMING_SNAKE_CASE   :				Any							    = sinusoidal_embeddings
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = causal
       _SCREAMING_SNAKE_CASE   :				str							    = asm
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = n_langs
       _SCREAMING_SNAKE_CASE   :				Dict							    = vocab_size
       _SCREAMING_SNAKE_CASE   :				Tuple							    = n_special
       _SCREAMING_SNAKE_CASE   :				Dict							    = hidden_size
       _SCREAMING_SNAKE_CASE   :				str							    = num_hidden_layers
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = num_attention_heads
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = hidden_dropout_prob
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = attention_probs_dropout_prob
       _SCREAMING_SNAKE_CASE   :				List[str]							    = max_position_embeddings
       _SCREAMING_SNAKE_CASE   :				int							    = type_sequence_label_size
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = initializer_range
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = num_labels
       _SCREAMING_SNAKE_CASE   :				Tuple							    = num_choices
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = summary_type
       _SCREAMING_SNAKE_CASE   :				Dict							    = use_proj
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = scope
       _SCREAMING_SNAKE_CASE   :				int							    = bos_token_id
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


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

       _SCREAMING_SNAKE_CASE   :				List[str]							    = None
       if self.use_input_lengths:
         _SCREAMING_SNAKE_CASE   :				Any							    = (
             ids_tensor([self.batch_size]  ,     vocab_size=2				) + self.seq_length - 2
         )  # small variation of seq_length

       _SCREAMING_SNAKE_CASE   :				Any							    = None
       if self.use_token_type_ids:
         _SCREAMING_SNAKE_CASE   :				List[Any]							    = ids_tensor([self.batch_size, self.seq_length]  ,     self.n_langs				)

       _SCREAMING_SNAKE_CASE   :				int							    = None
       _SCREAMING_SNAKE_CASE   :				Dict							    = None
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = None
       if self.use_labels:
         _SCREAMING_SNAKE_CASE   :				Dict							    = ids_tensor([self.batch_size]  ,     self.type_sequence_label_size				)
         _SCREAMING_SNAKE_CASE   :				List[Any]							    = ids_tensor([self.batch_size, self.seq_length]  ,     self.num_labels				)
         _SCREAMING_SNAKE_CASE   :				Dict							    = ids_tensor([self.batch_size]  ,     2				).float()
         _SCREAMING_SNAKE_CASE   :				List[Any]							    = ids_tensor([self.batch_size]  ,     self.num_choices				)

       _SCREAMING_SNAKE_CASE   :				Dict							    = self.get_config()

       return (
           config,
           input_ids,
           token_type_ids,
           input_lengths,
           sequence_labels,
           token_labels,
           is_impossible_labels,
           choice_labels,
           input_mask,
       )
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       return XLMConfig(
           vocab_size=self.vocab_size  ,     n_special=self.n_special  ,     emb_dim=self.hidden_size  ,     n_layers=self.num_hidden_layers  ,     n_heads=self.num_attention_heads  ,     dropout=self.hidden_dropout_prob  ,     attention_dropout=self.attention_probs_dropout_prob  ,     gelu_activation=self.gelu_activation  ,     sinusoidal_embeddings=self.sinusoidal_embeddings  ,     asm=self.asm  ,     causal=self.causal  ,     n_langs=self.n_langs  ,     max_position_embeddings=self.max_position_embeddings  ,     initializer_range=self.initializer_range  ,     summary_type=self.summary_type  ,     use_proj=self.use_proj  ,     num_labels=self.num_labels  ,     bos_token_id=self.bos_token_id  ,     )
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     ):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Tuple							    = XLMModel(config=snake_case__				)
       model.to(snake_case__				)
       model.eval()
       _SCREAMING_SNAKE_CASE   :				Any							    = model(snake_case__  ,     lengths=snake_case__  ,     langs=snake_case__				)
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = model(snake_case__  ,     langs=snake_case__				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = model(snake_case__				)
       self.parent.assertEqual(result.last_hidden_state.shape  ,     (self.batch_size, self.seq_length, self.hidden_size)				)
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     ):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				int							    = XLMWithLMHeadModel(snake_case__				)
       model.to(snake_case__				)
       model.eval()

       _SCREAMING_SNAKE_CASE   :				Any							    = model(snake_case__  ,     token_type_ids=snake_case__  ,     labels=snake_case__				)
       self.parent.assertEqual(result.loss.shape  ,     ()				)
       self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.seq_length, self.vocab_size)				)
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     ):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Tuple							    = XLMForQuestionAnsweringSimple(snake_case__				)
       model.to(snake_case__				)
       model.eval()

       _SCREAMING_SNAKE_CASE   :				List[str]							    = model(snake_case__				)

       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = model(snake_case__  ,     start_positions=snake_case__  ,     end_positions=snake_case__				)
       _SCREAMING_SNAKE_CASE   :				Tuple							    = outputs
       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 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     ):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				List[str]							    = XLMForQuestionAnswering(snake_case__				)
       model.to(snake_case__				)
       model.eval()

       _SCREAMING_SNAKE_CASE   :				str							    = model(snake_case__				)

       _SCREAMING_SNAKE_CASE   :				Tuple							    = model(
           snake_case__  ,     start_positions=snake_case__  ,     end_positions=snake_case__  ,     cls_index=snake_case__  ,     is_impossible=snake_case__  ,     p_mask=snake_case__  ,     )

       _SCREAMING_SNAKE_CASE   :				List[str]							    = model(
           snake_case__  ,     start_positions=snake_case__  ,     end_positions=snake_case__  ,     cls_index=snake_case__  ,     is_impossible=snake_case__  ,     )

       (_SCREAMING_SNAKE_CASE		)   :				Any							    = result_with_labels.to_tuple()

       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = model(snake_case__  ,     start_positions=snake_case__  ,     end_positions=snake_case__				)

       (_SCREAMING_SNAKE_CASE		)   :				List[Any]							    = result_with_labels.to_tuple()

       self.parent.assertEqual(result_with_labels.loss.shape  ,     ()				)
       self.parent.assertEqual(result.start_top_log_probs.shape  ,     (self.batch_size, model.config.start_n_top)				)
       self.parent.assertEqual(result.start_top_index.shape  ,     (self.batch_size, model.config.start_n_top)				)
       self.parent.assertEqual(
           result.end_top_log_probs.shape  ,     (self.batch_size, model.config.start_n_top * model.config.end_n_top)				)
       self.parent.assertEqual(
           result.end_top_index.shape  ,     (self.batch_size, model.config.start_n_top * model.config.end_n_top)				)
       self.parent.assertEqual(result.cls_logits.shape  ,     (self.batch_size,)				)
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     ):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = XLMForSequenceClassification(snake_case__				)
       model.to(snake_case__				)
       model.eval()

       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = model(snake_case__				)
       _SCREAMING_SNAKE_CASE   :				Tuple							    = model(snake_case__  ,     labels=snake_case__				)
       self.parent.assertEqual(result.loss.shape  ,     ()				)
       self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.type_sequence_label_size)				)
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     ):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self.num_labels
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = XLMForTokenClassification(snake_case__				)
       model.to(snake_case__				)
       model.eval()

       _SCREAMING_SNAKE_CASE   :				str							    = model(snake_case__  ,     attention_mask=snake_case__  ,     labels=snake_case__				)
       self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.seq_length, self.num_labels)				)
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     ):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = self.num_choices
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = XLMForMultipleChoice(config=snake_case__				)
       model.to(snake_case__				)
       model.eval()
       _SCREAMING_SNAKE_CASE   :				Any							    = input_ids.unsqueeze(1				).expand(-1  ,     self.num_choices  ,     -1				).contiguous()
       _SCREAMING_SNAKE_CASE   :				Dict							    = token_type_ids.unsqueeze(1				).expand(-1  ,     self.num_choices  ,     -1				).contiguous()
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = input_mask.unsqueeze(1				).expand(-1  ,     self.num_choices  ,     -1				).contiguous()
       _SCREAMING_SNAKE_CASE   :				List[Any]							    = model(
           snake_case__  ,     attention_mask=snake_case__  ,     token_type_ids=snake_case__  ,     labels=snake_case__  ,     )
       self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.num_choices)				)




     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				int							    = self.prepare_config_and_inputs()
       (
           _SCREAMING_SNAKE_CASE		
       )   :				Dict							    = config_and_inputs
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = {"input_ids": input_ids, "token_type_ids": token_type_ids, "lengths": input_lengths}
       return config, inputs_dict
@require_torch
class 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE      ,       __SCREAMING_SNAKE_CASE      ,       __SCREAMING_SNAKE_CASE      ,       unittest.TestCase  ):
     A__      =			(
         (
             XLMModel,
             XLMWithLMHeadModel,
             XLMForQuestionAnswering,
             XLMForSequenceClassification,
             XLMForQuestionAnsweringSimple,
             XLMForTokenClassification,
             XLMForMultipleChoice,
         )
         if is_torch_available()
         else ()
     )
     A__      =			(
         (XLMWithLMHeadModel,) if is_torch_available() else ()
     )  # TODO (PVP): Check other models whether language generation is also applicable
     A__      =			(
         {
             """feature-extraction""": XLMModel,
             """fill-mask""": XLMWithLMHeadModel,
             """question-answering""": XLMForQuestionAnsweringSimple,
             """text-classification""": XLMForSequenceClassification,
             """text-generation""": XLMWithLMHeadModel,
             """token-classification""": XLMForTokenClassification,
             """zero-shot""": XLMForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__				):




       """simple docstring"""


       if (
           pipeline_test_casse_name == "QAPipelineTests"
           and tokenizer_name is not None
           and not tokenizer_name.endswith("Fast"				)
       ):
         # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
         # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
         # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
         return True

       return False
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__=False				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				List[Any]							    = super()._prepare_for_class(snake_case__  ,     snake_case__  ,     return_labels=snake_case__				)

       if return_labels:
         if model_class.__name__ == "XLMForQuestionAnswering":
           _SCREAMING_SNAKE_CASE   :				int							    = torch.zeros(
               self.model_tester.batch_size  ,     dtype=torch.long  ,     device=snake_case__				)
           _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = torch.zeros(
               self.model_tester.batch_size  ,     dtype=torch.long  ,     device=snake_case__				)

       return inputs_dict
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = XLMModelTester(self				)
       _SCREAMING_SNAKE_CASE   :				Any							    = ConfigTester(self  ,     config_class=snake_case__  ,     emb_dim=37				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       self.config_tester.run_common_tests()
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				List[Any]							    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_xlm_model(*snake_case__				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				int							    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_xlm_lm_head(*snake_case__				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_xlm_simple_qa(*snake_case__				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				List[str]							    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_xlm_qa(*snake_case__				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				int							    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_xlm_sequence_classif(*snake_case__				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				str							    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_xlm_token_classif(*snake_case__				)
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_xlm_for_multiple_choice(*snake_case__				)
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__=False  ,     snake_case__=1				):




       """simple docstring"""


       self.assertIsInstance(snake_case__  ,     snake_case__				)
       self.assertListEqual(
           [isinstance(snake_case__  ,     snake_case__				) for iter_attentions in attentions]  ,     [True] * len(snake_case__				)				)
       self.assertEqual(len(snake_case__				)  ,     (max_length - min_length) * num_beam_groups				)

       for idx, iter_attentions in enumerate(snake_case__				):
         # adds PAD dummy token
         _SCREAMING_SNAKE_CASE   :				Dict							    = min_length + idx + 1
         _SCREAMING_SNAKE_CASE   :				List[Any]							    = min_length + idx + 1

         _SCREAMING_SNAKE_CASE   :				Optional[int]							    = (
             batch_size * num_beam_groups,
             config.num_attention_heads,
             tgt_len,
             src_len,
         )
         # check attn size
         self.assertListEqual(
             [layer_attention.shape for layer_attention in iter_attentions]  ,     [expected_shape] * len(snake_case__				)				)
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__  ,     snake_case__=False  ,     snake_case__=1				):




       """simple docstring"""


       self.assertIsInstance(snake_case__  ,     snake_case__				)
       self.assertListEqual(
           [isinstance(snake_case__  ,     snake_case__				) for iter_hidden_states in hidden_states]  ,     [True] * len(snake_case__				)  ,     )
       self.assertEqual(len(snake_case__				)  ,     (max_length - min_length) * num_beam_groups				)

       for idx, iter_hidden_states in enumerate(snake_case__				):
         # adds PAD dummy token
         _SCREAMING_SNAKE_CASE   :				Dict							    = min_length + idx + 1
         _SCREAMING_SNAKE_CASE   :				List[str]							    = (batch_size * num_beam_groups, seq_len, config.hidden_size)
         # check hidden size
         self.assertListEqual(
             [layer_hidden_states.shape for layer_hidden_states in iter_hidden_states]  ,     [expected_shape] * len(snake_case__				)  ,     )
       pass




     @slow
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         _SCREAMING_SNAKE_CASE   :				List[str]							    = XLMModel.from_pretrained(snake_case__				)
         self.assertIsNotNone(snake_case__				)
@require_torch
class 			UpperCamelCase      (							unittest.TestCase  ):
     @slow
     def 				__SCREAMING_SNAKE_CASE							(     self				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = XLMWithLMHeadModel.from_pretrained("xlm-mlm-en-2048"				)
       model.to(snake_case__				)
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = torch.tensor([[14, 447]]  ,     dtype=torch.long  ,     device=snake_case__				)  # the president
       _SCREAMING_SNAKE_CASE   :				List[str]							    = [
           14,
           447,
           14,
           447,
           14,
           447,
           14,
           447,
           14,
           447,
           14,
           447,
           14,
           447,
           14,
           447,
           14,
           447,
           14,
           447,
       ]  # the president the president the president the president the president the president the president the president the president the president
       # TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
       _SCREAMING_SNAKE_CASE   :				Any							    = model.generate(snake_case__  ,     do_sample=snake_case__				)
       self.assertListEqual(output_ids[0].cpu().numpy().tolist()  ,     snake_case__				)





"""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 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE  ):
     A__      =			42
class 			UpperCamelCase      (							__SCREAMING_SNAKE_CASE      ,       __SCREAMING_SNAKE_CASE  ):
     @register_to_config
     def __init__(     self  ,     snake_case__ = 3  ,     snake_case__ = 3  ,     snake_case__ = ("DownEncoderBlock2D",)  ,     snake_case__ = ("UpDecoderBlock2D",)  ,     snake_case__ = (64,)  ,     snake_case__ = 1  ,     snake_case__ = "silu"  ,     snake_case__ = 3  ,     snake_case__ = 32  ,     snake_case__ = 256  ,     snake_case__ = 32  ,     snake_case__ = None  ,     snake_case__ = 0.18_215  ,     snake_case__ = "group"  ,     ):




       """simple docstring"""


       super().__init__()

       # pass init params to Encoder
       _SCREAMING_SNAKE_CASE   :				Optional[Any]							    = Encoder(
           in_channels=snake_case__  ,     out_channels=snake_case__  ,     down_block_types=snake_case__  ,     block_out_channels=snake_case__  ,     layers_per_block=snake_case__  ,     act_fn=snake_case__  ,     norm_num_groups=snake_case__  ,     double_z=snake_case__  ,     )

       _SCREAMING_SNAKE_CASE   :				Dict							    = vq_embed_dim if vq_embed_dim is not None else latent_channels

       _SCREAMING_SNAKE_CASE   :				Tuple							    = nn.Convad(snake_case__  ,     snake_case__  ,     1				)
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = VectorQuantizer(snake_case__  ,     snake_case__  ,     beta=0.25  ,     remap=snake_case__  ,     sane_index_shape=snake_case__				)
       _SCREAMING_SNAKE_CASE   :				List[str]							    = nn.Convad(snake_case__  ,     snake_case__  ,     1				)

       # pass init params to Decoder
       _SCREAMING_SNAKE_CASE   :				Optional[int]							    = Decoder(
           in_channels=snake_case__  ,     out_channels=snake_case__  ,     up_block_types=snake_case__  ,     block_out_channels=snake_case__  ,     layers_per_block=snake_case__  ,     act_fn=snake_case__  ,     norm_num_groups=snake_case__  ,     norm_type=snake_case__  ,     )
     @apply_forward_hook
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__ = True				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				List[str]							    = self.encoder(snake_case__				)
       _SCREAMING_SNAKE_CASE   :				Dict							    = self.quant_conv(snake_case__				)

       if not return_dict:
         return (h,)

       return VQEncoderOutput(latents=snake_case__				)
     @apply_forward_hook
     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__ = False  ,     snake_case__ = True				):




       """simple docstring"""


       if not force_not_quantize:
         _SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE    ,		_SCREAMING_SNAKE_CASE   :				int							    = self.quantize(snake_case__				)
       else:
         _SCREAMING_SNAKE_CASE   :				int							    = h
       _SCREAMING_SNAKE_CASE   :				Dict							    = self.post_quant_conv(snake_case__				)
       _SCREAMING_SNAKE_CASE   :				int							    = self.decoder(snake_case__  ,     quant if self.config.norm_type == "spatial" else None				)

       if not return_dict:
         return (dec,)

       return DecoderOutput(sample=snake_case__				)




     def 				__SCREAMING_SNAKE_CASE							(     self  ,     snake_case__  ,     snake_case__ = True				):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE   :				Tuple							    = sample
       _SCREAMING_SNAKE_CASE   :				Union[str, Any]							    = self.encode(snake_case__				).latents
       _SCREAMING_SNAKE_CASE   :				Any							    = self.decode(snake_case__				).sample

       if not return_dict:
         return (dec,)

       return DecoderOutput(sample=snake_case__				)



