
import shutil
import tempfile
import unittest

from transformers import SPIECE_UNDERLINE, BatchEncoding, MBartaaTokenizer, MBartaaTokenizerFast, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
)

from ...test_tokenization_common import TokenizerTesterMixin


__lowerCamelCase							=       get_tests_dir("""fixtures/test_sentencepiece.model""")

if is_torch_available():
   from transformers.models.mbart.modeling_mbart import shift_tokens_right

__lowerCamelCase							=       25_00_04
__lowerCamelCase							=       25_00_20






@require_sentencepiece
@require_tokenizers
class   UpperCAmelCase		(						A_       ,unittest.TestCase ):
   A__					:		Dict						     =     MBartaaTokenizer
   A__					:		List[str]						     =     MBartaaTokenizerFast
   A__					:		Optional[int]						     =     True
   A__					:		Union[str, Any]						     =     True



   def 				_SCREAMING_SNAKE_CASE   (self		:					int )		->	Any:


          '''simple docstring'''


          super().setUp()

          # We have a SentencePiece fixture for testing
          snake_case       :   str							=      MBartaaTokenizer(snake_case__    ,						src_lang="en_XX"    ,						tgt_lang="ro_RO"    ,						keep_accents=snake_case__ )
          tokenizer.save_pretrained(self.tmpdirname )



   def 				_SCREAMING_SNAKE_CASE   (self		:					List[str] )		->	Tuple:


          '''simple docstring'''


          snake_case       :   List[str]							=      "<s>"
          snake_case       :   Dict							=      0

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



   def 				_SCREAMING_SNAKE_CASE   (self		:					str )		->	Optional[int]:


          '''simple docstring'''


          snake_case       :   Tuple							=      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(snake_case__ )    ,						10_54 )



   def 				_SCREAMING_SNAKE_CASE   (self		:					int )		->	Tuple:


          '''simple docstring'''


          self.assertEqual(self.get_tokenizer().vocab_size    ,						10_54 )



   def 				_SCREAMING_SNAKE_CASE   (self		:					Any )		->	int:


          '''simple docstring'''


          snake_case       :   Tuple							=      MBartaaTokenizer(snake_case__    ,						src_lang="en_XX"    ,						tgt_lang="ro_RO"    ,						keep_accents=snake_case__ )

          snake_case       :   Dict							=      tokenizer.tokenize("This is a test" )
          self.assertListEqual(snake_case__    ,						["▁This", "▁is", "▁a", "▁t", "est"] )

          self.assertListEqual(
              tokenizer.convert_tokens_to_ids(snake_case__ )    ,						[value + tokenizer.fairseq_offset for value in [2_85, 46, 10, 1_70, 3_82]]    ,						)

          snake_case       :   Optional[int]							=      tokenizer.tokenize("I was born in 92000, and this is falsé." )
          self.assertListEqual(
              snake_case__    ,						[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", "é", "."]    ,						)
          snake_case       :   Optional[int]							=      tokenizer.convert_tokens_to_ids(snake_case__ )
          self.assertListEqual(
              snake_case__    ,						[
                  value + tokenizer.fairseq_offset
                  for value in [8, 21, 84, 55, 24, 19, 7, 2, 6_02, 3_47, 3_47, 3_47, 3, 12, 66, 46, 72, 80, 6, 2, 4]
              ]    ,						)

          snake_case       :   Tuple							=      tokenizer.convert_ids_to_tokens(snake_case__ )
          self.assertListEqual(
              snake_case__    ,						[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>", "."]    ,						)



   @slow
   def 				_SCREAMING_SNAKE_CASE   (self		:					Union[str, Any] )		->	Optional[Any]:


          '''simple docstring'''


          snake_case       :   Any							=      {"input_ids": [[25_00_04, 1_10_62, 8_27_72, 7, 15, 8_27_72, 5_38, 5_15_29, 2_37, 1_71_98, 12_90, 2_06, 9, 21_51_75, 13_14, 1_36, 1_71_98, 12_90, 2_06, 9, 5_63_59, 42, 12_20_09, 9, 1_64_66, 16, 8_73_44, 45_37, 9, 47_17, 7_83_81, 6, 15_99_58, 7, 15, 2_44_80, 6_18, 4, 5_27, 2_26_93, 54_28, 4, 27_77, 2_44_80, 98_74, 4, 4_35_23, 5_94, 4, 8_03, 1_83_92, 3_31_89, 18, 4, 4_35_23, 2_44_47, 1_23_99, 1_00, 2_49_55, 8_36_58, 96_26, 14_40_57, 15, 8_39, 2_23_35, 16, 1_36, 2_49_55, 8_36_58, 8_34_79, 15, 3_91_02, 7_24, 16, 6_78, 6_45, 27_89, 13_28, 45_89, 42, 12_20_09, 11_57_74, 23, 8_05, 13_28, 4_68_76, 7, 1_36, 5_38_94, 19_40, 4_22_27, 4_11_59, 1_77_21, 8_23, 4_25, 4, 2_75_12, 9_87_22, 2_06, 1_36, 55_31, 49_70, 9_19, 1_73_36, 5, 2], [25_00_04, 2_00_80, 6_18, 83, 8_27_75, 47, 4_79, 9, 15_17, 73, 5_38_94, 3_33, 8_05_81, 11_01_17, 1_88_11, 52_56, 12_95, 51, 15_25_26, 2_97, 79_86, 3_90, 12_44_16, 5_38, 3_54_31, 2_14, 98, 1_50_44, 2_57_37, 1_36, 71_08, 4_37_01, 23, 7_56, 13_53_55, 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], [25_00_04, 5_81, 6_37_73, 11_94_55, 6, 14_77_97, 8_82_03, 7, 6_45, 70, 21, 32_85, 1_02_69, 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=snake_case__    ,						model_name="facebook/mbart-large-50"    ,						revision="d3913889c59cd5c9e456b269c376325eabad57e2"    ,						)



   def 				_SCREAMING_SNAKE_CASE   (self		:					List[str] )		->	List[Any]:


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

          snake_case       :   int							=      (self.rust_tokenizer_class, "hf-internal-testing/tiny-random-mbart50", {})
          for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
                 with self.subTest(f"""{tokenizer.__class__.__name__} ({pretrained_name})""" ):
                        snake_case       :   Union[str, Any]							=      self.rust_tokenizer_class.from_pretrained(snake_case__    ,						**snake_case__ )
                        snake_case       :   Tuple							=      self.tokenizer_class.from_pretrained(snake_case__    ,						**snake_case__ )

                        snake_case       :   List[str]							=      tempfile.mkdtemp()

                        snake_case       :   str							=      tokenizer_r.save_pretrained(snake_case__ )
                        snake_case       :   Optional[Any]							=      tokenizer_p.save_pretrained(snake_case__ )

                        # 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 ) )
                        snake_case       :   int							=      tuple(f for f in tokenizer_r_files if "tokenizer.json" not in f )
                        self.assertSequenceEqual(snake_case__    ,						snake_case__ )

                        # Checks everything loads correctly in the same way
                        snake_case       :   str							=      tokenizer_r.from_pretrained(snake_case__ )
                        snake_case       :   Dict							=      tokenizer_p.from_pretrained(snake_case__ )

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

                        shutil.rmtree(snake_case__ )

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

                        snake_case       :   List[Any]							=      tokenizer_r.save_pretrained(snake_case__    ,						legacy_format=snake_case__ )
                        snake_case       :   List[str]							=      tokenizer_p.save_pretrained(snake_case__ )

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

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

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

                        shutil.rmtree(snake_case__ )

                        # Save tokenizer rust, legacy_format=False
                        snake_case       :   List[str]							=      tempfile.mkdtemp()

                        snake_case       :   int							=      tokenizer_r.save_pretrained(snake_case__    ,						legacy_format=snake_case__ )
                        snake_case       :   List[str]							=      tokenizer_p.save_pretrained(snake_case__ )

                        # 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
                        snake_case       :   List[str]							=      tokenizer_r.from_pretrained(snake_case__ )
                        snake_case       :   int							=      tokenizer_p.from_pretrained(snake_case__ )

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

                        shutil.rmtree(snake_case__ )






@require_torch
@require_sentencepiece
@require_tokenizers
class   UpperCAmelCase		(						unittest.TestCase ):
   A__					:		str						     =     "facebook/mbart-large-50-one-to-many-mmt"
   A__					:		str						     =     [
       " UN Chief Says There Is No Military Solution in Syria",
       " Secretary-General Ban Ki-moon says his response to Russia's stepped up military support for Syria is that \"there is no military solution\" to the nearly five-year conflict and more weapons will only worsen the violence and misery for millions of people.",
   ]
   A__					:		Tuple						     =     [
       "Şeful ONU declară că nu există o soluţie militară în Siria",
       "Secretarul General Ban Ki-moon declară că răspunsul său la intensificarea sprijinului militar al Rusiei"
       " pentru Siria este că \"nu există o soluţie militară\" la conflictul de aproape cinci ani şi că noi arme nu vor"
       " face decât să înrăutăţească violenţele şi mizeria pentru milioane de oameni.",
   ]
   A__					:		List[str]						     =     [EN_CODE, 82_74, 12_78_73, 2_59_16, 7, 86_22, 20_71, 4_38, 6_74_85, 53, 18_78_95, 23, 5_17_12, 2]



   @classmethod
   def 				_SCREAMING_SNAKE_CASE   (cls		:					Optional[int] )		->	Dict:


          '''simple docstring'''


          snake_case       :   MBartaaTokenizer							=      MBartaaTokenizer.from_pretrained(
              cls.checkpoint_name    ,						src_lang="en_XX"    ,						tgt_lang="ro_RO" )
          snake_case       :   List[Any]							=      1
          return cls



   def 				_SCREAMING_SNAKE_CASE   (self		:					int )		->	Any:


          '''simple docstring'''


          self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["ar_AR"]    ,						25_00_01 )
          self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["en_EN"]    ,						25_00_04 )
          self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["ro_RO"]    ,						25_00_20 )
          self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["mr_IN"]    ,						25_00_38 )



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


          '''simple docstring'''


          snake_case       :   Optional[int]							=      self.tokenizer.batch_encode_plus(self.src_text ).input_ids[0]
          self.assertListEqual(self.expected_src_tokens    ,						snake_case__ )



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


          '''simple docstring'''


          self.assertIn(snake_case__    ,						self.tokenizer.all_special_ids )
          snake_case       :   List[str]							=      [RO_CODE, 8_84, 90_19, 96, 9, 9_16, 8_67_92, 36, 1_87_43, 1_55_96, 5, 2]
          snake_case       :   Optional[int]							=      self.tokenizer.decode(snake_case__    ,						skip_special_tokens=snake_case__ )
          snake_case       :   int							=      self.tokenizer.decode(generated_ids[1:]    ,						skip_special_tokens=snake_case__ )
          self.assertEqual(snake_case__    ,						snake_case__ )
          self.assertNotIn(self.tokenizer.eos_token    ,						snake_case__ )



   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[Any] )		->	Dict:


          '''simple docstring'''


          snake_case       :   Union[str, Any]							=      ["this is gunna be a long sentence " * 20]
          assert isinstance(src_text[0]    ,						snake_case__ )
          snake_case       :   List[str]							=      10
          snake_case       :   Tuple							=      self.tokenizer(snake_case__    ,						max_length=snake_case__    ,						truncation=snake_case__ ).input_ids[0]
          self.assertEqual(ids[0]    ,						snake_case__ )
          self.assertEqual(ids[-1]    ,						2 )
          self.assertEqual(len(snake_case__ )    ,						snake_case__ )



   def 				_SCREAMING_SNAKE_CASE   (self		:					str )		->	Optional[int]:


          '''simple docstring'''


          self.assertListEqual(self.tokenizer.convert_tokens_to_ids(["<mask>", "ar_AR"] )    ,						[25_00_53, 25_00_01] )



   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[Any] )		->	List[str]:


          '''simple docstring'''


          snake_case       :   Optional[Any]							=      tempfile.mkdtemp()
          snake_case       :   Union[str, Any]							=      self.tokenizer.fairseq_tokens_to_ids
          self.tokenizer.save_pretrained(snake_case__ )
          snake_case       :   Optional[Any]							=      MBartaaTokenizer.from_pretrained(snake_case__ )
          self.assertDictEqual(new_tok.fairseq_tokens_to_ids    ,						snake_case__ )



   @require_torch
   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[int] )		->	Optional[Any]:


          '''simple docstring'''


          snake_case       :   Union[str, Any]							=      self.tokenizer(self.src_text    ,						text_target=self.tgt_text    ,						padding=snake_case__    ,						return_tensors="pt" )
          snake_case       :   int							=      shift_tokens_right(batch["labels"]    ,						self.tokenizer.pad_token_id )

          # fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
          assert batch.input_ids[1][0] == EN_CODE
          assert batch.input_ids[1][-1] == 2
          assert batch.labels[1][0] == RO_CODE
          assert batch.labels[1][-1] == 2
          assert batch.decoder_input_ids[1][:2].tolist() == [2, RO_CODE]



   @require_torch
   def 				_SCREAMING_SNAKE_CASE   (self		:					Union[str, Any] )		->	Optional[Any]:


          '''simple docstring'''


          snake_case       :   Union[str, Any]							=      self.tokenizer(
              self.src_text    ,						text_target=self.tgt_text    ,						padding=snake_case__    ,						truncation=snake_case__    ,						max_length=len(self.expected_src_tokens )    ,						return_tensors="pt"    ,						)
          snake_case       :   List[Any]							=      shift_tokens_right(batch["labels"]    ,						self.tokenizer.pad_token_id )

          self.assertIsInstance(snake_case__    ,						snake_case__ )

          self.assertEqual((2, 14)    ,						batch.input_ids.shape )
          self.assertEqual((2, 14)    ,						batch.attention_mask.shape )
          snake_case       :   Optional[int]							=      batch.input_ids.tolist()[0]
          self.assertListEqual(self.expected_src_tokens    ,						snake_case__ )
          self.assertEqual(2    ,						batch.decoder_input_ids[0, 0] )  # decoder_start_token_id
          # Test that special tokens are reset
          self.assertEqual(self.tokenizer.prefix_tokens    ,						[EN_CODE] )
          self.assertEqual(self.tokenizer.suffix_tokens    ,						[self.tokenizer.eos_token_id] )



   def 				_SCREAMING_SNAKE_CASE   (self		:					List[Any] )		->	Optional[Any]:


          '''simple docstring'''


          snake_case       :   Optional[Any]							=      self.tokenizer(self.src_text    ,						padding=snake_case__    ,						truncation=snake_case__    ,						max_length=3    ,						return_tensors="pt" )
          snake_case       :   Optional[int]							=      self.tokenizer(
              text_target=self.tgt_text    ,						padding=snake_case__    ,						truncation=snake_case__    ,						max_length=10    ,						return_tensors="pt" )
          snake_case       :   str							=      targets["input_ids"]
          snake_case       :   List[str]							=      shift_tokens_right(snake_case__    ,						self.tokenizer.pad_token_id )

          self.assertEqual(batch.input_ids.shape[1]    ,						3 )
          self.assertEqual(batch.decoder_input_ids.shape[1]    ,						10 )



   @require_torch
   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[Any] )		->	Any:


          '''simple docstring'''


          snake_case       :   List[Any]							=      self.tokenizer._build_translation_inputs(
              "A test"    ,						return_tensors="pt"    ,						src_lang="en_XX"    ,						tgt_lang="ar_AR" )

          self.assertEqual(
              nested_simplify(snake_case__ )    ,						{
                  # en_XX, A, test, EOS
                  "input_ids": [[25_00_04, 62, 30_34, 2]],
                  "attention_mask": [[1, 1, 1, 1]],
                  # ar_AR
                  "forced_bos_token_id": 25_00_01,
              }    ,						)







