






import unittest

from transformers import DebertaVaConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import (
     DebertaVaForMaskedLM,
     DebertaVaForMultipleChoice,
     DebertaVaForQuestionAnswering,
     DebertaVaForSequenceClassification,
     DebertaVaForTokenClassification,
     DebertaVaModel,
 )
 from transformers.models.deberta_va.modeling_deberta_va import DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST


class      lowercase__     (			_UpperCAmelCase	):

 def __init__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase=13		,    __UpperCAmelCase=7		,    __UpperCAmelCase=True		,    __UpperCAmelCase=True		,    __UpperCAmelCase=True		,    __UpperCAmelCase=True		,    __UpperCAmelCase=99		,    __UpperCAmelCase=32		,    __UpperCAmelCase=5		,    __UpperCAmelCase=4		,    __UpperCAmelCase=37		,    __UpperCAmelCase="gelu"		,    __UpperCAmelCase=0.1		,    __UpperCAmelCase=0.1		,    __UpperCAmelCase=512		,    __UpperCAmelCase=16		,    __UpperCAmelCase=2		,    __UpperCAmelCase=0.02		,    __UpperCAmelCase=False		,    __UpperCAmelCase=True		,    __UpperCAmelCase="None"		,    __UpperCAmelCase=3		,    __UpperCAmelCase=4		,    __UpperCAmelCase=None		,    )->					List[str]:







       '''simple docstring'''





       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__       =						relative_attention
       lowerCAmelCase__       =						position_biased_input
       lowerCAmelCase__       =						pos_att_type
       lowerCAmelCase__       =						scope

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







       '''simple docstring'''





       lowerCAmelCase__       =						ids_tensor([self.batch_size, self.seq_length]		,    self.vocab_size       )

       lowerCAmelCase__       =						None
       if self.use_input_mask:
             lowerCAmelCase__       =						ids_tensor([self.batch_size, self.seq_length]		,    vocab_size=2       )

       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 		UpperCAmelCase  (							self       )->					List[Any]:







       '''simple docstring'''





       return DebertaVaConfig(
           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		,    relative_attention=self.relative_attention		,    position_biased_input=self.position_biased_input		,    pos_att_type=self.pos_att_type		,    )

 def 		UpperCAmelCase  (							self		,    __UpperCAmelCase       )->					int:







       '''simple docstring'''





       self.parent.assertListEqual(list(result.loss.size()       )		,    []       )

 def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )->					Tuple:







       '''simple docstring'''





       lowerCAmelCase__       =						DebertaVaModel(config=UpperCamelCase__       )
       model.to(UpperCamelCase__       )
       model.eval()
       lowerCAmelCase__       =						model(UpperCamelCase__		,    attention_mask=UpperCamelCase__		,    token_type_ids=UpperCamelCase__       )[0]
       lowerCAmelCase__       =						model(UpperCamelCase__		,    token_type_ids=UpperCamelCase__       )[0]
       lowerCAmelCase__       =						model(UpperCamelCase__       )[0]

       self.parent.assertListEqual(list(sequence_output.size()       )		,    [self.batch_size, self.seq_length, self.hidden_size]       )

 def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )->					str:







       '''simple docstring'''





       lowerCAmelCase__       =						DebertaVaForMaskedLM(config=UpperCamelCase__       )
       model.to(UpperCamelCase__       )
       model.eval()
       lowerCAmelCase__       =						model(UpperCamelCase__		,    attention_mask=UpperCamelCase__		,    token_type_ids=UpperCamelCase__		,    labels=UpperCamelCase__       )

       self.parent.assertEqual(result.logits.shape		,    (self.batch_size, self.seq_length, self.vocab_size)       )

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







       '''simple docstring'''





       lowerCAmelCase__       =						self.num_labels
       lowerCAmelCase__       =						DebertaVaForSequenceClassification(UpperCamelCase__       )
       model.to(UpperCamelCase__       )
       model.eval()
       lowerCAmelCase__       =						model(UpperCamelCase__		,    attention_mask=UpperCamelCase__		,    token_type_ids=UpperCamelCase__		,    labels=UpperCamelCase__       )
       self.parent.assertListEqual(list(result.logits.size()       )		,    [self.batch_size, self.num_labels]       )
       self.check_loss_output(UpperCamelCase__       )

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







       '''simple docstring'''





       lowerCAmelCase__       =						self.num_labels
       lowerCAmelCase__       =						DebertaVaForTokenClassification(config=UpperCamelCase__       )
       model.to(UpperCamelCase__       )
       model.eval()
       lowerCAmelCase__       =						model(UpperCamelCase__		,    attention_mask=UpperCamelCase__		,    token_type_ids=UpperCamelCase__		,    labels=UpperCamelCase__       )
       self.parent.assertEqual(result.logits.shape		,    (self.batch_size, self.seq_length, self.num_labels)       )

 def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )->					Optional[int]:







       '''simple docstring'''





       lowerCAmelCase__       =						DebertaVaForQuestionAnswering(config=UpperCamelCase__       )
       model.to(UpperCamelCase__       )
       model.eval()
       lowerCAmelCase__       =						model(
           UpperCamelCase__		,    attention_mask=UpperCamelCase__		,    token_type_ids=UpperCamelCase__		,    start_positions=UpperCamelCase__		,    end_positions=UpperCamelCase__		,    )
       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 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )->					str:







       '''simple docstring'''





       lowerCAmelCase__       =						DebertaVaForMultipleChoice(config=UpperCamelCase__       )
       model.to(UpperCamelCase__       )
       model.eval()
       lowerCAmelCase__       =						input_ids.unsqueeze(1       ).expand(-1		,    self.num_choices		,    -1       ).contiguous()
       lowerCAmelCase__       =						token_type_ids.unsqueeze(1       ).expand(-1		,    self.num_choices		,    -1       ).contiguous()
       lowerCAmelCase__       =						input_mask.unsqueeze(1       ).expand(-1		,    self.num_choices		,    -1       ).contiguous()
       lowerCAmelCase__       =						model(
           UpperCamelCase__		,    attention_mask=UpperCamelCase__		,    token_type_ids=UpperCamelCase__		,    labels=UpperCamelCase__		,    )
       self.parent.assertEqual(result.logits.shape		,    (self.batch_size, self.num_choices)       )


 def 		UpperCAmelCase  (							self       )->					int:







       '''simple docstring'''





       lowerCAmelCase__       =						self.prepare_config_and_inputs()
       (
           (
           lowerCAmelCase__       
       )      ,							(
           lowerCAmelCase__       
       )      ,							(
           lowerCAmelCase__       
       )      ,							(
           lowerCAmelCase__       
       )      ,							(
           lowerCAmelCase__       
       )      ,							(
           lowerCAmelCase__       
       )      ,							(
           lowerCAmelCase__       
       )      ,							       
       )       =						config_and_inputs
       lowerCAmelCase__       =						{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
       return config, inputs_dict





@require_torch
class      lowercase__     (			_UpperCAmelCase,					_UpperCAmelCase,					unittest.TestCase	):
 a_         =(
     (
         DebertaVaModel,
         DebertaVaForMaskedLM,
         DebertaVaForSequenceClassification,
         DebertaVaForTokenClassification,
         DebertaVaForQuestionAnswering,
         DebertaVaForMultipleChoice,
     )
     if is_torch_available()
     else ()
 )
 a_         =(
     {
         "feature-extraction": DebertaVaModel,
         "fill-mask": DebertaVaForMaskedLM,
         "question-answering": DebertaVaForQuestionAnswering,
         "text-classification": DebertaVaForSequenceClassification,
         "token-classification": DebertaVaForTokenClassification,
         "zero-shot": DebertaVaForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )

 a_         =True
 a_         =False
 a_         =False
 a_         =False
 a_         =False

 def 		UpperCAmelCase  (							self       )->					Any:







       '''simple docstring'''





       lowerCAmelCase__       =						DebertaVaModelTester(self       )
       lowerCAmelCase__       =						ConfigTester(self		,    config_class=UpperCamelCase__		,    hidden_size=37       )

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







       '''simple docstring'''





       self.config_tester.run_common_tests()

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







       '''simple docstring'''





       lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_deberta_model(*UpperCamelCase__       )

 def 		UpperCAmelCase  (							self       )->					str:







       '''simple docstring'''





       lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_deberta_for_sequence_classification(*UpperCamelCase__       )

 def 		UpperCAmelCase  (							self       )->					int:







       '''simple docstring'''





       lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_deberta_for_masked_lm(*UpperCamelCase__       )

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







       '''simple docstring'''





       lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_deberta_for_question_answering(*UpperCamelCase__       )

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







       '''simple docstring'''





       lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_deberta_for_token_classification(*UpperCamelCase__       )

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







       '''simple docstring'''





       lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_deberta_for_multiple_choice(*UpperCamelCase__       )


 @slow
 def 		UpperCAmelCase  (							self       )->					Optional[Any]:







       '''simple docstring'''





       for model_name in DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             lowerCAmelCase__       =						DebertaVaModel.from_pretrained(UpperCamelCase__       )
             self.assertIsNotNone(UpperCamelCase__       )





@require_torch
@require_sentencepiece
@require_tokenizers
class      lowercase__     (			unittest.TestCase	):

 @unittest.skip(reason="Model not available yet"       )
 def 		UpperCAmelCase  (							self       )->					Optional[int]:







       '''simple docstring'''





       pass


 @slow
 def 		UpperCAmelCase  (							self       )->					Any:







       '''simple docstring'''





       lowerCAmelCase__       =						DebertaVaModel.from_pretrained("microsoft/deberta-v2-xlarge"       )

       lowerCAmelCase__       =						torch.tensor([[0, 31414, 232, 328, 740, 1140, 12695, 69, 46078, 1588, 2]]       )
       lowerCAmelCase__       =						torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]       )
       with torch.no_grad():
             lowerCAmelCase__       =						model(UpperCamelCase__		,    attention_mask=UpperCamelCase__       )[0]
       # compare the actual values for a slice.
       lowerCAmelCase__       =						torch.tensor(
           [[[0.2_356, 0.1_948, 0.0_369], [-0.1_063, 0.3_586, -0.5_152], [-0.6_399, -0.0_259, -0.2_525]]]       )
       self.assertTrue(torch.allclose(output[:, 1:4, 1:4]		,    UpperCamelCase__		,    atol=1E-4       )		,    F"{output[:, 1:4, 1:4]}"       )














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

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


