
'''simple docstring'''






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   :							str,__A   :							int,__A   :							Union[str, Any]=1_3,__A   :							Optional[int]=7,__A   :							Dict=True,__A   :							Dict=True,__A   :							Optional[int]=True,__A   :							List[str]=True,__A   :							int=True,__A   :							int=False,__A   :							List[Any]=False,__A   :							Union[str, Any]=False,__A   :							Union[str, Any]=2,__A   :							str=9_9,__A   :							List[str]=0,__A   :							Any=3_2,__A   :							Optional[Any]=5,__A   :							Union[str, Any]=4,__A   :							List[Any]=0.1,__A   :							Tuple=0.1,__A   :							Dict=5_1_2,__A   :							Optional[int]=2,__A   :							List[str]=0.02,__A   :							Tuple=2,__A   :							Optional[int]=4,__A   :							List[Any]="last",__A   :							Optional[Any]=True,__A   :							Any=None,__A   :							Optional[Any]=0,):
          _lowerCamelCase					: Any		    = parent
          _lowerCamelCase					: Union[str, Any]		    = batch_size
          _lowerCamelCase					: int		    = seq_length
          _lowerCamelCase					: Optional[int]		    = is_training
          _lowerCamelCase					: Any		    = use_input_lengths
          _lowerCamelCase					: Any		    = use_token_type_ids
          _lowerCamelCase					: Optional[Any]		    = use_labels
          _lowerCamelCase					: Union[str, Any]		    = gelu_activation
          _lowerCamelCase					: Optional[Any]		    = sinusoidal_embeddings
          _lowerCamelCase					: Optional[int]		    = causal
          _lowerCamelCase					: str		    = asm
          _lowerCamelCase					: Optional[int]		    = n_langs
          _lowerCamelCase					: List[Any]		    = vocab_size
          _lowerCamelCase					: Union[str, Any]		    = n_special
          _lowerCamelCase					: Any		    = hidden_size
          _lowerCamelCase					: Any		    = num_hidden_layers
          _lowerCamelCase					: Optional[int]		    = num_attention_heads
          _lowerCamelCase					: str		    = hidden_dropout_prob
          _lowerCamelCase					: Dict		    = attention_probs_dropout_prob
          _lowerCamelCase					: Dict		    = max_position_embeddings
          _lowerCamelCase					: Tuple		    = type_sequence_label_size
          _lowerCamelCase					: Any		    = initializer_range
          _lowerCamelCase					: Tuple		    = num_labels
          _lowerCamelCase					: Any		    = num_choices
          _lowerCamelCase					: List[str]		    = summary_type
          _lowerCamelCase					: Union[str, Any]		    = use_proj
          _lowerCamelCase					: Dict		    = scope
          _lowerCamelCase					: List[Any]		    = bos_token_id



    def 	lowerCamelCase_	(						self   :							List[Any]  ):
          _lowerCamelCase					: Optional[int]		    = ids_tensor([self.batch_size, self.seq_length],self.vocab_size  )
          _lowerCamelCase					: Dict		    = random_attention_mask([self.batch_size, self.seq_length]  )

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

          _lowerCamelCase					: Union[str, Any]		    = None
          if self.use_token_type_ids:
                _lowerCamelCase					: Any		    = ids_tensor([self.batch_size, self.seq_length],self.n_langs  )

          _lowerCamelCase					: Optional[int]		    = None
          _lowerCamelCase					: Dict		    = None
          _lowerCamelCase					: int		    = None
          if self.use_labels:
                _lowerCamelCase					: str		    = ids_tensor([self.batch_size],self.type_sequence_label_size  )
                _lowerCamelCase					: Dict		    = ids_tensor([self.batch_size, self.seq_length],self.num_labels  )
                _lowerCamelCase					: int		    = ids_tensor([self.batch_size],2  ).float()
                _lowerCamelCase					: Dict		    = ids_tensor([self.batch_size],self.num_choices  )

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

          return (
              config,
              input_ids,
              token_type_ids,
              input_lengths,
              sequence_labels,
              token_labels,
              is_impossible_labels,
              choice_labels,
              input_mask,
          )



    def 	lowerCamelCase_	(						self   :							Union[str, Any]  ):
          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 	lowerCamelCase_	(						self   :							Union[str, Any],__A   :							Optional[int],__A   :							Optional[int],__A   :							int,__A   :							List[Any],__A   :							List[str],__A   :							Optional[Any],__A   :							Optional[int],__A   :							Union[str, Any],__A   :							List[Any],):
          _lowerCamelCase					: List[Any]		    = XLMModel(config=__A  )
          model.to(__A  )
          model.eval()
          _lowerCamelCase					: int		    = model(__A,lengths=__A,langs=__A  )
          _lowerCamelCase					: str		    = model(__A,langs=__A  )
          _lowerCamelCase					: str		    = model(__A  )
          self.parent.assertEqual(result.last_hidden_state.shape,(self.batch_size, self.seq_length, self.hidden_size)  )



    def 	lowerCamelCase_	(						self   :							Dict,__A   :							Optional[int],__A   :							int,__A   :							Optional[int],__A   :							Union[str, Any],__A   :							Union[str, Any],__A   :							int,__A   :							Optional[int],__A   :							int,__A   :							List[str],):
          _lowerCamelCase					: List[str]		    = XLMWithLMHeadModel(__A  )
          model.to(__A  )
          model.eval()

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



    def 	lowerCamelCase_	(						self   :							Optional[int],__A   :							int,__A   :							List[Any],__A   :							Optional[int],__A   :							Union[str, Any],__A   :							str,__A   :							Tuple,__A   :							List[Any],__A   :							str,__A   :							List[Any],):
          _lowerCamelCase					: List[str]		    = XLMForQuestionAnsweringSimple(__A  )
          model.to(__A  )
          model.eval()

          _lowerCamelCase					: Dict		    = model(__A  )

          _lowerCamelCase					: str		    = model(__A,start_positions=__A,end_positions=__A  )
          _lowerCamelCase					: Optional[int]		    = 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 	lowerCamelCase_	(						self   :							Tuple,__A   :							int,__A   :							Union[str, Any],__A   :							Optional[Any],__A   :							Optional[int],__A   :							Optional[Any],__A   :							Dict,__A   :							int,__A   :							Union[str, Any],__A   :							List[Any],):
          _lowerCamelCase					: List[str]		    = XLMForQuestionAnswering(__A  )
          model.to(__A  )
          model.eval()

          _lowerCamelCase					: Dict		    = model(__A  )

          _lowerCamelCase					: Dict		    = model(
              __A,start_positions=__A,end_positions=__A,cls_index=__A,is_impossible=__A,p_mask=__A,)

          _lowerCamelCase					: str		    = model(
              __A,start_positions=__A,end_positions=__A,cls_index=__A,is_impossible=__A,)

          ((_lowerCamelCase)    ,		)					: int		    = result_with_labels.to_tuple()

          _lowerCamelCase					: int		    = model(__A,start_positions=__A,end_positions=__A  )

          ((_lowerCamelCase)    ,		)					: Tuple		    = 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 	lowerCamelCase_	(						self   :							List[Any],__A   :							Optional[int],__A   :							str,__A   :							int,__A   :							Any,__A   :							str,__A   :							str,__A   :							Tuple,__A   :							int,__A   :							List[str],):
          _lowerCamelCase					: Union[str, Any]		    = XLMForSequenceClassification(__A  )
          model.to(__A  )
          model.eval()

          _lowerCamelCase					: Optional[Any]		    = model(__A  )
          _lowerCamelCase					: Any		    = model(__A,labels=__A  )
          self.parent.assertEqual(result.loss.shape,()  )
          self.parent.assertEqual(result.logits.shape,(self.batch_size, self.type_sequence_label_size)  )



    def 	lowerCamelCase_	(						self   :							List[Any],__A   :							Dict,__A   :							Union[str, Any],__A   :							List[str],__A   :							List[str],__A   :							Optional[int],__A   :							Optional[Any],__A   :							List[str],__A   :							str,__A   :							Optional[int],):
          _lowerCamelCase					: int		    = self.num_labels
          _lowerCamelCase					: List[Any]		    = XLMForTokenClassification(__A  )
          model.to(__A  )
          model.eval()

          _lowerCamelCase					: Optional[int]		    = model(__A,attention_mask=__A,labels=__A  )
          self.parent.assertEqual(result.logits.shape,(self.batch_size, self.seq_length, self.num_labels)  )



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



    def 	lowerCamelCase_	(						self   :							Optional[int]  ):
          _lowerCamelCase					: str		    = self.prepare_config_and_inputs()
          (
              (
              _lowerCamelCase
          )    ,		(
              _lowerCamelCase
          )    ,		(
              _lowerCamelCase
          )    ,		(
              _lowerCamelCase
          )    ,		(
              _lowerCamelCase
          )    ,		(
              _lowerCamelCase
          )    ,		(
              _lowerCamelCase
          )    ,		(
              _lowerCamelCase
          )    ,		(
              _lowerCamelCase
          )    ,		
          )					: Optional[int]		    = config_and_inputs
          _lowerCamelCase					: List[Any]		    = {"input_ids": input_ids, "token_type_ids": token_type_ids, "lengths": input_lengths}
          return config, inputs_dict


