


"""simple docstring"""



import unittest

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


if is_torch_available():
  import torch

  from transformers import BertGenerationDecoder, BertGenerationEncoder







class _lowerCAmelCase       :







   def __init__(			self		,			UpperCamelCase__		,			UpperCamelCase__=13		,			UpperCamelCase__=7		,			UpperCamelCase__=True		,			UpperCamelCase__=True		,			UpperCamelCase__=99		,			UpperCamelCase__=32		,			UpperCamelCase__=5		,			UpperCamelCase__=4		,			UpperCamelCase__=37		,			UpperCamelCase__="gelu"		,			UpperCamelCase__=0.1		,			UpperCamelCase__=0.1		,			UpperCamelCase__=50		,			UpperCamelCase__=0.02		,			UpperCamelCase__=True		,			UpperCamelCase__=None		,			)			->	Optional[Any]:

       '''simple docstring'''




       snake_case					:   str         =					parent
       snake_case					:   int         =					batch_size
       snake_case					:   str         =					seq_length
       snake_case					:   Any         =					is_training
       snake_case					:   List[str]         =					use_input_mask
       snake_case					:   str         =					vocab_size
       snake_case					:   int         =					hidden_size
       snake_case					:   Optional[int]         =					num_hidden_layers
       snake_case					:   List[str]         =					num_attention_heads
       snake_case					:   List[Any]         =					intermediate_size
       snake_case					:   Dict         =					hidden_act
       snake_case					:   Any         =					hidden_dropout_prob
       snake_case					:   Optional[Any]         =					attention_probs_dropout_prob
       snake_case					:   Optional[int]         =					max_position_embeddings
       snake_case					:   int         =					initializer_range
       snake_case					:   Tuple         =					use_labels
       snake_case					:   int         =					scope







   def 				lowerCamelCase	(			self		)			->	int:

       '''simple docstring'''




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

       snake_case					:   Tuple         =					None
       if self.use_input_mask:
           snake_case					:   Optional[int]         =					random_attention_mask([self.batch_size, self.seq_length]		)

       if self.use_labels:
           snake_case					:   int         =					ids_tensor([self.batch_size, self.seq_length]		,			self.vocab_size		)

       snake_case					:   Dict         =					self.get_config()

       return config, input_ids, input_mask, token_labels







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

       '''simple docstring'''




       return BertGenerationConfig(
           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		,			is_decoder=UpperCamelCase__		,			initializer_range=self.initializer_range		,			)







   def 				lowerCamelCase	(			self		)			->	int:

       '''simple docstring'''




       (
           (
           snake_case
       )      ,(
           snake_case
       )      ,(
           snake_case
       )      ,(
           snake_case
       )      ,
       )					:   Any         =					self.prepare_config_and_inputs()

       snake_case					:   Dict         =					True
       snake_case					:   List[str]         =					floats_tensor([self.batch_size, self.seq_length, self.hidden_size]		)
       snake_case					:   Optional[Any]         =					ids_tensor([self.batch_size, self.seq_length]		,			vocab_size=2		)

       return (
           config,
           input_ids,
           input_mask,
           token_labels,
           encoder_hidden_states,
           encoder_attention_mask,
       )







   def 				lowerCamelCase	(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			**UpperCamelCase__		,			)			->	List[str]:

       '''simple docstring'''




       snake_case					:   str         =					BertGenerationEncoder(config=UpperCamelCase__		)
       model.to(UpperCamelCase__		)
       model.eval()
       snake_case					:   List[str]         =					model(UpperCamelCase__		,			attention_mask=UpperCamelCase__		)
       snake_case					:   Optional[Any]         =					model(UpperCamelCase__		)
       self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, self.seq_length, self.hidden_size)		)







   def 				lowerCamelCase	(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			**UpperCamelCase__		,			)			->	Optional[Any]:

       '''simple docstring'''




       snake_case					:   Dict         =					True
       snake_case					:   Tuple         =					BertGenerationEncoder(config=UpperCamelCase__		)
       model.to(UpperCamelCase__		)
       model.eval()
       snake_case					:   Optional[int]         =					model(
           UpperCamelCase__		,			attention_mask=UpperCamelCase__		,			encoder_hidden_states=UpperCamelCase__		,			encoder_attention_mask=UpperCamelCase__		,			)
       snake_case					:   Tuple         =					model(
           UpperCamelCase__		,			attention_mask=UpperCamelCase__		,			encoder_hidden_states=UpperCamelCase__		,			)
       self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, self.seq_length, self.hidden_size)		)







   def 				lowerCamelCase	(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			**UpperCamelCase__		,			)			->	Optional[Any]:

       '''simple docstring'''




       snake_case					:   Optional[int]         =					True
       snake_case					:   str         =					True
       snake_case					:   Optional[int]         =					BertGenerationDecoder(config=UpperCamelCase__		).to(UpperCamelCase__		).eval()

       # first forward pass
       snake_case					:   Union[str, Any]         =					model(
           UpperCamelCase__		,			attention_mask=UpperCamelCase__		,			encoder_hidden_states=UpperCamelCase__		,			encoder_attention_mask=UpperCamelCase__		,			use_cache=UpperCamelCase__		,			)
       snake_case					:   Dict         =					outputs.past_key_values

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

       # append to next input_ids and
       snake_case					:   int         =					torch.cat([input_ids, next_tokens]		,			dim=-1		)
       snake_case					:   int         =					torch.cat([input_mask, next_mask]		,			dim=-1		)

       snake_case					:   Union[str, Any]         =					model(
           UpperCamelCase__		,			attention_mask=UpperCamelCase__		,			encoder_hidden_states=UpperCamelCase__		,			encoder_attention_mask=UpperCamelCase__		,			output_hidden_states=UpperCamelCase__		,			)["hidden_states"][0]
       snake_case					:   Any         =					model(
           UpperCamelCase__		,			attention_mask=UpperCamelCase__		,			encoder_hidden_states=UpperCamelCase__		,			encoder_attention_mask=UpperCamelCase__		,			past_key_values=UpperCamelCase__		,			output_hidden_states=UpperCamelCase__		,			)["hidden_states"][0]

       # select random slice
       snake_case					:   Dict         =					ids_tensor((1,)		,			output_from_past.shape[-1]		).item()
       snake_case					:   Optional[int]         =					output_from_no_past[:, -3:, random_slice_idx].detach()
       snake_case					:   Tuple         =					output_from_past[:, :, random_slice_idx].detach()

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

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







   def 				lowerCamelCase	(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			*UpperCamelCase__		,			)			->	Tuple:

       '''simple docstring'''




       snake_case					:   int         =					BertGenerationDecoder(UpperCamelCase__		)
       model.to(UpperCamelCase__		)
       model.eval()
       snake_case					:   List[str]         =					model(UpperCamelCase__		,			attention_mask=UpperCamelCase__		,			labels=UpperCamelCase__		)
       self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.seq_length, self.vocab_size)		)







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

       '''simple docstring'''




       snake_case      ,snake_case      ,snake_case      ,snake_case					:   Union[str, Any]         =					self.prepare_config_and_inputs()
       snake_case					:   Optional[int]         =					{"input_ids": input_ids, "attention_mask": input_mask}
       return config, inputs_dict







