
'''simple docstring'''




from __future__ import annotations

import unittest

from transformers import MobileBertConfig, is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
      import tensorflow as tf

      from transformers import (
          TF_MODEL_FOR_PRETRAINING_MAPPING,
          TFMobileBertForMaskedLM,
          TFMobileBertForMultipleChoice,
          TFMobileBertForNextSentencePrediction,
          TFMobileBertForPreTraining,
          TFMobileBertForQuestionAnswering,
          TFMobileBertForSequenceClassification,
          TFMobileBertForTokenClassification,
          TFMobileBertModel,
      )





@require_tf
class  A				(							__snake_case    ,			__snake_case    ,			unittest.TestCase						):
     __magic_name__						=		(
         (
             TFMobileBertModel,
             TFMobileBertForMaskedLM,
             TFMobileBertForNextSentencePrediction,
             TFMobileBertForPreTraining,
             TFMobileBertForQuestionAnswering,
             TFMobileBertForSequenceClassification,
             TFMobileBertForTokenClassification,
             TFMobileBertForMultipleChoice,
         )
         if is_tf_available()
         else ()
     )
     __magic_name__						=		(
         {
             '''feature-extraction''': TFMobileBertModel,
             '''fill-mask''': TFMobileBertForMaskedLM,
             '''question-answering''': TFMobileBertForQuestionAnswering,
             '''text-classification''': TFMobileBertForSequenceClassification,
             '''token-classification''': TFMobileBertForTokenClassification,
             '''zero-shot''': TFMobileBertForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     __magic_name__						=		False
     __magic_name__						=		False





     def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=False	) ->				List[str]:






         """simple docstring"""






         A       :   Tuple					=						super()._prepare_for_class(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			return_labels=SCREAMING_SNAKE_CASE	)

         if return_labels:
             if model_class in get_values(SCREAMING_SNAKE_CASE	):
                 A       :   int					=						tf.zeros(self.model_tester.batch_size						,			dtype=tf.intaa	)

         return inputs_dict





     class  A				(							__snake_case						):





          def __init__(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=13						,			SCREAMING_SNAKE_CASE=7						,			SCREAMING_SNAKE_CASE=True						,			SCREAMING_SNAKE_CASE=True						,			SCREAMING_SNAKE_CASE=True						,			SCREAMING_SNAKE_CASE=True						,			SCREAMING_SNAKE_CASE=99						,			SCREAMING_SNAKE_CASE=32						,			SCREAMING_SNAKE_CASE=32						,			SCREAMING_SNAKE_CASE=2						,			SCREAMING_SNAKE_CASE=4						,			SCREAMING_SNAKE_CASE=37						,			SCREAMING_SNAKE_CASE="gelu"						,			SCREAMING_SNAKE_CASE=0.1						,			SCREAMING_SNAKE_CASE=0.1						,			SCREAMING_SNAKE_CASE=512						,			SCREAMING_SNAKE_CASE=16						,			SCREAMING_SNAKE_CASE=2						,			SCREAMING_SNAKE_CASE=0.02						,			SCREAMING_SNAKE_CASE=3						,			SCREAMING_SNAKE_CASE=4						,			SCREAMING_SNAKE_CASE=None						,			) ->				str:






              """simple docstring"""






              A       :   List[str]					=						parent
              A       :   str					=						batch_size
              A       :   Optional[Any]					=						seq_length
              A       :   List[str]					=						is_training
              A       :   List[Any]					=						use_input_mask
              A       :   Optional[Any]					=						use_token_type_ids
              A       :   Optional[Any]					=						use_labels
              A       :   List[str]					=						vocab_size
              A       :   Dict					=						hidden_size
              A       :   Union[str, Any]					=						num_hidden_layers
              A       :   Tuple					=						num_attention_heads
              A       :   Dict					=						intermediate_size
              A       :   Tuple					=						hidden_act
              A       :   List[Any]					=						hidden_dropout_prob
              A       :   Tuple					=						attention_probs_dropout_prob
              A       :   int					=						max_position_embeddings
              A       :   int					=						type_vocab_size
              A       :   str					=						type_sequence_label_size
              A       :   int					=						initializer_range
              A       :   Optional[Any]					=						num_labels
              A       :   Optional[int]					=						num_choices
              A       :   Tuple					=						scope
              A       :   Dict					=						embedding_size





          def 				__lowerCAmelCase							(							self	) ->				List[str]:






              """simple docstring"""






              A       :   List[Any]					=						ids_tensor([self.batch_size, self.seq_length]						,			self.vocab_size	)

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

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

              A       :   Tuple					=						None
              A       :   str					=						None
              A       :   Any					=						None
              if self.use_labels:
                  A       :   int					=						ids_tensor([self.batch_size]						,			self.type_sequence_label_size	)
                  A       :   List[Any]					=						ids_tensor([self.batch_size, self.seq_length]						,			self.num_labels	)
                  A       :   Union[str, Any]					=						ids_tensor([self.batch_size]						,			self.num_choices	)

              A       :   Optional[Any]					=						MobileBertConfig(
                  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						,			initializer_range=self.initializer_range						,			embedding_size=self.embedding_size						,			)

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





          def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Optional[Any]:






              """simple docstring"""






              A       :   int					=						TFMobileBertModel(config=SCREAMING_SNAKE_CASE	)
              A       :   Union[str, Any]					=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
              A       :   List[str]					=						model(SCREAMING_SNAKE_CASE	)

              A       :   str					=						[input_ids, input_mask]
              A       :   List[str]					=						model(SCREAMING_SNAKE_CASE	)

              A       :   Dict					=						model(SCREAMING_SNAKE_CASE	)

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





          def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Optional[int]:






              """simple docstring"""






              A       :   Optional[int]					=						TFMobileBertForMaskedLM(config=SCREAMING_SNAKE_CASE	)
              A       :   Any					=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
              A       :   Dict					=						model(SCREAMING_SNAKE_CASE	)
              self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.seq_length, self.vocab_size)	)





          def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Optional[Any]:






              """simple docstring"""






              A       :   Union[str, Any]					=						TFMobileBertForNextSentencePrediction(config=SCREAMING_SNAKE_CASE	)
              A       :   Tuple					=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
              A       :   Optional[Any]					=						model(SCREAMING_SNAKE_CASE	)
              self.parent.assertEqual(result.logits.shape						,			(self.batch_size, 2)	)





          def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Tuple:






              """simple docstring"""






              A       :   str					=						TFMobileBertForPreTraining(config=SCREAMING_SNAKE_CASE	)
              A       :   int					=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
              A       :   str					=						model(SCREAMING_SNAKE_CASE	)
              self.parent.assertEqual(
                  result.prediction_logits.shape						,			(self.batch_size, self.seq_length, self.vocab_size)	)
              self.parent.assertEqual(result.seq_relationship_logits.shape						,			(self.batch_size, 2)	)





          def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Optional[int]:






              """simple docstring"""






              A       :   str					=						self.num_labels
              A       :   Tuple					=						TFMobileBertForSequenceClassification(config=SCREAMING_SNAKE_CASE	)
              A       :   Optional[Any]					=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
              A       :   Dict					=						model(SCREAMING_SNAKE_CASE	)
              self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.num_labels)	)





          def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Dict:






              """simple docstring"""






              A       :   Dict					=						self.num_choices
              A       :   Dict					=						TFMobileBertForMultipleChoice(config=SCREAMING_SNAKE_CASE	)
              A       :   Tuple					=						tf.tile(tf.expand_dims(SCREAMING_SNAKE_CASE						,			1	)						,			(1, self.num_choices, 1)	)
              A       :   int					=						tf.tile(tf.expand_dims(SCREAMING_SNAKE_CASE						,			1	)						,			(1, self.num_choices, 1)	)
              A       :   Union[str, Any]					=						tf.tile(tf.expand_dims(SCREAMING_SNAKE_CASE						,			1	)						,			(1, self.num_choices, 1)	)
              A       :   Tuple					=						{
                  '''input_ids''': multiple_choice_inputs_ids,
                  '''attention_mask''': multiple_choice_input_mask,
                  '''token_type_ids''': multiple_choice_token_type_ids,
              }
              A       :   Optional[Any]					=						model(SCREAMING_SNAKE_CASE	)
              self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.num_choices)	)





          def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Tuple:






              """simple docstring"""






              A       :   Dict					=						self.num_labels
              A       :   Optional[Any]					=						TFMobileBertForTokenClassification(config=SCREAMING_SNAKE_CASE	)
              A       :   Any					=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
              A       :   str					=						model(SCREAMING_SNAKE_CASE	)
              self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.seq_length, self.num_labels)	)





          def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				List[str]:






              """simple docstring"""






              A       :   Union[str, Any]					=						TFMobileBertForQuestionAnswering(config=SCREAMING_SNAKE_CASE	)
              A       :   Dict					=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
              A       :   Any					=						model(SCREAMING_SNAKE_CASE	)
              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	) ->				List[str]:






              """simple docstring"""






              A       :   List[str]					=						self.prepare_config_and_inputs()
              (
                  (
                  A
              ),  (
                  A
              ),  (
                  A
              ),  (
                  A
              ),  (
                  A
              ),  (
                  A
              ),  (
                  A
              ),  
              )       :   Any					=						config_and_inputs
              A       :   List[str]					=						{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
              return config, inputs_dict





     def 				__lowerCAmelCase							(							self	) ->				Tuple:






         """simple docstring"""






         A       :   Optional[int]					=						TFMobileBertModelTest.TFMobileBertModelTester(self	)
         A       :   Union[str, Any]					=						ConfigTester(self						,			config_class=SCREAMING_SNAKE_CASE						,			hidden_size=37	)





     def 				__lowerCAmelCase							(							self	) ->				Any:






         """simple docstring"""






         self.config_tester.run_common_tests()





     def 				__lowerCAmelCase							(							self	) ->				Optional[int]:






         """simple docstring"""






         A       :   List[Any]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_model(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				Optional[int]:






         """simple docstring"""






         A       :   Optional[int]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_masked_lm(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				Dict:






         """simple docstring"""






         A       :   Tuple					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_multiple_choice(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				Dict:






         """simple docstring"""






         A       :   str					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				str:






         """simple docstring"""






         A       :   Union[str, Any]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_pretraining(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				Union[str, Any]:






         """simple docstring"""






         A       :   List[Any]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_question_answering(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				Optional[Any]:






         """simple docstring"""






         A       :   Optional[int]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_sequence_classification(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				List[str]:






         """simple docstring"""






         A       :   Union[str, Any]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_token_classification(*SCREAMING_SNAKE_CASE	)





     @slow
     def 				__lowerCAmelCase							(							self	) ->				Any:






         """simple docstring"""






         for model_name in ["google/mobilebert-uncased"]:
             A       :   Optional[int]					=						TFMobileBertModel.from_pretrained(SCREAMING_SNAKE_CASE	)
             self.assertIsNotNone(SCREAMING_SNAKE_CASE	)





@require_tf
class  A				(							unittest.TestCase						):





     @slow
     def 				__lowerCAmelCase							(							self	) ->				Dict:






         """simple docstring"""






         A       :   Optional[Any]					=						TFMobileBertForPreTraining.from_pretrained('''google/mobilebert-uncased'''	)
         A       :   Tuple					=						tf.constant([[0, 1, 2, 3, 4, 5]]	)
         A       :   str					=						model(SCREAMING_SNAKE_CASE	)[0]

         A       :   Dict					=						[1, 6, 30522]
         self.assertEqual(output.shape						,			SCREAMING_SNAKE_CASE	)

         A       :   Union[str, Any]					=						tf.constant(
             [
                 [
                     [-4.5_919_547, -9.248_295, -9.645_256],
                     [-6.7_306_175, -6.440_284, -6.6_052_837],
                     [-7.2_743_506, -6.7_847_915, -6.024_673],
                 ]
             ]	)
         tf.debugging.assert_near(output[:, :3, :3]						,			SCREAMING_SNAKE_CASE						,			atol=1e-4	)






'''simple docstring'''




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

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





class  A				(							nn.Module						):
     __magic_name__						=		42
     __magic_name__						=		42
     __magic_name__						=		0.0
     __magic_name__						=		1
     __magic_name__						=		1
     __magic_name__						=		True
     __magic_name__						=		False
     __magic_name__						=		False
     __magic_name__						=		False
     __magic_name__						=		jnp.floataa





     def 				__lowerCAmelCase							(							self	) ->				Tuple:






         """simple docstring"""






         A       :   Union[str, Any]					=						[]
         A       :   Union[str, Any]					=						[]

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

             A       :   Optional[Any]					=						FlaxResnetBlockaD(
                 in_channels=SCREAMING_SNAKE_CASE						,			out_channels=self.out_channels						,			dropout_prob=self.dropout						,			dtype=self.dtype						,			)
             resnets.append(SCREAMING_SNAKE_CASE	)

             A       :   Optional[int]					=						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(SCREAMING_SNAKE_CASE	)

         A       :   Union[str, Any]					=						resnets
         A       :   Union[str, Any]					=						attentions

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





     def __call__(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=True	) ->				Union[str, Any]:






         """simple docstring"""






         A       :   Optional[Any]					=						()

         for resnet, attn in zip(self.resnets						,			self.attentions	):
             A       :   int					=						resnet(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			deterministic=SCREAMING_SNAKE_CASE	)
             A       :   Dict					=						attn(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			deterministic=SCREAMING_SNAKE_CASE	)
             output_states += (hidden_states,)

         if self.add_downsample:
             A       :   Optional[Any]					=						self.downsamplers_a(SCREAMING_SNAKE_CASE	)
             output_states += (hidden_states,)

         return hidden_states, output_states





class  A				(							nn.Module						):
     __magic_name__						=		42
     __magic_name__						=		42
     __magic_name__						=		0.0
     __magic_name__						=		1
     __magic_name__						=		True
     __magic_name__						=		jnp.floataa





     def 				__lowerCAmelCase							(							self	) ->				Optional[Any]:






         """simple docstring"""






         A       :   Optional[Any]					=						[]

         for i in range(self.num_layers	):
             A       :   Optional[Any]					=						self.in_channels if i == 0 else self.out_channels

             A       :   List[str]					=						FlaxResnetBlockaD(
                 in_channels=SCREAMING_SNAKE_CASE						,			out_channels=self.out_channels						,			dropout_prob=self.dropout						,			dtype=self.dtype						,			)
             resnets.append(SCREAMING_SNAKE_CASE	)
         A       :   Dict					=						resnets

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





     def __call__(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=True	) ->				Optional[Any]:






         """simple docstring"""






         A       :   str					=						()

         for resnet in self.resnets:
             A       :   Optional[int]					=						resnet(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			deterministic=SCREAMING_SNAKE_CASE	)
             output_states += (hidden_states,)

         if self.add_downsample:
             A       :   Optional[int]					=						self.downsamplers_a(SCREAMING_SNAKE_CASE	)
             output_states += (hidden_states,)

         return hidden_states, output_states





class  A				(							nn.Module						):
     __magic_name__						=		42
     __magic_name__						=		42
     __magic_name__						=		42
     __magic_name__						=		0.0
     __magic_name__						=		1
     __magic_name__						=		1
     __magic_name__						=		True
     __magic_name__						=		False
     __magic_name__						=		False
     __magic_name__						=		False
     __magic_name__						=		jnp.floataa





     def 				__lowerCAmelCase							(							self	) ->				Tuple:






         """simple docstring"""






         A       :   Optional[Any]					=						[]
         A       :   Optional[int]					=						[]

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

             A       :   List[str]					=						FlaxResnetBlockaD(
                 in_channels=resnet_in_channels + res_skip_channels						,			out_channels=self.out_channels						,			dropout_prob=self.dropout						,			dtype=self.dtype						,			)
             resnets.append(SCREAMING_SNAKE_CASE	)

             A       :   int					=						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(SCREAMING_SNAKE_CASE	)

         A       :   Dict					=						resnets
         A       :   Optional[Any]					=						attentions

         if self.add_upsample:
             A       :   Optional[int]					=						FlaxUpsampleaD(self.out_channels						,			dtype=self.dtype	)





     def __call__(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=True	) ->				Optional[int]:






         """simple docstring"""






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

             A       :   Union[str, Any]					=						resnet(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			deterministic=SCREAMING_SNAKE_CASE	)
             A       :   Tuple					=						attn(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			deterministic=SCREAMING_SNAKE_CASE	)

         if self.add_upsample:
             A       :   Dict					=						self.upsamplers_a(SCREAMING_SNAKE_CASE	)

         return hidden_states





class  A				(							nn.Module						):
     __magic_name__						=		42
     __magic_name__						=		42
     __magic_name__						=		42
     __magic_name__						=		0.0
     __magic_name__						=		1
     __magic_name__						=		True
     __magic_name__						=		jnp.floataa





     def 				__lowerCAmelCase							(							self	) ->				Dict:






         """simple docstring"""






         A       :   int					=						[]

         for i in range(self.num_layers	):
             A       :   List[Any]					=						self.in_channels if (i == self.num_layers - 1) else self.out_channels
             A       :   List[str]					=						self.prev_output_channel if i == 0 else self.out_channels

             A       :   str					=						FlaxResnetBlockaD(
                 in_channels=resnet_in_channels + res_skip_channels						,			out_channels=self.out_channels						,			dropout_prob=self.dropout						,			dtype=self.dtype						,			)
             resnets.append(SCREAMING_SNAKE_CASE	)

         A       :   List[Any]					=						resnets

         if self.add_upsample:
             A       :   Optional[Any]					=						FlaxUpsampleaD(self.out_channels						,			dtype=self.dtype	)





     def __call__(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=True	) ->				Tuple:






         """simple docstring"""






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

             A       :   Optional[Any]					=						resnet(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			deterministic=SCREAMING_SNAKE_CASE	)

         if self.add_upsample:
             A       :   List[str]					=						self.upsamplers_a(SCREAMING_SNAKE_CASE	)

         return hidden_states





class  A				(							nn.Module						):
     __magic_name__						=		42
     __magic_name__						=		0.0
     __magic_name__						=		1
     __magic_name__						=		1
     __magic_name__						=		False
     __magic_name__						=		False
     __magic_name__						=		jnp.floataa





     def 				__lowerCAmelCase							(							self	) ->				Optional[int]:






         """simple docstring"""






         A       :   str					=						[
             FlaxResnetBlockaD(
                 in_channels=self.in_channels						,			out_channels=self.in_channels						,			dropout_prob=self.dropout						,			dtype=self.dtype						,			)
         ]

         A       :   List[Any]					=						[]

         for _ in range(self.num_layers	):
             A       :   int					=						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(SCREAMING_SNAKE_CASE	)

             A       :   Union[str, Any]					=						FlaxResnetBlockaD(
                 in_channels=self.in_channels						,			out_channels=self.in_channels						,			dropout_prob=self.dropout						,			dtype=self.dtype						,			)
             resnets.append(SCREAMING_SNAKE_CASE	)

         A       :   List[str]					=						resnets
         A       :   List[str]					=						attentions





     def __call__(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=True	) ->				Dict:






         """simple docstring"""






         A       :   Optional[Any]					=						self.resnets[0](SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	)
         for attn, resnet in zip(self.attentions						,			self.resnets[1:]	):
             A       :   Optional[int]					=						attn(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			deterministic=SCREAMING_SNAKE_CASE	)
             A       :   Union[str, Any]					=						resnet(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			deterministic=SCREAMING_SNAKE_CASE	)

         return hidden_states