from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)






@flax.struct.dataclass
class   UpperCAmelCase		(						A_ ):
   A__					:		jnp.ndarray






@flax_register_to_config
class   UpperCAmelCase		(						nn.Module       ,A_       ,A_ ):
   A__					:		int						     =     32
   A__					:		int						     =     4
   A__					:		int						     =     4
   A__					:		Tuple[str]						     =     (
       "CrossAttnDownBlock2D",
       "CrossAttnDownBlock2D",
       "CrossAttnDownBlock2D",
       "DownBlock2D",
   )
   A__					:		Tuple[str]						     =     ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
   A__					:		Union[bool, Tuple[bool]]						     =     False
   A__					:		Tuple[int]						     =     (3_20, 6_40, 12_80, 12_80)
   A__					:		int						     =     2
   A__					:		Union[int, Tuple[int]]						     =     8
   A__					:		Optional[Union[int, Tuple[int]]]						     =     None
   A__					:		int						     =     12_80
   A__					:		float						     =     0.0
   A__					:		bool						     =     False
   A__					:		jnp.dtype						     =     jnp.floataa
   A__					:		bool						     =     True
   A__					:		int						     =     0
   A__					:		bool						     =     False



   def 				_SCREAMING_SNAKE_CASE   (self		:					Optional[int]    ,						snake_case__		:					jax.random.KeyArray )		->	FrozenDict:


          '''simple docstring'''


          snake_case       :   Dict							=      (1, self.in_channels, self.sample_size, self.sample_size)
          snake_case       :   Any							=      jnp.zeros(snake_case__    ,						dtype=jnp.floataa )
          snake_case       :   List[str]							=      jnp.ones((1,)    ,						dtype=jnp.intaa )
          snake_case       :   str							=      jnp.zeros((1, 1, self.cross_attention_dim)    ,						dtype=jnp.floataa )

          snake_case				,					snake_case       :   Optional[int]							=      jax.random.split(snake_case__ )
          snake_case       :   Union[str, Any]							=      {"params": params_rng, "dropout": dropout_rng}

          return self.init(snake_case__    ,						snake_case__    ,						snake_case__    ,						snake_case__ )["params"]



   def 				_SCREAMING_SNAKE_CASE   (self		:					str )		->	Tuple:


          '''simple docstring'''


          snake_case       :   str							=      self.block_out_channels
          snake_case       :   Optional[Any]							=      block_out_channels[0] * 4

          if self.num_attention_heads is not None:
                 raise ValueError(
                     "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19." )

          # If `num_attention_heads` is not defined (which is the case for most models)
          # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
          # The reason for this behavior is to correct for incorrectly named variables that were introduced
          # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
          # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
          # which is why we correct for the naming here.
          snake_case       :   Tuple							=      self.num_attention_heads or self.attention_head_dim

          # input
          snake_case       :   Tuple							=      nn.Conv(
              block_out_channels[0]    ,						kernel_size=(3, 3)    ,						strides=(1, 1)    ,						padding=((1, 1), (1, 1))    ,						dtype=self.dtype    ,						)

          # time
          snake_case       :   Union[str, Any]							=      FlaxTimesteps(
              block_out_channels[0]    ,						flip_sin_to_cos=self.flip_sin_to_cos    ,						freq_shift=self.config.freq_shift )
          snake_case       :   Dict							=      FlaxTimestepEmbedding(snake_case__    ,						dtype=self.dtype )

          snake_case       :   List[str]							=      self.only_cross_attention
          if isinstance(snake_case__    ,						snake_case__ ):
                 snake_case       :   List[Any]							=      (only_cross_attention,) * len(self.down_block_types )

          if isinstance(snake_case__    ,						snake_case__ ):
                 snake_case       :   List[Any]							=      (num_attention_heads,) * len(self.down_block_types )

          # down
          snake_case       :   List[Any]							=      []
          snake_case       :   Optional[int]							=      block_out_channels[0]
          for i, down_block_type in enumerate(self.down_block_types ):
                 snake_case       :   List[Any]							=      output_channel
                 snake_case       :   Dict							=      block_out_channels[i]
                 snake_case       :   Optional[Any]							=      i == len(snake_case__ ) - 1

                 if down_block_type == "CrossAttnDownBlock2D":
                        snake_case       :   List[Any]							=      FlaxCrossAttnDownBlockaD(
                            in_channels=snake_case__    ,						out_channels=snake_case__    ,						dropout=self.dropout    ,						num_layers=self.layers_per_block    ,						num_attention_heads=num_attention_heads[i]    ,						add_downsample=not is_final_block    ,						use_linear_projection=self.use_linear_projection    ,						only_cross_attention=only_cross_attention[i]    ,						use_memory_efficient_attention=self.use_memory_efficient_attention    ,						dtype=self.dtype    ,						)
                 else:
                        snake_case       :   Union[str, Any]							=      FlaxDownBlockaD(
                            in_channels=snake_case__    ,						out_channels=snake_case__    ,						dropout=self.dropout    ,						num_layers=self.layers_per_block    ,						add_downsample=not is_final_block    ,						dtype=self.dtype    ,						)

                 down_blocks.append(snake_case__ )
          snake_case       :   Dict							=      down_blocks

          # mid
          snake_case       :   Optional[int]							=      FlaxUNetMidBlockaDCrossAttn(
              in_channels=block_out_channels[-1]    ,						dropout=self.dropout    ,						num_attention_heads=num_attention_heads[-1]    ,						use_linear_projection=self.use_linear_projection    ,						use_memory_efficient_attention=self.use_memory_efficient_attention    ,						dtype=self.dtype    ,						)

          # up
          snake_case       :   Optional[Any]							=      []
          snake_case       :   Optional[int]							=      list(reversed(snake_case__ ) )
          snake_case       :   Dict							=      list(reversed(snake_case__ ) )
          snake_case       :   Tuple							=      list(reversed(snake_case__ ) )
          snake_case       :   Optional[Any]							=      reversed_block_out_channels[0]
          for i, up_block_type in enumerate(self.up_block_types ):
                 snake_case       :   Optional[int]							=      output_channel
                 snake_case       :   List[Any]							=      reversed_block_out_channels[i]
                 snake_case       :   Union[str, Any]							=      reversed_block_out_channels[min(i + 1    ,						len(snake_case__ ) - 1 )]

                 snake_case       :   int							=      i == len(snake_case__ ) - 1

                 if up_block_type == "CrossAttnUpBlock2D":
                        snake_case       :   Any							=      FlaxCrossAttnUpBlockaD(
                            in_channels=snake_case__    ,						out_channels=snake_case__    ,						prev_output_channel=snake_case__    ,						num_layers=self.layers_per_block + 1    ,						num_attention_heads=reversed_num_attention_heads[i]    ,						add_upsample=not is_final_block    ,						dropout=self.dropout    ,						use_linear_projection=self.use_linear_projection    ,						only_cross_attention=only_cross_attention[i]    ,						use_memory_efficient_attention=self.use_memory_efficient_attention    ,						dtype=self.dtype    ,						)
                 else:
                        snake_case       :   Optional[int]							=      FlaxUpBlockaD(
                            in_channels=snake_case__    ,						out_channels=snake_case__    ,						prev_output_channel=snake_case__    ,						num_layers=self.layers_per_block + 1    ,						add_upsample=not is_final_block    ,						dropout=self.dropout    ,						dtype=self.dtype    ,						)

                 up_blocks.append(snake_case__ )
                 snake_case       :   Optional[int]							=      output_channel
          snake_case       :   Tuple							=      up_blocks

          # out
          snake_case       :   Optional[int]							=      nn.GroupNorm(num_groups=32    ,						epsilon=1e-5 )
          snake_case       :   List[str]							=      nn.Conv(
              self.out_channels    ,						kernel_size=(3, 3)    ,						strides=(1, 1)    ,						padding=((1, 1), (1, 1))    ,						dtype=self.dtype    ,						)



   def __call__(self		:					Dict    ,						snake_case__		:					Dict    ,						snake_case__		:					Dict    ,						snake_case__		:					Optional[int]    ,						snake_case__		:					Tuple=None    ,						snake_case__		:					Union[str, Any]=None    ,						snake_case__		:					bool = True    ,						snake_case__		:					bool = False    ,						)		->	Union[FlaxUNetaDConditionOutput, Tuple]:


          '''simple docstring'''


          if not isinstance(snake_case__    ,						jnp.ndarray ):
                 snake_case       :   List[Any]							=      jnp.array([timesteps]    ,						dtype=jnp.intaa )
          elif isinstance(snake_case__    ,						jnp.ndarray ) and len(timesteps.shape ) == 0:
                 snake_case       :   Any							=      timesteps.astype(dtype=jnp.floataa )
                 snake_case       :   int							=      jnp.expand_dims(snake_case__    ,						0 )

          snake_case       :   str							=      self.time_proj(snake_case__ )
          snake_case       :   str							=      self.time_embedding(snake_case__ )

          # 2. pre-process
          snake_case       :   int							=      jnp.transpose(snake_case__    ,						(0, 2, 3, 1) )
          snake_case       :   List[Any]							=      self.conv_in(snake_case__ )

          # 3. down
          snake_case       :   Optional[int]							=      (sample,)
          for down_block in self.down_blocks:
                 if isinstance(snake_case__    ,						snake_case__ ):
                        snake_case				,					snake_case       :   List[Any]							=      down_block(snake_case__    ,						snake_case__    ,						snake_case__    ,						deterministic=not train )
                 else:
                        snake_case				,					snake_case       :   str							=      down_block(snake_case__    ,						snake_case__    ,						deterministic=not train )
                 down_block_res_samples += res_samples

          if down_block_additional_residuals is not None:
                 snake_case       :   Tuple							=      ()

                 for down_block_res_sample, down_block_additional_residual in zip(
                     snake_case__    ,						snake_case__ ):
                        down_block_res_sample += down_block_additional_residual
                        new_down_block_res_samples += (down_block_res_sample,)

                 snake_case       :   Optional[int]							=      new_down_block_res_samples

          # 4. mid
          snake_case       :   Optional[int]							=      self.mid_block(snake_case__    ,						snake_case__    ,						snake_case__    ,						deterministic=not train )

          if mid_block_additional_residual is not None:
                 sample += mid_block_additional_residual

          # 5. up
          for up_block in self.up_blocks:
                 snake_case       :   int							=      down_block_res_samples[-(self.layers_per_block + 1) :]
                 snake_case       :   Optional[Any]							=      down_block_res_samples[: -(self.layers_per_block + 1)]
                 if isinstance(snake_case__    ,						snake_case__ ):
                        snake_case       :   Optional[Any]							=      up_block(
                            snake_case__    ,						temb=snake_case__    ,						encoder_hidden_states=snake_case__    ,						res_hidden_states_tuple=snake_case__    ,						deterministic=not train    ,						)
                 else:
                        snake_case       :   Dict							=      up_block(snake_case__    ,						temb=snake_case__    ,						res_hidden_states_tuple=snake_case__    ,						deterministic=not train )

        # 6. post-process
          snake_case       :   List[str]							=      self.conv_norm_out(snake_case__ )
          snake_case       :   Any							=      nn.silu(snake_case__ )
          snake_case       :   Optional[int]							=      self.conv_out(snake_case__ )
          snake_case       :   Union[str, Any]							=      jnp.transpose(snake_case__    ,						(0, 3, 1, 2) )

          if not return_dict:
                 return (sample,)

          return FlaxUNetaDConditionOutput(sample=snake_case__ )