@require_torch
class _lowerCAmelCase       (   snake_case_					,							snake_case_					,							snake_case_					,							unittest.TestCase			):
   __UpperCAmelCase				: Tuple													=						(BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
   __UpperCAmelCase				: Any													=						(BertGenerationDecoder,) if is_torch_available() else ()
   __UpperCAmelCase				: Union[str, Any]													=						(
       {'''feature-extraction''': BertGenerationEncoder, '''text-generation''': BertGenerationDecoder}
       if is_torch_available()
       else {}
   )







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   int         =					BertGenerationEncoderTester(self		)
       snake_case					:   str         =					ConfigTester(self		,			config_class=UpperCamelCase__		,			hidden_size=37		)







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       self.config_tester.run_common_tests()







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   Optional[int]         =					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*UpperCamelCase__		)







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

       '''simple docstring'''




       snake_case      ,snake_case      ,snake_case      ,snake_case					:   Dict         =					self.model_tester.prepare_config_and_inputs()
       snake_case					:   List[Any]         =					"bert"
       self.model_tester.create_and_check_model(UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		)







   def 				lowerCamelCase	(			self		)			->	List[Any]:

       '''simple docstring'''




       snake_case					:   str         =					self.model_tester.prepare_config_and_inputs_for_decoder()
       self.model_tester.create_and_check_model_as_decoder(*UpperCamelCase__		)







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

       '''simple docstring'''




       snake_case					:   str         =					self.model_tester.prepare_config_and_inputs_for_decoder()
       self.model_tester.create_and_check_decoder_model_past_large_inputs(*UpperCamelCase__		)







   def 				lowerCamelCase	(			self		)			->	Dict:

       '''simple docstring'''




       (
           (
           snake_case
       )      ,(
           snake_case
       )      ,(
           snake_case
       )      ,(
           snake_case
       )      ,(
           snake_case
       )      ,(
           snake_case
       )      ,
       )					:   Optional[Any]         =					self.model_tester.prepare_config_and_inputs_for_decoder()

       snake_case					:   str         =					None

       self.model_tester.create_and_check_model_as_decoder(
           UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			)







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   Optional[Any]         =					self.model_tester.prepare_config_and_inputs_for_decoder()
       self.model_tester.create_and_check_for_causal_lm(*UpperCamelCase__		)







   @slow
   def 				lowerCamelCase	(			self		)			->	List[str]:

       '''simple docstring'''




       snake_case					:   str         =					BertGenerationEncoder.from_pretrained("google/bert_for_seq_generation_L-24_bbc_encoder"		)
       self.assertIsNotNone(UpperCamelCase__		)