@require_torch
class 			UpperCAmelCase__    (							A    ,						A    ,						A    ,						unittest.TestCase			):
    lowerCAmelCase_					     =		(
        (
            XLMModel,
            XLMWithLMHeadModel,
            XLMForQuestionAnswering,
            XLMForSequenceClassification,
            XLMForQuestionAnsweringSimple,
            XLMForTokenClassification,
            XLMForMultipleChoice,
        )
        if is_torch_available()
        else ()
    )
    lowerCAmelCase_					     =		(
        (XLMWithLMHeadModel,) if is_torch_available() else ()
    )  # TODO (PVP): Check other models whether language generation is also applicable
    lowerCAmelCase_					     =		(
        {
            '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 	lowerCamelCase_	(						self   :							int,__A   :							Tuple,__A   :							Union[str, Any],__A   :							Optional[int],__A   :							Union[str, Any],__A   :							Any  ):
          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 	lowerCamelCase_	(						self   :							List[Any],__A   :							List[Any],__A   :							Dict,__A   :							Union[str, Any]=False  ):
          _lowerCamelCase					: List[str]		    = super()._prepare_for_class(__A,__A,return_labels=__A  )

          if return_labels:
                if model_class.__name__ == "XLMForQuestionAnswering":
                      _lowerCamelCase					: Optional[Any]		    = torch.zeros(
                          self.model_tester.batch_size,dtype=torch.long,device=__A  )
                      _lowerCamelCase					: Dict		    = torch.zeros(
                          self.model_tester.batch_size,dtype=torch.long,device=__A  )

          return inputs_dict



    def 	lowerCamelCase_	(						self   :							str  ):
          _lowerCamelCase					: int		    = XLMModelTester(self  )
          _lowerCamelCase					: Optional[int]		    = ConfigTester(self,config_class=__A,emb_dim=3_7  )



    def 	lowerCamelCase_	(						self   :							str  ):
          self.config_tester.run_common_tests()



    def 	lowerCamelCase_	(						self   :							Union[str, Any]  ):
          _lowerCamelCase					: Optional[int]		    = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_model(*__A  )



    def 	lowerCamelCase_	(						self   :							Union[str, Any]  ):
          _lowerCamelCase					: int		    = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_lm_head(*__A  )



    def 	lowerCamelCase_	(						self   :							int  ):
          _lowerCamelCase					: List[Any]		    = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_simple_qa(*__A  )



    def 	lowerCamelCase_	(						self   :							List[str]  ):
          _lowerCamelCase					: Tuple		    = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_qa(*__A  )



    def 	lowerCamelCase_	(						self   :							List[str]  ):
          _lowerCamelCase					: Dict		    = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_sequence_classif(*__A  )



    def 	lowerCamelCase_	(						self   :							Dict  ):
          _lowerCamelCase					: Any		    = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_token_classif(*__A  )



    def 	lowerCamelCase_	(						self   :							List[Any]  ):
          _lowerCamelCase					: Any		    = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_for_multiple_choice(*__A  )



    def 	lowerCamelCase_	(						self   :							str,__A   :							Optional[Any],__A   :							str,__A   :							List[str],__A   :							Tuple,__A   :							Optional[Any],__A   :							Optional[int]=False,__A   :							Optional[int]=1  ):
          self.assertIsInstance(__A,__A  )
          self.assertListEqual(
              [isinstance(__A,__A  ) for iter_attentions in attentions],[True] * len(__A  )  )
          self.assertEqual(len(__A  ),(max_length - min_length) * num_beam_groups  )

          for idx, iter_attentions in enumerate(__A  ):
                # adds PAD dummy token
                _lowerCamelCase					: Optional[Any]		    = min_length + idx + 1
                _lowerCamelCase					: int		    = min_length + idx + 1

                _lowerCamelCase					: 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(__A  )  )



    def 	lowerCamelCase_	(						self   :							Optional[Any],__A   :							Union[str, Any],__A   :							Dict,__A   :							List[Any],__A   :							Dict,__A   :							List[str],__A   :							Union[str, Any]=False,__A   :							Dict=1  ):
          self.assertIsInstance(__A,__A  )
          self.assertListEqual(
              [isinstance(__A,__A  ) for iter_hidden_states in hidden_states],[True] * len(__A  ),)
          self.assertEqual(len(__A  ),(max_length - min_length) * num_beam_groups  )

          for idx, iter_hidden_states in enumerate(__A  ):
                # adds PAD dummy token
                _lowerCamelCase					: int		    = min_length + idx + 1
                _lowerCamelCase					: Any		    = (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(__A  ),)
          pass



    @slow
    def 	lowerCamelCase_	(						self   :							Tuple  ):
          for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                _lowerCamelCase					: Dict		    = XLMModel.from_pretrained(__A  )
                self.assertIsNotNone(__A  )


@require_torch
class 			UpperCAmelCase__    (							unittest.TestCase			):



    @slow
    def 	lowerCamelCase_	(						self   :							str  ):
          _lowerCamelCase					: int		    = XLMWithLMHeadModel.from_pretrained("xlm-mlm-en-2048"  )
          model.to(__A  )
          _lowerCamelCase					: int		    = torch.tensor([[1_4, 4_4_7]],dtype=torch.long,device=__A  )  # the president
          _lowerCamelCase					: Optional[Any]		    = [
              1_4,
              4_4_7,
              1_4,
              4_4_7,
              1_4,
              4_4_7,
              1_4,
              4_4_7,
              1_4,
              4_4_7,
              1_4,
              4_4_7,
              1_4,
              4_4_7,
              1_4,
              4_4_7,
              1_4,
              4_4_7,
              1_4,
              4_4_7,
          ]  # 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
          _lowerCamelCase					: Optional[Any]		    = model.generate(__A,do_sample=__A  )
          self.assertListEqual(output_ids[0].cpu().numpy().tolist(),__A  )

import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD




class 			UpperCAmelCase_     (					nn.Module	):


  '''simple docstring'''







  a__       			=						42
  a__       			=						42
  a__       			=						0.0
  a__       			=						1
  a__       			=						1
  a__       			=						True
  a__       			=						False
  a__       			=						False
  a__       			=						False
  a__       			=						jnp.floataa
  def 				_lowercase				(					self     :					Optional[Any]       )     ->      Tuple:







    """simple docstring"""






    __magic_name__								=			[]
    __magic_name__								=			[]

    for i in range(self.num_layers       ):
      __magic_name__								=			self.in_channels if i == 0 else self.out_channels

      __magic_name__								=			FlaxResnetBlockaD(
          in_channels=UpperCamelCase__     ,				out_channels=self.out_channels     ,				dropout_prob=self.dropout     ,				dtype=self.dtype     ,				)
      resnets.append(UpperCamelCase__       )

      __magic_name__								=			FlaxTransformeraDModel(
          in_channels=self.out_channels     ,				n_heads=self.num_attention_heads     ,				d_head=self.out_channels // self.num_attention_heads     ,				depth=1     ,				use_linear_projection=self.use_linear_projection     ,				only_cross_attention=self.only_cross_attention     ,				use_memory_efficient_attention=self.use_memory_efficient_attention     ,				dtype=self.dtype     ,				)
      attentions.append(UpperCamelCase__       )

    __magic_name__								=			resnets
    __magic_name__								=			attentions

    if self.add_downsample:
      __magic_name__								=			FlaxDownsampleaD(self.out_channels     ,				dtype=self.dtype       )





  def __call__(					self     :					List[Any]     ,				UpperCamelCase__     :					Optional[int]     ,				UpperCamelCase__     :					Tuple     ,				UpperCamelCase__     :					Dict     ,				UpperCamelCase__     :					Optional[int]=True       )     ->      Any:







    """simple docstring"""






    __magic_name__								=			()

    for resnet, attn in zip(self.resnets     ,				self.attentions       ):
      __magic_name__								=			resnet(UpperCamelCase__     ,				UpperCamelCase__     ,				deterministic=UpperCamelCase__       )
      __magic_name__								=			attn(UpperCamelCase__     ,				UpperCamelCase__     ,				deterministic=UpperCamelCase__       )
      output_states += (hidden_states,)

    if self.add_downsample:
      __magic_name__								=			self.downsamplers_a(UpperCamelCase__       )
      output_states += (hidden_states,)

    return hidden_states, output_states







class 			UpperCAmelCase_     (					nn.Module	):


  '''simple docstring'''







  a__       			=						42
  a__       			=						42
  a__       			=						0.0
  a__       			=						1
  a__       			=						True
  a__       			=						jnp.floataa
  def 				_lowercase				(					self     :					str       )     ->      Optional[Any]:







    """simple docstring"""






    __magic_name__								=			[]

    for i in range(self.num_layers       ):
      __magic_name__								=			self.in_channels if i == 0 else self.out_channels

      __magic_name__								=			FlaxResnetBlockaD(
          in_channels=UpperCamelCase__     ,				out_channels=self.out_channels     ,				dropout_prob=self.dropout     ,				dtype=self.dtype     ,				)
      resnets.append(UpperCamelCase__       )
    __magic_name__								=			resnets

    if self.add_downsample:
      __magic_name__								=			FlaxDownsampleaD(self.out_channels     ,				dtype=self.dtype       )





  def __call__(					self     :					str     ,				UpperCamelCase__     :					Tuple     ,				UpperCamelCase__     :					Union[str, Any]     ,				UpperCamelCase__     :					int=True       )     ->      Dict:







    """simple docstring"""






    __magic_name__								=			()

    for resnet in self.resnets:
      __magic_name__								=			resnet(UpperCamelCase__     ,				UpperCamelCase__     ,				deterministic=UpperCamelCase__       )
      output_states += (hidden_states,)

    if self.add_downsample:
      __magic_name__								=			self.downsamplers_a(UpperCamelCase__       )
      output_states += (hidden_states,)

    return hidden_states, output_states







class 			UpperCAmelCase_     (					nn.Module	):


  '''simple docstring'''







  a__       			=						42
  a__       			=						42
  a__       			=						42
  a__       			=						0.0
  a__       			=						1
  a__       			=						1
  a__       			=						True
  a__       			=						False
  a__       			=						False
  a__       			=						False
  a__       			=						jnp.floataa
  def 				_lowercase				(					self     :					Optional[Any]       )     ->      Any:







    """simple docstring"""






    __magic_name__								=			[]
    __magic_name__								=			[]

    for i in range(self.num_layers       ):
      __magic_name__								=			self.in_channels if (i == self.num_layers - 1) else self.out_channels
      __magic_name__								=			self.prev_output_channel if i == 0 else self.out_channels

      __magic_name__								=			FlaxResnetBlockaD(
          in_channels=resnet_in_channels + res_skip_channels     ,				out_channels=self.out_channels     ,				dropout_prob=self.dropout     ,				dtype=self.dtype     ,				)
      resnets.append(UpperCamelCase__       )

      __magic_name__								=			FlaxTransformeraDModel(
          in_channels=self.out_channels     ,				n_heads=self.num_attention_heads     ,				d_head=self.out_channels // self.num_attention_heads     ,				depth=1     ,				use_linear_projection=self.use_linear_projection     ,				only_cross_attention=self.only_cross_attention     ,				use_memory_efficient_attention=self.use_memory_efficient_attention     ,				dtype=self.dtype     ,				)
      attentions.append(UpperCamelCase__       )

    __magic_name__								=			resnets
    __magic_name__								=			attentions

    if self.add_upsample:
      __magic_name__								=			FlaxUpsampleaD(self.out_channels     ,				dtype=self.dtype       )





  def __call__(					self     :					Optional[int]     ,				UpperCamelCase__     :					Any     ,				UpperCamelCase__     :					Dict     ,				UpperCamelCase__     :					Optional[int]     ,				UpperCamelCase__     :					Optional[Any]     ,				UpperCamelCase__     :					Dict=True       )     ->      Union[str, Any]:







    """simple docstring"""






    for resnet, attn in zip(self.resnets     ,				self.attentions       ):
      # pop res hidden states
      __magic_name__								=			res_hidden_states_tuple[-1]
      __magic_name__								=			res_hidden_states_tuple[:-1]
      __magic_name__								=			jnp.concatenate((hidden_states, res_hidden_states)     ,				axis=-1       )

      __magic_name__								=			resnet(UpperCamelCase__     ,				UpperCamelCase__     ,				deterministic=UpperCamelCase__       )
      __magic_name__								=			attn(UpperCamelCase__     ,				UpperCamelCase__     ,				deterministic=UpperCamelCase__       )

    if self.add_upsample:
      __magic_name__								=			self.upsamplers_a(UpperCamelCase__       )

    return hidden_states







class 			UpperCAmelCase_     (					nn.Module	):


  '''simple docstring'''







  a__       			=						42
  a__       			=						42
  a__       			=						42
  a__       			=						0.0
  a__       			=						1
  a__       			=						True
  a__       			=						jnp.floataa
  def 				_lowercase				(					self     :					List[Any]       )     ->      str:







    """simple docstring"""






    __magic_name__								=			[]

    for i in range(self.num_layers       ):
      __magic_name__								=			self.in_channels if (i == self.num_layers - 1) else self.out_channels
      __magic_name__								=			self.prev_output_channel if i == 0 else self.out_channels

      __magic_name__								=			FlaxResnetBlockaD(
          in_channels=resnet_in_channels + res_skip_channels     ,				out_channels=self.out_channels     ,				dropout_prob=self.dropout     ,				dtype=self.dtype     ,				)
      resnets.append(UpperCamelCase__       )

    __magic_name__								=			resnets

    if self.add_upsample:
      __magic_name__								=			FlaxUpsampleaD(self.out_channels     ,				dtype=self.dtype       )





  def __call__(					self     :					List[Any]     ,				UpperCamelCase__     :					str     ,				UpperCamelCase__     :					Union[str, Any]     ,				UpperCamelCase__     :					Optional[Any]     ,				UpperCamelCase__     :					str=True       )     ->      Dict:







    """simple docstring"""






    for resnet in self.resnets:
      # pop res hidden states
      __magic_name__								=			res_hidden_states_tuple[-1]
      __magic_name__								=			res_hidden_states_tuple[:-1]
      __magic_name__								=			jnp.concatenate((hidden_states, res_hidden_states)     ,				axis=-1       )

      __magic_name__								=			resnet(UpperCamelCase__     ,				UpperCamelCase__     ,				deterministic=UpperCamelCase__       )

    if self.add_upsample:
      __magic_name__								=			self.upsamplers_a(UpperCamelCase__       )

    return hidden_states







class 			UpperCAmelCase_     (					nn.Module	):


  '''simple docstring'''







  a__       			=						42
  a__       			=						0.0
  a__       			=						1
  a__       			=						1
  a__       			=						False
  a__       			=						False
  a__       			=						jnp.floataa
  def 				_lowercase				(					self     :					List[Any]       )     ->      Optional[Any]:







    """simple docstring"""






    __magic_name__								=			[
        FlaxResnetBlockaD(
            in_channels=self.in_channels     ,				out_channels=self.in_channels     ,				dropout_prob=self.dropout     ,				dtype=self.dtype     ,				)
    ]

    __magic_name__								=			[]

    for _ in range(self.num_layers       ):
      __magic_name__								=			FlaxTransformeraDModel(
          in_channels=self.in_channels     ,				n_heads=self.num_attention_heads     ,				d_head=self.in_channels // self.num_attention_heads     ,				depth=1     ,				use_linear_projection=self.use_linear_projection     ,				use_memory_efficient_attention=self.use_memory_efficient_attention     ,				dtype=self.dtype     ,				)
      attentions.append(UpperCamelCase__       )

      __magic_name__								=			FlaxResnetBlockaD(
          in_channels=self.in_channels     ,				out_channels=self.in_channels     ,				dropout_prob=self.dropout     ,				dtype=self.dtype     ,				)
      resnets.append(UpperCamelCase__       )

    __magic_name__								=			resnets
    __magic_name__								=			attentions





  def __call__(					self     :					Optional[int]     ,				UpperCamelCase__     :					Any     ,				UpperCamelCase__     :					Any     ,				UpperCamelCase__     :					List[Any]     ,				UpperCamelCase__     :					int=True       )     ->      str:







    """simple docstring"""






    __magic_name__								=			self.resnets[0](UpperCamelCase__     ,				UpperCamelCase__       )
    for attn, resnet in zip(self.attentions     ,				self.resnets[1:]       ):
      __magic_name__								=			attn(UpperCamelCase__     ,				UpperCamelCase__     ,				deterministic=UpperCamelCase__       )
      __magic_name__								=			resnet(UpperCamelCase__     ,				UpperCamelCase__     ,				deterministic=UpperCamelCase__       )

    return hidden_states

