

import pickle
import shutil
import tempfile
import unittest

from transformers import SPIECE_UNDERLINE, XLMRobertaTokenizer, XLMRobertaTokenizerFast
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow
from transformers.utils import cached_property

from ...test_tokenization_common import TokenizerTesterMixin


UpperCAmelCase__ :     List[Any]				    =    get_tests_dir("""fixtures/test_sentencepiece.model""")







@require_sentencepiece
@require_tokenizers
class        a__       (   UpperCAmelCase	,			unittest.TestCase       ):







  """simple docstring"""




  UpperCAmelCase__      : Union[str, Any]			       =XLMRobertaTokenizer
  UpperCAmelCase__      : Optional[Any]			       =XLMRobertaTokenizerFast
  UpperCAmelCase__      : Union[str, Any]			       =True
  UpperCAmelCase__      : List[Any]			       =True


  def  _lowercase      (			self    : int		)					->Dict:
        """simple docstring"""







        super().setUp()

        # We have a SentencePiece fixture for testing
        SCREAMING_SNAKE_CASE :      int       =   XLMRobertaTokenizer(UpperCAmelCase__				,				keep_accents=UpperCAmelCase__		)
        tokenizer.save_pretrained(self.tmpdirname		)


  def  _lowercase      (			self    : List[Any]		)					->Tuple:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[Any]       =   """<pad>"""
        SCREAMING_SNAKE_CASE :      Dict       =   1

        self.assertEqual(self.get_tokenizer()._convert_token_to_id(UpperCAmelCase__		)				,				UpperCAmelCase__		)
        self.assertEqual(self.get_tokenizer()._convert_id_to_token(UpperCAmelCase__		)				,				UpperCAmelCase__		)


  def  _lowercase      (			self    : Any		)					->Dict:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Any       =   list(self.get_tokenizer().get_vocab().keys()		)

        self.assertEqual(vocab_keys[0]				,				"""<s>"""		)
        self.assertEqual(vocab_keys[1]				,				"""<pad>"""		)
        self.assertEqual(vocab_keys[-1]				,				"""<mask>"""		)
        self.assertEqual(len(UpperCAmelCase__		)				,				1_0_0_2		)


  def  _lowercase      (			self    : Optional[Any]		)					->Optional[Any]:
        """simple docstring"""







        self.assertEqual(self.get_tokenizer().vocab_size				,				1_0_0_2		)


  def  _lowercase      (			self    : int		)					->str:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Dict       =   XLMRobertaTokenizer(UpperCAmelCase__				,				keep_accents=UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      Optional[Any]       =   tokenizer.tokenize("""This is a test"""		)
        self.assertListEqual(UpperCAmelCase__				,				["""▁This""", """▁is""", """▁a""", """▁t""", """est"""]		)

        self.assertListEqual(
            tokenizer.convert_tokens_to_ids(UpperCAmelCase__		)				,				[value + tokenizer.fairseq_offset for value in [2_8_5, 4_6, 1_0, 1_7_0, 3_8_2]]				,				)

        SCREAMING_SNAKE_CASE :      Optional[Any]       =   tokenizer.tokenize("""I was born in 92000, and this is falsé."""		)
        self.assertListEqual(
            UpperCAmelCase__				,				[
                SPIECE_UNDERLINE + """I""",
                SPIECE_UNDERLINE + """was""",
                SPIECE_UNDERLINE + """b""",
                """or""",
                """n""",
                SPIECE_UNDERLINE + """in""",
                SPIECE_UNDERLINE + """""",
                """9""",
                """2""",
                """0""",
                """0""",
                """0""",
                """,""",
                SPIECE_UNDERLINE + """and""",
                SPIECE_UNDERLINE + """this""",
                SPIECE_UNDERLINE + """is""",
                SPIECE_UNDERLINE + """f""",
                """al""",
                """s""",
                """é""",
                """.""",
            ]				,				)
        SCREAMING_SNAKE_CASE :      Optional[Any]       =   tokenizer.convert_tokens_to_ids(UpperCAmelCase__		)
        self.assertListEqual(
            UpperCAmelCase__				,				[
                value + tokenizer.fairseq_offset
                for value in [8, 2_1, 8_4, 5_5, 2_4, 1_9, 7, 2, 6_0_2, 3_4_7, 3_4_7, 3_4_7, 3, 1_2, 6_6, 4_6, 7_2, 8_0, 6, 2, 4]
                #                                       ^ unk: 2 + 1 = 3                  unk: 2 + 1 = 3 ^
            ]				,				)

        SCREAMING_SNAKE_CASE :      Any       =   tokenizer.convert_ids_to_tokens(UpperCAmelCase__		)
        self.assertListEqual(
            UpperCAmelCase__				,				[
                SPIECE_UNDERLINE + """I""",
                SPIECE_UNDERLINE + """was""",
                SPIECE_UNDERLINE + """b""",
                """or""",
                """n""",
                SPIECE_UNDERLINE + """in""",
                SPIECE_UNDERLINE + """""",
                """<unk>""",
                """2""",
                """0""",
                """0""",
                """0""",
                """,""",
                SPIECE_UNDERLINE + """and""",
                SPIECE_UNDERLINE + """this""",
                SPIECE_UNDERLINE + """is""",
                SPIECE_UNDERLINE + """f""",
                """al""",
                """s""",
                """<unk>""",
                """.""",
            ]				,				)


  def  _lowercase      (			self    : List[Any]		)					->Dict:
        """simple docstring"""







        if not self.test_slow_tokenizer:
              # as we don't have a slow version, we can't compare the outputs between slow and fast versions
              return

        SCREAMING_SNAKE_CASE :      Tuple       =   (self.rust_tokenizer_class, """hf-internal-testing/tiny-xlm-roberta""", {})
        for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
              with self.subTest(f"{tokenizer.__class__.__name__} ({pretrained_name})"		):
                    SCREAMING_SNAKE_CASE :      List[str]       =   self.rust_tokenizer_class.from_pretrained(UpperCAmelCase__				,				**UpperCAmelCase__		)
                    SCREAMING_SNAKE_CASE :      Tuple       =   self.tokenizer_class.from_pretrained(UpperCAmelCase__				,				**UpperCAmelCase__		)

                    SCREAMING_SNAKE_CASE :      str       =   tempfile.mkdtemp()

                    SCREAMING_SNAKE_CASE :      str       =   tokenizer_r.save_pretrained(UpperCAmelCase__		)
                    SCREAMING_SNAKE_CASE :      str       =   tokenizer_p.save_pretrained(UpperCAmelCase__		)

                    # Checks it save with the same files + the tokenizer.json file for the fast one
                    self.assertTrue(any("""tokenizer.json""" in f for f in tokenizer_r_files		)		)
                    SCREAMING_SNAKE_CASE :      List[str]       =   tuple(f for f in tokenizer_r_files if """tokenizer.json""" not in f		)
                    self.assertSequenceEqual(UpperCAmelCase__				,				UpperCAmelCase__		)

                    # Checks everything loads correctly in the same way
                    SCREAMING_SNAKE_CASE :      Union[str, Any]       =   tokenizer_r.from_pretrained(UpperCAmelCase__		)
                    SCREAMING_SNAKE_CASE :      List[Any]       =   tokenizer_p.from_pretrained(UpperCAmelCase__		)

                    # Check special tokens are set accordingly on Rust and Python
                    for key in tokenizer_pp.special_tokens_map:
                          self.assertTrue(hasattr(UpperCAmelCase__				,				UpperCAmelCase__		)		)
                          # self.assertEqual(getattr(tokenizer_rp, key), getattr(tokenizer_pp, key))
                          # self.assertEqual(getattr(tokenizer_rp, key + "_id"), getattr(tokenizer_pp, key + "_id"))

                    shutil.rmtree(UpperCAmelCase__		)

                    # Save tokenizer rust, legacy_format=True
                    SCREAMING_SNAKE_CASE :      Optional[Any]       =   tempfile.mkdtemp()

                    SCREAMING_SNAKE_CASE :      Optional[Any]       =   tokenizer_r.save_pretrained(UpperCAmelCase__				,				legacy_format=UpperCAmelCase__		)
                    SCREAMING_SNAKE_CASE :      Tuple       =   tokenizer_p.save_pretrained(UpperCAmelCase__		)

                    # Checks it save with the same files
                    self.assertSequenceEqual(UpperCAmelCase__				,				UpperCAmelCase__		)

                    # Checks everything loads correctly in the same way
                    SCREAMING_SNAKE_CASE :      Optional[int]       =   tokenizer_r.from_pretrained(UpperCAmelCase__		)
                    SCREAMING_SNAKE_CASE :      Dict       =   tokenizer_p.from_pretrained(UpperCAmelCase__		)

                    # Check special tokens are set accordingly on Rust and Python
                    for key in tokenizer_pp.special_tokens_map:
                          self.assertTrue(hasattr(UpperCAmelCase__				,				UpperCAmelCase__		)		)

                    shutil.rmtree(UpperCAmelCase__		)

                    # Save tokenizer rust, legacy_format=False
                    SCREAMING_SNAKE_CASE :      Optional[int]       =   tempfile.mkdtemp()

                    SCREAMING_SNAKE_CASE :      Tuple       =   tokenizer_r.save_pretrained(UpperCAmelCase__				,				legacy_format=UpperCAmelCase__		)
                    SCREAMING_SNAKE_CASE :      Tuple       =   tokenizer_p.save_pretrained(UpperCAmelCase__		)

                    # Checks it saved the tokenizer.json file
                    self.assertTrue(any("""tokenizer.json""" in f for f in tokenizer_r_files		)		)

                    # Checks everything loads correctly in the same way
                    SCREAMING_SNAKE_CASE :      Optional[int]       =   tokenizer_r.from_pretrained(UpperCAmelCase__		)
                    SCREAMING_SNAKE_CASE :      int       =   tokenizer_p.from_pretrained(UpperCAmelCase__		)

                    # Check special tokens are set accordingly on Rust and Python
                    for key in tokenizer_pp.special_tokens_map:
                          self.assertTrue(hasattr(UpperCAmelCase__				,				UpperCAmelCase__		)		)

                    shutil.rmtree(UpperCAmelCase__		)


  @cached_property
  def  _lowercase      (			self    : List[Any]		)					->Dict:
        """simple docstring"""







        return XLMRobertaTokenizer.from_pretrained("""xlm-roberta-base"""		)


  def  _lowercase      (			self    : str		)					->Any:
        """simple docstring"""







        with tempfile.NamedTemporaryFile() as f:
              shutil.copyfile(UpperCAmelCase__				,				f.name		)
              SCREAMING_SNAKE_CASE :      Union[str, Any]       =   XLMRobertaTokenizer(f.name				,				keep_accents=UpperCAmelCase__		)
              SCREAMING_SNAKE_CASE :      str       =   pickle.dumps(UpperCAmelCase__		)
        pickle.loads(UpperCAmelCase__		)


  def  _lowercase      (			self    : int		)					->Optional[Any]:
        """simple docstring"""







        if not self.test_rust_tokenizer:
              return

        SCREAMING_SNAKE_CASE :      List[Any]       =   self.get_tokenizer()
        SCREAMING_SNAKE_CASE :      int       =   self.get_rust_tokenizer()

        SCREAMING_SNAKE_CASE :      Tuple       =   """I was born in 92000, and this is falsé."""

        SCREAMING_SNAKE_CASE :      Optional[int]       =   tokenizer.tokenize(UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE :      Optional[int]       =   rust_tokenizer.tokenize(UpperCAmelCase__		)
        self.assertListEqual(UpperCAmelCase__				,				UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      Tuple       =   tokenizer.encode(UpperCAmelCase__				,				add_special_tokens=UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE :      Dict       =   rust_tokenizer.encode(UpperCAmelCase__				,				add_special_tokens=UpperCAmelCase__		)
        self.assertListEqual(UpperCAmelCase__				,				UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      int       =   self.get_rust_tokenizer()
        SCREAMING_SNAKE_CASE :      List[Any]       =   tokenizer.encode(UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE :      Dict       =   rust_tokenizer.encode(UpperCAmelCase__		)
        self.assertListEqual(UpperCAmelCase__				,				UpperCAmelCase__		)


  @slow
  def  _lowercase      (			self    : Dict		)					->Tuple:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Any       =   """Hello World!"""
        SCREAMING_SNAKE_CASE :      str       =   [0, 3_5_3_7_8, 6_6_6_1, 3_8, 2]
        # xlmr = torch.hub.load('pytorch/fairseq', 'xlmr.base')  # xlmr.large has same tokenizer
        # xlmr.eval()
        # xlmr.encode(symbols)

        self.assertListEqual(UpperCAmelCase__				,				self.big_tokenizer.encode(UpperCAmelCase__		)		)


  @slow
  def  _lowercase      (			self    : int		)					->List[str]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      List[str]       =   (
            """This is a very long text with a lot of weird characters, such as: . , ~ ? ( ) \" [ ] ! : - . Also we will"""
            """ add words that should not exsist and be tokenized to <unk>, such as saoneuhaoesuth"""
        )
        SCREAMING_SNAKE_CASE :      str       =   [
            0,
            3_2_9_3,
            8_3,
            1_0,
            4_5_5_2,
            4_9_8_9,
            7_9_8_6,
            6_7_8,
            1_0,
            5_9_1_5,
            1_1_1,
            1_7_9_4_5_9,
            1_2_4_8_5_0,
            4,
            6_0_4_4,
            2_3_7,
            1_2,
            6,
            5,
            6,
            4,
            6_7_8_0,
            7_0_5,
            1_5,
            1_3_8_8,
            4_4,
            3_7_8,
            1_0_1_1_4,
            7_1_1,
            1_5_2,
            2_0,
            6,
            5,
            2_2_3_7_6,
            6_4_2,
            1_2_2_1,
            1_5_1_9_0,
            3_4_1_5_3,
            4_5_0,
            5_6_0_8,
            9_5_9,
            1_1_1_9,
            5_7_7_0_2,
            1_3_6,
            1_8_6,
            4_7,
            1_0_9_8,
            2_9_3_6_7,
            4_7,
            # 4426, # What fairseq tokenizes from "<unk>": "_<"
            # 3678, # What fairseq tokenizes from "<unk>": "unk"
            # 2740, # What fairseq tokenizes from "<unk>": ">"
            3,  # What we tokenize from "<unk>": "<unk>"
            6,  # Residue from the tokenization: an extra sentencepiece underline
            4,
            6_0_4_4,
            2_3_7,
            6_2_8_4,
            5_0_9_0_1,
            5_2_8,
            3_1,
            9_0,
            3_4,
            9_2_7,
            2,
        ]
        # xlmr = torch.hub.load('pytorch/fairseq', 'xlmr.base')  # xlmr.large has same tokenizer
        # xlmr.eval()
        # xlmr.encode(symbols)

        self.assertListEqual(UpperCAmelCase__				,				self.big_tokenizer.encode(UpperCAmelCase__		)		)




  @slow
  def  _lowercase      (			self    : List[str]		)					->Union[str, Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[Any]       =   {"""input_ids""": [[0, 1_1_0_6_2, 8_2_7_7_2, 7, 1_5, 8_2_7_7_2, 5_3_8, 5_1_5_2_9, 2_3_7, 1_7_1_9_8, 1_2_9_0, 2_0_6, 9, 2_1_5_1_7_5, 1_3_1_4, 1_3_6, 1_7_1_9_8, 1_2_9_0, 2_0_6, 9, 5_6_3_5_9, 4_2, 1_2_2_0_0_9, 9, 1_6_4_6_6, 1_6, 8_7_3_4_4, 4_5_3_7, 9, 4_7_1_7, 7_8_3_8_1, 6, 1_5_9_9_5_8, 7, 1_5, 2_4_4_8_0, 6_1_8, 4, 5_2_7, 2_2_6_9_3, 5_4_2_8, 4, 2_7_7_7, 2_4_4_8_0, 9_8_7_4, 4, 4_3_5_2_3, 5_9_4, 4, 8_0_3, 1_8_3_9_2, 3_3_1_8_9, 1_8, 4, 4_3_5_2_3, 2_4_4_4_7, 1_2_3_9_9, 1_0_0, 2_4_9_5_5, 8_3_6_5_8, 9_6_2_6, 1_4_4_0_5_7, 1_5, 8_3_9, 2_2_3_3_5, 1_6, 1_3_6, 2_4_9_5_5, 8_3_6_5_8, 8_3_4_7_9, 1_5, 3_9_1_0_2, 7_2_4, 1_6, 6_7_8, 6_4_5, 2_7_8_9, 1_3_2_8, 4_5_8_9, 4_2, 1_2_2_0_0_9, 1_1_5_7_7_4, 2_3, 8_0_5, 1_3_2_8, 4_6_8_7_6, 7, 1_3_6, 5_3_8_9_4, 1_9_4_0, 4_2_2_2_7, 4_1_1_5_9, 1_7_7_2_1, 8_2_3, 4_2_5, 4, 2_7_5_1_2, 9_8_7_2_2, 2_0_6, 1_3_6, 5_5_3_1, 4_9_7_0, 9_1_9, 1_7_3_3_6, 5, 2], [0, 2_0_0_8_0, 6_1_8, 8_3, 8_2_7_7_5, 4_7, 4_7_9, 9, 1_5_1_7, 7_3, 5_3_8_9_4, 3_3_3, 8_0_5_8_1, 1_1_0_1_1_7, 1_8_8_1_1, 5_2_5_6, 1_2_9_5, 5_1, 1_5_2_5_2_6, 2_9_7, 7_9_8_6, 3_9_0, 1_2_4_4_1_6, 5_3_8, 3_5_4_3_1, 2_1_4, 9_8, 1_5_0_4_4, 2_5_7_3_7, 1_3_6, 7_1_0_8, 4_3_7_0_1, 2_3, 7_5_6, 1_3_5_3_5_5, 7, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 5_8_1, 6_3_7_7_3, 1_1_9_4_5_5, 6, 1_4_7_7_9_7, 8_8_2_0_3, 7, 6_4_5, 7_0, 2_1, 3_2_8_5, 1_0_2_6_9, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], """attention_mask""": [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
        # fmt: on

        self.tokenizer_integration_test_util(
            expected_encoding=UpperCAmelCase__				,				model_name="""xlm-roberta-base"""				,				revision="""d9d8a8ea5eb94b1c6654ae9249df7793cd2933d3"""				,				)







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       ):







  """simple docstring"""




  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  _lowercase      (			self    : Tuple		)					->Optional[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      List[str]       =   []
        SCREAMING_SNAKE_CASE :      Any       =   []

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

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

              SCREAMING_SNAKE_CASE :      List[Any]       =   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__		)

        SCREAMING_SNAKE_CASE :      Optional[Any]       =   resnets
        SCREAMING_SNAKE_CASE :      Optional[int]       =   attentions

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




  def __call__(			self    : int				,				UpperCAmelCase__    : List[str]				,				UpperCAmelCase__    : Union[str, Any]				,				UpperCAmelCase__    : Optional[Any]				,				UpperCAmelCase__    : List[str]=True		)					->str:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[int]       =   ()

        for resnet, attn in zip(self.resnets				,				self.attentions		):
              SCREAMING_SNAKE_CASE :      List[str]       =   resnet(UpperCAmelCase__				,				UpperCAmelCase__				,				deterministic=UpperCAmelCase__		)
              SCREAMING_SNAKE_CASE :      List[str]       =   attn(UpperCAmelCase__				,				UpperCAmelCase__				,				deterministic=UpperCAmelCase__		)
              output_states += (hidden_states,)

        if self.add_downsample:
              SCREAMING_SNAKE_CASE :      int       =   self.downsamplers_a(UpperCAmelCase__		)
              output_states += (hidden_states,)

        return hidden_states, output_states







class        a__       (   nn.Module       ):







  """simple docstring"""




  UpperCAmelCase__      : int
  UpperCAmelCase__      : int
  UpperCAmelCase__      : float			       =0.0
  UpperCAmelCase__      : int			       =1
  UpperCAmelCase__      : bool			       =True
  UpperCAmelCase__      : jnp.dtype			       =jnp.floataa


  def  _lowercase      (			self    : Dict		)					->Union[str, Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Dict       =   []

        for i in range(self.num_layers		):
              SCREAMING_SNAKE_CASE :      Union[str, Any]       =   self.in_channels if i == 0 else self.out_channels

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

        if self.add_downsample:
              SCREAMING_SNAKE_CASE :      List[str]       =   FlaxDownsampleaD(self.out_channels				,				dtype=self.dtype		)




  def __call__(			self    : Tuple				,				UpperCAmelCase__    : int				,				UpperCAmelCase__    : Optional[Any]				,				UpperCAmelCase__    : Any=True		)					->Optional[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[int]       =   ()

        for resnet in self.resnets:
              SCREAMING_SNAKE_CASE :      List[str]       =   resnet(UpperCAmelCase__				,				UpperCAmelCase__				,				deterministic=UpperCAmelCase__		)
              output_states += (hidden_states,)

        if self.add_downsample:
              SCREAMING_SNAKE_CASE :      Dict       =   self.downsamplers_a(UpperCAmelCase__		)
              output_states += (hidden_states,)

        return hidden_states, output_states







class        a__       (   nn.Module       ):







  """simple docstring"""




  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  _lowercase      (			self    : int		)					->Any:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Dict       =   []
        SCREAMING_SNAKE_CASE :      List[Any]       =   []

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

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

              SCREAMING_SNAKE_CASE :      Tuple       =   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__		)

        SCREAMING_SNAKE_CASE :      Optional[Any]       =   resnets
        SCREAMING_SNAKE_CASE :      List[Any]       =   attentions

        if self.add_upsample:
              SCREAMING_SNAKE_CASE :      Union[str, Any]       =   FlaxUpsampleaD(self.out_channels				,				dtype=self.dtype		)




  def __call__(			self    : Union[str, Any]				,				UpperCAmelCase__    : str				,				UpperCAmelCase__    : List[Any]				,				UpperCAmelCase__    : Any				,				UpperCAmelCase__    : Union[str, Any]				,				UpperCAmelCase__    : Dict=True		)					->Tuple:
        """simple docstring"""







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

              SCREAMING_SNAKE_CASE :      Tuple       =   resnet(UpperCAmelCase__				,				UpperCAmelCase__				,				deterministic=UpperCAmelCase__		)
              SCREAMING_SNAKE_CASE :      Tuple       =   attn(UpperCAmelCase__				,				UpperCAmelCase__				,				deterministic=UpperCAmelCase__		)

        if self.add_upsample:
              SCREAMING_SNAKE_CASE :      Tuple       =   self.upsamplers_a(UpperCAmelCase__		)

        return hidden_states







class        a__       (   nn.Module       ):







  """simple docstring"""




  UpperCAmelCase__      : int
  UpperCAmelCase__      : int
  UpperCAmelCase__      : int
  UpperCAmelCase__      : float			       =0.0
  UpperCAmelCase__      : int			       =1
  UpperCAmelCase__      : bool			       =True
  UpperCAmelCase__      : jnp.dtype			       =jnp.floataa


  def  _lowercase      (			self    : Union[str, Any]		)					->Optional[int]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[Any]       =   []

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

              SCREAMING_SNAKE_CASE :      Optional[int]       =   FlaxResnetBlockaD(
                  in_channels=resnet_in_channels + res_skip_channels				,				out_channels=self.out_channels				,				dropout_prob=self.dropout				,				dtype=self.dtype				,				)
              resnets.append(UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      Dict       =   resnets

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




  def __call__(			self    : Optional[Any]				,				UpperCAmelCase__    : List[str]				,				UpperCAmelCase__    : Dict				,				UpperCAmelCase__    : Tuple				,				UpperCAmelCase__    : Optional[int]=True		)					->Union[str, Any]:
        """simple docstring"""







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

              SCREAMING_SNAKE_CASE :      Optional[int]       =   resnet(UpperCAmelCase__				,				UpperCAmelCase__				,				deterministic=UpperCAmelCase__		)

        if self.add_upsample:
              SCREAMING_SNAKE_CASE :      Optional[int]       =   self.upsamplers_a(UpperCAmelCase__		)

        return hidden_states







class        a__       (   nn.Module       ):







  """simple docstring"""




  UpperCAmelCase__      : int
  UpperCAmelCase__      : float			       =0.0
  UpperCAmelCase__      : int			       =1
  UpperCAmelCase__      : int			       =1
  UpperCAmelCase__      : bool			       =False
  UpperCAmelCase__      : bool			       =False
  UpperCAmelCase__      : jnp.dtype			       =jnp.floataa


  def  _lowercase      (			self    : Optional[Any]		)					->List[Any]:
        """simple docstring"""







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

        SCREAMING_SNAKE_CASE :      List[str]       =   []

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

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

        SCREAMING_SNAKE_CASE :      int       =   resnets
        SCREAMING_SNAKE_CASE :      Tuple       =   attentions




  def __call__(			self    : Union[str, Any]				,				UpperCAmelCase__    : List[str]				,				UpperCAmelCase__    : Union[str, Any]				,				UpperCAmelCase__    : Any				,				UpperCAmelCase__    : str=True		)					->Optional[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Tuple       =   self.resnets[0](UpperCAmelCase__				,				UpperCAmelCase__		)
        for attn, resnet in zip(self.attentions				,				self.resnets[1:]		):
              SCREAMING_SNAKE_CASE :      Optional[int]       =   attn(UpperCAmelCase__				,				UpperCAmelCase__				,				deterministic=UpperCAmelCase__		)
              SCREAMING_SNAKE_CASE :      List[str]       =   resnet(UpperCAmelCase__				,				UpperCAmelCase__				,				deterministic=UpperCAmelCase__		)

        return hidden_states