class      lowercase__     (			nn.Module	):
	a_         =42
	a_         =42
	a_         =0.0
	a_         =1
	a_         =1
	a_         =True
	a_         =False
	a_         =False
	a_         =False
	a_         =jnp.floataa

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







							'''simple docstring'''





							lowerCAmelCase__       =						[]
							lowerCAmelCase__       =						[]

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

													lowerCAmelCase__       =						FlaxResnetBlockaD(
													    in_channels=__UpperCAmelCase		,    out_channels=self.out_channels		,    dropout_prob=self.dropout		,    dtype=self.dtype		,    )
													resnets.append(__UpperCAmelCase       )

													lowerCAmelCase__       =						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       )

							lowerCAmelCase__       =						resnets
							lowerCAmelCase__       =						attentions

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


	def __call__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase=True       )->					str:







							'''simple docstring'''





							lowerCAmelCase__       =						()

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

							if self.add_downsample:
													lowerCAmelCase__       =						self.downsamplers_a(__UpperCAmelCase       )
													output_states += (hidden_states,)

							return hidden_states, output_states





class      lowercase__     (			nn.Module	):
	a_         =42
	a_         =42
	a_         =0.0
	a_         =1
	a_         =True
	a_         =jnp.floataa

	def 		UpperCAmelCase  (							self       )->					Dict:







							'''simple docstring'''





							lowerCAmelCase__       =						[]

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

													lowerCAmelCase__       =						FlaxResnetBlockaD(
													    in_channels=__UpperCAmelCase		,    out_channels=self.out_channels		,    dropout_prob=self.dropout		,    dtype=self.dtype		,    )
													resnets.append(__UpperCAmelCase       )
							lowerCAmelCase__       =						resnets

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


	def __call__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase=True       )->					int:







							'''simple docstring'''





							lowerCAmelCase__       =						()

							for resnet in self.resnets:
													lowerCAmelCase__       =						resnet(__UpperCAmelCase		,    __UpperCAmelCase		,    deterministic=__UpperCAmelCase       )
													output_states += (hidden_states,)

							if self.add_downsample:
													lowerCAmelCase__       =						self.downsamplers_a(__UpperCAmelCase       )
													output_states += (hidden_states,)

							return hidden_states, output_states





