


"""simple docstring"""



import os
import unittest

from transformers import MobileBertTokenizer, MobileBertTokenizerFast
from transformers.models.bert.tokenization_bert import (
    VOCAB_FILES_NAMES,
    BasicTokenizer,
    WordpieceTokenizer,
    _is_control,
    _is_punctuation,
    _is_whitespace,
)
from transformers.testing_utils import require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin, filter_non_english







@require_tokenizers
class _lowerCAmelCase       (   _lowerCamelCase					,							unittest.TestCase			):
   __UpperCAmelCase				: List[Any]													=						MobileBertTokenizer
   __UpperCAmelCase				: str													=						MobileBertTokenizerFast
   __UpperCAmelCase				: List[str]													=						True
   __UpperCAmelCase				: Dict													=						True
   __UpperCAmelCase				: Optional[int]													=						filter_non_english
   __UpperCAmelCase				: List[Any]													=						'''google/mobilebert-uncased'''







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       super().setUp()

       snake_case					:   str         =					[
           "[UNK]",
           "[CLS]",
           "[SEP]",
           "[PAD]",
           "[MASK]",
           "want",
           "##want",
           "##ed",
           "wa",
           "un",
           "runn",
           "##ing",
           ",",
           "low",
           "lowest",
       ]
       snake_case					:   List[Any]         =					os.path.join(self.tmpdirname		,			VOCAB_FILES_NAMES["vocab_file"]		)
       with open(self.vocab_file		,			"w"		,			encoding="utf-8"		) as vocab_writer:
           vocab_writer.write("".join([x + "\n" for x in vocab_tokens]		)		)

       snake_case					:   Dict         =					[
           (tokenizer_def[0], self.pre_trained_model_path, tokenizer_def[2])  # else the 'google/' prefix is stripped
           for tokenizer_def in self.tokenizers_list
       ]







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

       '''simple docstring'''




       snake_case					:   Union[str, Any]         =					"UNwant\u00E9d,running"
       snake_case					:   List[Any]         =					"unwanted, running"
       return input_text, output_text







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   Tuple         =					self.tokenizer_class(self.vocab_file		)

       snake_case					:   Tuple         =					tokenizer.tokenize("UNwant\u00E9d,running"		)
       self.assertListEqual(A__		,			["un", "##want", "##ed", ",", "runn", "##ing"]		)
       self.assertListEqual(tokenizer.convert_tokens_to_ids(A__		)		,			[9, 6, 7, 12, 10, 11]		)







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

       '''simple docstring'''




       if not self.test_rust_tokenizer:
           return

       snake_case					:   List[Any]         =					self.get_tokenizer()
       snake_case					:   Optional[Any]         =					self.get_rust_tokenizer()

       snake_case					:   Optional[int]         =					"UNwant\u00E9d,running"

       snake_case					:   Any         =					tokenizer.tokenize(A__		)
       snake_case					:   Any         =					rust_tokenizer.tokenize(A__		)
       self.assertListEqual(A__		,			A__		)

       snake_case					:   str         =					tokenizer.encode(A__		,			add_special_tokens=A__		)
       snake_case					:   List[Any]         =					rust_tokenizer.encode(A__		,			add_special_tokens=A__		)
       self.assertListEqual(A__		,			A__		)

       snake_case					:   Union[str, Any]         =					self.get_rust_tokenizer()
       snake_case					:   List[str]         =					tokenizer.encode(A__		)
       snake_case					:   Optional[Any]         =					rust_tokenizer.encode(A__		)
       self.assertListEqual(A__		,			A__		)

       # With lower casing
       snake_case					:   Union[str, Any]         =					self.get_tokenizer(do_lower_case=A__		)
       snake_case					:   List[str]         =					self.get_rust_tokenizer(do_lower_case=A__		)

       snake_case					:   str         =					"UNwant\u00E9d,running"

       snake_case					:   Union[str, Any]         =					tokenizer.tokenize(A__		)
       snake_case					:   Optional[int]         =					rust_tokenizer.tokenize(A__		)
       self.assertListEqual(A__		,			A__		)

       snake_case					:   Optional[Any]         =					tokenizer.encode(A__		,			add_special_tokens=A__		)
       snake_case					:   str         =					rust_tokenizer.encode(A__		,			add_special_tokens=A__		)
       self.assertListEqual(A__		,			A__		)

       snake_case					:   List[Any]         =					self.get_rust_tokenizer()
       snake_case					:   Dict         =					tokenizer.encode(A__		)
       snake_case					:   str         =					rust_tokenizer.encode(A__		)
       self.assertListEqual(A__		,			A__		)







   def 				lowerCamelCase	(			self		)			->	Tuple:

       '''simple docstring'''




       snake_case					:   Optional[int]         =					BasicTokenizer()

       self.assertListEqual(tokenizer.tokenize("ah\u535A\u63A8zz"		)		,			["ah", "\u535A", "\u63A8", "zz"]		)







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

       '''simple docstring'''




       snake_case					:   str         =					BasicTokenizer(do_lower_case=A__		)

       self.assertListEqual(
           tokenizer.tokenize(" \tHeLLo!how  \n Are yoU?  "		)		,			["hello", "!", "how", "are", "you", "?"]		)
       self.assertListEqual(tokenizer.tokenize("H\u00E9llo"		)		,			["hello"]		)







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

       '''simple docstring'''




       snake_case					:   List[Any]         =					BasicTokenizer(do_lower_case=A__		,			strip_accents=A__		)

       self.assertListEqual(
           tokenizer.tokenize(" \tHäLLo!how  \n Are yoU?  "		)		,			["hällo", "!", "how", "are", "you", "?"]		)
       self.assertListEqual(tokenizer.tokenize("H\u00E9llo"		)		,			["h\u00E9llo"]		)







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   Union[str, Any]         =					BasicTokenizer(do_lower_case=A__		,			strip_accents=A__		)

       self.assertListEqual(
           tokenizer.tokenize(" \tHäLLo!how  \n Are yoU?  "		)		,			["hallo", "!", "how", "are", "you", "?"]		)
       self.assertListEqual(tokenizer.tokenize("H\u00E9llo"		)		,			["hello"]		)







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   List[Any]         =					BasicTokenizer(do_lower_case=A__		)

       self.assertListEqual(
           tokenizer.tokenize(" \tHäLLo!how  \n Are yoU?  "		)		,			["hallo", "!", "how", "are", "you", "?"]		)
       self.assertListEqual(tokenizer.tokenize("H\u00E9llo"		)		,			["hello"]		)







   def 				lowerCamelCase	(			self		)			->	str:

       '''simple docstring'''




       snake_case					:   Union[str, Any]         =					BasicTokenizer(do_lower_case=A__		)

       self.assertListEqual(
           tokenizer.tokenize(" \tHeLLo!how  \n Are yoU?  "		)		,			["HeLLo", "!", "how", "Are", "yoU", "?"]		)







   def 				lowerCamelCase	(			self		)			->	Dict:

       '''simple docstring'''




       snake_case					:   Optional[Any]         =					BasicTokenizer(do_lower_case=A__		,			strip_accents=A__		)

       self.assertListEqual(
           tokenizer.tokenize(" \tHäLLo!how  \n Are yoU?  "		)		,			["HäLLo", "!", "how", "Are", "yoU", "?"]		)







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

       '''simple docstring'''




       snake_case					:   Tuple         =					BasicTokenizer(do_lower_case=A__		,			strip_accents=A__		)

       self.assertListEqual(
           tokenizer.tokenize(" \tHäLLo!how  \n Are yoU?  "		)		,			["HaLLo", "!", "how", "Are", "yoU", "?"]		)







   def 				lowerCamelCase	(			self		)			->	str:

       '''simple docstring'''




       snake_case					:   Any         =					BasicTokenizer(do_lower_case=A__		,			never_split=["[UNK]"]		)

       self.assertListEqual(
           tokenizer.tokenize(" \tHeLLo!how  \n Are yoU? [UNK]"		)		,			["HeLLo", "!", "how", "Are", "yoU", "?", "[UNK]"]		)







   def 				lowerCamelCase	(			self		)			->	Tuple:

       '''simple docstring'''




       snake_case					:   Any         =					["[UNK]", "[CLS]", "[SEP]", "want", "##want", "##ed", "wa", "un", "runn", "##ing"]

       snake_case					:   Optional[int]         =					{}
       for i, token in enumerate(A__		):
           snake_case					:   Tuple         =					i
       snake_case					:   Any         =					WordpieceTokenizer(vocab=A__		,			unk_token="[UNK]"		)

       self.assertListEqual(tokenizer.tokenize(""		)		,			[]		)

       self.assertListEqual(tokenizer.tokenize("unwanted running"		)		,			["un", "##want", "##ed", "runn", "##ing"]		)

       self.assertListEqual(tokenizer.tokenize("unwantedX running"		)		,			["[UNK]", "runn", "##ing"]		)







   def 				lowerCamelCase	(			self		)			->	str:

       '''simple docstring'''




       self.assertTrue(_is_whitespace(" "		)		)
       self.assertTrue(_is_whitespace("\t"		)		)
       self.assertTrue(_is_whitespace("\r"		)		)
       self.assertTrue(_is_whitespace("\n"		)		)
       self.assertTrue(_is_whitespace("\u00A0"		)		)

       self.assertFalse(_is_whitespace("A"		)		)
       self.assertFalse(_is_whitespace("-"		)		)







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

       '''simple docstring'''




       self.assertTrue(_is_control("\u0005"		)		)

       self.assertFalse(_is_control("A"		)		)
       self.assertFalse(_is_control(" "		)		)
       self.assertFalse(_is_control("\t"		)		)
       self.assertFalse(_is_control("\r"		)		)







   def 				lowerCamelCase	(			self		)			->	Dict:

       '''simple docstring'''




       self.assertTrue(_is_punctuation("-"		)		)
       self.assertTrue(_is_punctuation("$"		)		)
       self.assertTrue(_is_punctuation("`"		)		)
       self.assertTrue(_is_punctuation("."		)		)

       self.assertFalse(_is_punctuation("A"		)		)
       self.assertFalse(_is_punctuation(" "		)		)







   def 				lowerCamelCase	(			self		)			->	str:

       '''simple docstring'''




       snake_case					:   str         =					self.get_tokenizer()
       snake_case					:   List[Any]         =					self.get_rust_tokenizer()

       # Example taken from the issue https://github.com/huggingface/tokenizers/issues/340
       self.assertListEqual([tokenizer.tokenize(A__		) for t in ["Test", "\xad", "test"]]		,			[["[UNK]"], [], ["[UNK]"]]		)

       self.assertListEqual(
           [rust_tokenizer.tokenize(A__		) for t in ["Test", "\xad", "test"]]		,			[["[UNK]"], [], ["[UNK]"]]		)







   @slow
   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   Union[str, Any]         =					self.tokenizer_class.from_pretrained("google/mobilebert-uncased"		)

       snake_case					:   Any         =					tokenizer.encode("sequence builders"		,			add_special_tokens=A__		)
       snake_case					:   List[str]         =					tokenizer.encode("multi-sequence build"		,			add_special_tokens=A__		)

       snake_case					:   int         =					tokenizer.build_inputs_with_special_tokens(A__		)
       snake_case					:   str         =					tokenizer.build_inputs_with_special_tokens(A__		,			A__		)

       assert encoded_sentence == [101] + text + [102]
       assert encoded_pair == [101] + text + [102] + text_a + [102]







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

       '''simple docstring'''




       for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
           with self.subTest(F'{tokenizer.__class__.__name__} ({pretrained_name})'		):
               snake_case					:   str         =					self.rust_tokenizer_class.from_pretrained(A__		,			**A__		)

               snake_case					:   Union[str, Any]         =					F'A, naïve {tokenizer_r.mask_token} AllenNLP sentence.'
               snake_case					:   Dict         =					tokenizer_r.encode_plus(
                   A__		,			return_attention_mask=A__		,			return_token_type_ids=A__		,			return_offsets_mapping=A__		,			add_special_tokens=A__		,			)

               snake_case					:   Dict         =					tokenizer_r.do_lower_case if hasattr(A__		,			"do_lower_case"		) else False
               snake_case					:   List[Any]         =					(
                   [
                       ((0, 0), tokenizer_r.cls_token),
                       ((0, 1), "A"),
                       ((1, 2), ","),
                       ((3, 5), "na"),
                       ((5, 6), "##ï"),
                       ((6, 8), "##ve"),
                       ((9, 15), tokenizer_r.mask_token),
                       ((16, 21), "Allen"),
                       ((21, 23), "##NL"),
                       ((23, 24), "##P"),
                       ((25, 33), "sentence"),
                       ((33, 34), "."),
                       ((0, 0), tokenizer_r.sep_token),
                   ]
                   if not do_lower_case
                   else [
                       ((0, 0), tokenizer_r.cls_token),
                       ((0, 1), "a"),
                       ((1, 2), ","),
                       ((3, 8), "naive"),
                       ((9, 15), tokenizer_r.mask_token),
                       ((16, 21), "allen"),
                       ((21, 23), "##nl"),
                       ((23, 24), "##p"),
                       ((25, 33), "sentence"),
                       ((33, 34), "."),
                       ((0, 0), tokenizer_r.sep_token),
                   ]
               )

               self.assertEqual(
                   [e[1] for e in expected_results]		,			tokenizer_r.convert_ids_to_tokens(tokens["input_ids"]		)		)
               self.assertEqual([e[0] for e in expected_results]		,			tokens["offset_mapping"]		)







   def 				lowerCamelCase	(			self		)			->	Any:

       '''simple docstring'''




       snake_case					:   Tuple         =					["的", "人", "有"]
       snake_case					:   str         =					"".join(A__		)
       for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
           with self.subTest(F'{tokenizer.__class__.__name__} ({pretrained_name})'		):
               snake_case					:   str         =					True
               snake_case					:   str         =					self.tokenizer_class.from_pretrained(A__		,			**A__		)
               snake_case					:   List[str]         =					self.rust_tokenizer_class.from_pretrained(A__		,			**A__		)

               snake_case					:   Dict         =					tokenizer_p.encode(A__		,			add_special_tokens=A__		)
               snake_case					:   Optional[Any]         =					tokenizer_r.encode(A__		,			add_special_tokens=A__		)

               snake_case					:   str         =					tokenizer_r.convert_ids_to_tokens(A__		)
               snake_case					:   Union[str, Any]         =					tokenizer_p.convert_ids_to_tokens(A__		)

               # it is expected that each Chinese character is not preceded by "##"
               self.assertListEqual(A__		,			A__		)
               self.assertListEqual(A__		,			A__		)

               snake_case					:   Any         =					False
               snake_case					:   Optional[int]         =					self.rust_tokenizer_class.from_pretrained(A__		,			**A__		)
               snake_case					:   Tuple         =					self.tokenizer_class.from_pretrained(A__		,			**A__		)

               snake_case					:   List[Any]         =					tokenizer_r.encode(A__		,			add_special_tokens=A__		)
               snake_case					:   List[str]         =					tokenizer_p.encode(A__		,			add_special_tokens=A__		)

               snake_case					:   int         =					tokenizer_r.convert_ids_to_tokens(A__		)
               snake_case					:   Dict         =					tokenizer_p.convert_ids_to_tokens(A__		)

               # it is expected that only the first Chinese character is not preceded by "##".
               snake_case					:   str         =					[
                   F'##{token}' if idx != 0 else token for idx, token in enumerate(A__		)
               ]
               self.assertListEqual(A__		,			A__		)
               self.assertListEqual(A__		,			A__		)









"""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