@require_torch
class _lowerCAmelCase       (   unittest.TestCase			):







   @slow
   def 				lowerCamelCase	(			self		)			->	List[Any]:

       '''simple docstring'''




       snake_case					:   List[str]         =					BertGenerationEncoder.from_pretrained("google/bert_for_seq_generation_L-24_bbc_encoder"		)
       snake_case					:   List[str]         =					torch.tensor([[101, 7592, 1010, 2026, 3899, 2003, 1_0140, 102]]		)
       with torch.no_grad():
           snake_case					:   Optional[int]         =					model(UpperCamelCase__		)[0]
       snake_case					:   Dict         =					torch.Size([1, 8, 1024]		)
       self.assertEqual(output.shape		,			UpperCamelCase__		)
       snake_case					:   Any         =					torch.tensor(
           [[[0.1775, 0.0083, -0.0321], [1.6002, 0.1287, 0.3912], [2.1473, 0.5791, 0.6066]]]		)
       self.assertTrue(torch.allclose(output[:, :3, :3]		,			UpperCamelCase__		,			atol=1e-4		)		)







@require_torch
class _lowerCAmelCase       (   unittest.TestCase			):







   @slow
   def 				lowerCamelCase	(			self		)			->	Tuple:

       '''simple docstring'''




       snake_case					:   Union[str, Any]         =					BertGenerationDecoder.from_pretrained("google/bert_for_seq_generation_L-24_bbc_encoder"		)
       snake_case					:   str         =					torch.tensor([[101, 7592, 1010, 2026, 3899, 2003, 1_0140, 102]]		)
       with torch.no_grad():
           snake_case					:   Any         =					model(UpperCamelCase__		)[0]
       snake_case					:   str         =					torch.Size([1, 8, 5_0358]		)
       self.assertEqual(output.shape		,			UpperCamelCase__		)
       snake_case					:   Dict         =					torch.tensor(
           [[[-0.5788, -2.5994, -3.7054], [0.0438, 4.7997, 1.8795], [1.5862, 6.6409, 4.4638]]]		)
       self.assertTrue(torch.allclose(output[:, :3, :3]		,			UpperCamelCase__		,			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 _lowerCAmelCase       (   nn.Module			):
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: float													=						0.0
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: bool													=						True
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: jnp.dtype													=						jnp.floataa







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

       '''simple docstring'''




       snake_case					:   Optional[int]         =					[]
       snake_case					:   Tuple         =					[]

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

           snake_case					:   Optional[int]         =					FlaxResnetBlockaD(
               in_channels=UpperCamelCase__		,			out_channels=self.out_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
           resnets.append(UpperCamelCase__		)

           snake_case					:   Dict         =					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__		)

       snake_case					:   List[Any]         =					resnets
       snake_case					:   List[Any]         =					attentions

       if self.add_downsample:
           snake_case					:   Optional[Any]         =					FlaxDownsampleaD(self.out_channels		,			dtype=self.dtype		)







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

       '''simple docstring'''




       snake_case					:   Tuple         =					()

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

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

       return hidden_states, output_states







class _lowerCAmelCase       (   nn.Module			):
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: float													=						0.0
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: bool													=						True
   __UpperCAmelCase				: jnp.dtype													=						jnp.floataa







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

       '''simple docstring'''




       snake_case					:   Dict         =					[]

       for i in range(self.num_layers		):
           snake_case					:   List[str]         =					self.in_channels if i == 0 else self.out_channels

           snake_case					:   str         =					FlaxResnetBlockaD(
               in_channels=UpperCamelCase__		,			out_channels=self.out_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
           resnets.append(UpperCamelCase__		)
       snake_case					:   str         =					resnets

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







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

       '''simple docstring'''




       snake_case					:   Tuple         =					()

       for resnet in self.resnets:
           snake_case					:   Optional[Any]         =					resnet(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)
           output_states += (hidden_states,)

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

       return hidden_states, output_states







class _lowerCAmelCase       (   nn.Module			):
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: float													=						0.0
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: bool													=						True
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: jnp.dtype													=						jnp.floataa







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

       '''simple docstring'''




       snake_case					:   Dict         =					[]
       snake_case					:   Union[str, Any]         =					[]

       for i in range(self.num_layers		):
           snake_case					:   Optional[Any]         =					self.in_channels if (i == self.num_layers - 1) else self.out_channels
           snake_case					:   Union[str, Any]         =					self.prev_output_channel if i == 0 else self.out_channels

           snake_case					:   Optional[Any]         =					FlaxResnetBlockaD(
               in_channels=resnet_in_channels + res_skip_channels		,			out_channels=self.out_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
           resnets.append(UpperCamelCase__		)

           snake_case					:   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(UpperCamelCase__		)

       snake_case					:   str         =					resnets
       snake_case					:   List[Any]         =					attentions

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







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

       '''simple docstring'''




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

           snake_case					:   List[Any]         =					resnet(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)
           snake_case					:   Any         =					attn(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)

       if self.add_upsample:
           snake_case					:   int         =					self.upsamplers_a(UpperCamelCase__		)

       return hidden_states







class _lowerCAmelCase       (   nn.Module			):
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: int
   __UpperCAmelCase				: float													=						0.0
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: bool													=						True
   __UpperCAmelCase				: jnp.dtype													=						jnp.floataa







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

       '''simple docstring'''




       snake_case					:   str         =					[]

       for i in range(self.num_layers		):
           snake_case					:   Union[str, Any]         =					self.in_channels if (i == self.num_layers - 1) else self.out_channels
           snake_case					:   Optional[Any]         =					self.prev_output_channel if i == 0 else self.out_channels

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

       snake_case					:   Optional[int]         =					resnets

       if self.add_upsample:
           snake_case					:   Dict         =					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
           snake_case					:   int         =					res_hidden_states_tuple[-1]
           snake_case					:   List[Any]         =					res_hidden_states_tuple[:-1]
           snake_case					:   Optional[int]         =					jnp.concatenate((hidden_states, res_hidden_states)		,			axis=-1		)

           snake_case					:   Dict         =					resnet(UpperCamelCase__		,			UpperCamelCase__		,			deterministic=UpperCamelCase__		)

       if self.add_upsample:
           snake_case					:   Optional[int]         =					self.upsamplers_a(UpperCamelCase__		)

       return hidden_states







class _lowerCAmelCase       (   nn.Module			):
   __UpperCAmelCase				: int
   __UpperCAmelCase				: float													=						0.0
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: int													=						1
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: bool													=						False
   __UpperCAmelCase				: jnp.dtype													=						jnp.floataa







   def 				lowerCamelCase	(			self		)			->	Tuple:

       '''simple docstring'''




       snake_case					:   Dict         =					[
           FlaxResnetBlockaD(
               in_channels=self.in_channels		,			out_channels=self.in_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
       ]

       snake_case					:   Any         =					[]

       for _ in range(self.num_layers		):
           snake_case					:   Union[str, Any]         =					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__		)

           snake_case					:   List[Any]         =					FlaxResnetBlockaD(
               in_channels=self.in_channels		,			out_channels=self.in_channels		,			dropout_prob=self.dropout		,			dtype=self.dtype		,			)
           resnets.append(UpperCamelCase__		)

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







   def __call__(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__=True		)			->	Any:

       '''simple docstring'''




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

       return hidden_states