class      lowercase__     (			nn.Module	):
	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 		UpperCAmelCase  (							self       )->					int:







							'''simple docstring'''





							lowerCAmelCase__       =						[]
							lowerCAmelCase__       =						[]

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

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

													lowerCAmelCase__       =						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       )

							lowerCAmelCase__       =						resnets
							lowerCAmelCase__       =						attentions

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


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







							'''simple docstring'''





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

													lowerCAmelCase__       =						resnet(__UpperCAmelCase		,    __UpperCAmelCase		,    deterministic=__UpperCAmelCase       )
													lowerCAmelCase__       =						attn(__UpperCAmelCase		,    __UpperCAmelCase		,    deterministic=__UpperCAmelCase       )

							if self.add_upsample:
													lowerCAmelCase__       =						self.upsamplers_a(__UpperCAmelCase       )

							return hidden_states





class      lowercase__     (			nn.Module	):
	a_         =42
	a_         =42
	a_         =42
	a_         =0.0
	a_         =1
	a_         =True
	a_         =jnp.floataa

	def 		UpperCAmelCase  (							self       )->					int:







							'''simple docstring'''





							lowerCAmelCase__       =						[]

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

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

							lowerCAmelCase__       =						resnets

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


	def __call__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase=True       )->					int:







							'''simple docstring'''





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

													lowerCAmelCase__       =						resnet(__UpperCAmelCase		,    __UpperCAmelCase		,    deterministic=__UpperCAmelCase       )

							if self.add_upsample:
													lowerCAmelCase__       =						self.upsamplers_a(__UpperCAmelCase       )

							return hidden_states





class      lowercase__     (			nn.Module	):
	a_         =42
	a_         =0.0
	a_         =1
	a_         =1
	a_         =False
	a_         =False
	a_         =jnp.floataa

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







							'''simple docstring'''





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

							lowerCAmelCase__       =						[]

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

													lowerCAmelCase__       =						FlaxResnetBlockaD(
													    in_channels=self.in_channels		,    out_channels=self.in_channels		,    dropout_prob=self.dropout		,    dtype=self.dtype		,    )
													resnets.append(__UpperCAmelCase       )

							lowerCAmelCase__       =						resnets
							lowerCAmelCase__       =						attentions


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







							'''simple docstring'''





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

							return hidden_states






