


'''simple docstring'''



import json
import os
import sys
import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from huggingface_hub import HfFolder, Repository, create_repo, delete_repo
from requests.exceptions import HTTPError

import transformers
from transformers import (
    CONFIG_MAPPING,
    FEATURE_EXTRACTOR_MAPPING,
    PROCESSOR_MAPPING,
    TOKENIZER_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    AutoProcessor,
    AutoTokenizer,
    BertTokenizer,
    ProcessorMixin,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
)
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test
from transformers.tokenization_utils import TOKENIZER_CONFIG_FILE
from transformers.utils import FEATURE_EXTRACTOR_NAME, is_tokenizers_available


sys.path.append(str(Path(__file__).parent.parent.parent.parent / """utils"""))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402
from test_module.custom_processing import CustomProcessor  # noqa E402
from test_module.custom_tokenization import CustomTokenizer  # noqa E402


_lowercase      :      Union[str, Any]            =					get_tests_dir("""fixtures/dummy_feature_extractor_config.json""")
_lowercase      :      int            =					get_tests_dir("""fixtures/vocab.json""")
_lowercase      :      List[str]            =					get_tests_dir("""fixtures""")






class 	UpperCamelCase__(       unittest.TestCase	):
				__magic_name__    :				Dict  		=	["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]", "bla", "blou"]




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

										"""simple docstring"""




										UpperCAmelCase					    =     0




				def 					a__(		self      :			List[Any]				)->					int:

										"""simple docstring"""




										UpperCAmelCase					    =     AutoProcessor.from_pretrained('''facebook/wav2vec2-base-960h'''				)
										self.assertIsInstance(lowerCAmelCase				,   lowerCAmelCase				)




				def 					a__(		self      :			str				)->					Tuple:

										"""simple docstring"""




										with tempfile.TemporaryDirectory() as tmpdirname:
																UpperCAmelCase					    =     WavaVecaConfig()
																UpperCAmelCase					    =     AutoProcessor.from_pretrained('''facebook/wav2vec2-base-960h'''				)

																# save in new folder
																model_config.save_pretrained(lowerCAmelCase				)
																processor.save_pretrained(lowerCAmelCase				)

																UpperCAmelCase					    =     AutoProcessor.from_pretrained(lowerCAmelCase				)

										self.assertIsInstance(lowerCAmelCase				,   lowerCAmelCase				)




				def 					a__(		self      :			int				)->					int:

										"""simple docstring"""




										with tempfile.TemporaryDirectory() as tmpdirname:
																# copy relevant files
																copyfile(lowerCAmelCase				,   os.path.join(lowerCAmelCase				,   lowerCAmelCase				)				)
																copyfile(lowerCAmelCase				,   os.path.join(lowerCAmelCase				,   '''vocab.json'''				)				)

																UpperCAmelCase					    =     AutoProcessor.from_pretrained(lowerCAmelCase				)

										self.assertIsInstance(lowerCAmelCase				,   lowerCAmelCase				)




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

										"""simple docstring"""




										with tempfile.TemporaryDirectory() as tmpdirname:
																UpperCAmelCase					    =     WavaVecaFeatureExtractor()
																UpperCAmelCase					    =     AutoTokenizer.from_pretrained('''facebook/wav2vec2-base-960h'''				)

																UpperCAmelCase					    =     WavaVecaProcessor(lowerCAmelCase				,   lowerCAmelCase				)

																# save in new folder
																processor.save_pretrained(lowerCAmelCase				)

																# drop `processor_class` in tokenizer
																with open(os.path.join(lowerCAmelCase				,   lowerCAmelCase				)				,   '''r'''				) as f:
																						UpperCAmelCase					    =     json.load(lowerCAmelCase				)
																						config_dict.pop('''processor_class'''				)

																with open(os.path.join(lowerCAmelCase				,   lowerCAmelCase				)				,   '''w'''				) as f:
																						f.write(json.dumps(lowerCAmelCase				)				)

																UpperCAmelCase					    =     AutoProcessor.from_pretrained(lowerCAmelCase				)

										self.assertIsInstance(lowerCAmelCase				,   lowerCAmelCase				)




				def 					a__(		self      :			List[str]				)->					int:

										"""simple docstring"""




										with tempfile.TemporaryDirectory() as tmpdirname:
																UpperCAmelCase					    =     WavaVecaFeatureExtractor()
																UpperCAmelCase					    =     AutoTokenizer.from_pretrained('''facebook/wav2vec2-base-960h'''				)

																UpperCAmelCase					    =     WavaVecaProcessor(lowerCAmelCase				,   lowerCAmelCase				)

																# save in new folder
																processor.save_pretrained(lowerCAmelCase				)

																# drop `processor_class` in feature extractor
																with open(os.path.join(lowerCAmelCase				,   lowerCAmelCase				)				,   '''r'''				) as f:
																						UpperCAmelCase					    =     json.load(lowerCAmelCase				)
																						config_dict.pop('''processor_class'''				)

																with open(os.path.join(lowerCAmelCase				,   lowerCAmelCase				)				,   '''w'''				) as f:
																						f.write(json.dumps(lowerCAmelCase				)				)

																UpperCAmelCase					    =     AutoProcessor.from_pretrained(lowerCAmelCase				)

										self.assertIsInstance(lowerCAmelCase				,   lowerCAmelCase				)




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

										"""simple docstring"""




										with tempfile.TemporaryDirectory() as tmpdirname:
																UpperCAmelCase					    =     WavaVecaConfig(processor_class='''Wav2Vec2Processor'''				)
																model_config.save_pretrained(lowerCAmelCase				)
																# copy relevant files
																copyfile(lowerCAmelCase				,   os.path.join(lowerCAmelCase				,   '''vocab.json'''				)				)
																# create emtpy sample processor
																with open(os.path.join(lowerCAmelCase				,   lowerCAmelCase				)				,   '''w'''				) as f:
																						f.write('''{}'''				)

																UpperCAmelCase					    =     AutoProcessor.from_pretrained(lowerCAmelCase				)

										self.assertIsInstance(lowerCAmelCase				,   lowerCAmelCase				)




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

										"""simple docstring"""




										with self.assertRaises(lowerCAmelCase				):
																UpperCAmelCase					    =     AutoProcessor.from_pretrained('''hf-internal-testing/test_dynamic_processor'''				)
										# If remote code is disabled, we can't load this config.
										with self.assertRaises(lowerCAmelCase				):
																UpperCAmelCase					    =     AutoProcessor.from_pretrained(
																    '''hf-internal-testing/test_dynamic_processor'''				,   trust_remote_code=lowerCAmelCase				)

										UpperCAmelCase					    =     AutoProcessor.from_pretrained('''hf-internal-testing/test_dynamic_processor'''				,   trust_remote_code=lowerCAmelCase				)
										self.assertTrue(processor.special_attribute_present				)
										self.assertEqual(processor.__class__.__name__				,   '''NewProcessor'''				)

										UpperCAmelCase					    =     processor.feature_extractor
										self.assertTrue(feature_extractor.special_attribute_present				)
										self.assertEqual(feature_extractor.__class__.__name__				,   '''NewFeatureExtractor'''				)

										UpperCAmelCase					    =     processor.tokenizer
										self.assertTrue(tokenizer.special_attribute_present				)
										if is_tokenizers_available():
																self.assertEqual(tokenizer.__class__.__name__				,   '''NewTokenizerFast'''				)

																# Test we can also load the slow version
																UpperCAmelCase					    =     AutoProcessor.from_pretrained(
																    '''hf-internal-testing/test_dynamic_processor'''				,   trust_remote_code=lowerCAmelCase				,   use_fast=lowerCAmelCase				)
																UpperCAmelCase					    =     new_processor.tokenizer
																self.assertTrue(new_tokenizer.special_attribute_present				)
																self.assertEqual(new_tokenizer.__class__.__name__				,   '''NewTokenizer'''				)
										else:
																self.assertEqual(tokenizer.__class__.__name__				,   '''NewTokenizer'''				)




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

										"""simple docstring"""




										try:
																AutoConfig.register('''custom'''				,   lowerCAmelCase				)
																AutoFeatureExtractor.register(lowerCAmelCase				,   lowerCAmelCase				)
																AutoTokenizer.register(lowerCAmelCase				,   slow_tokenizer_class=lowerCAmelCase				)
																AutoProcessor.register(lowerCAmelCase				,   lowerCAmelCase				)
																# Trying to register something existing in the Transformers library will raise an error
																with self.assertRaises(lowerCAmelCase				):
																						AutoProcessor.register(lowerCAmelCase				,   lowerCAmelCase				)

																# Now that the config is registered, it can be used as any other config with the auto-API
																UpperCAmelCase					    =     CustomFeatureExtractor.from_pretrained(lowerCAmelCase				)

																with tempfile.TemporaryDirectory() as tmp_dir:
																						UpperCAmelCase					    =     os.path.join(lowerCAmelCase				,   '''vocab.txt'''				)
																						with open(lowerCAmelCase				,   '''w'''				,   encoding='''utf-8'''				) as vocab_writer:
																												vocab_writer.write(''''''.join([x + '''\n''' for x in self.vocab_tokens]				)				)
																						UpperCAmelCase					    =     CustomTokenizer(lowerCAmelCase				)

																UpperCAmelCase					    =     CustomProcessor(lowerCAmelCase				,   lowerCAmelCase				)

																with tempfile.TemporaryDirectory() as tmp_dir:
																						processor.save_pretrained(lowerCAmelCase				)
																						UpperCAmelCase					    =     AutoProcessor.from_pretrained(lowerCAmelCase				)
																						self.assertIsInstance(lowerCAmelCase				,   lowerCAmelCase				)

										finally:
																if "custom" in CONFIG_MAPPING._extra_content:
																						del CONFIG_MAPPING._extra_content["custom"]
																if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
																						del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]
																if CustomConfig in TOKENIZER_MAPPING._extra_content:
																						del TOKENIZER_MAPPING._extra_content[CustomConfig]
																if CustomConfig in PROCESSOR_MAPPING._extra_content:
																						del PROCESSOR_MAPPING._extra_content[CustomConfig]




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

										"""simple docstring"""




										class 	UpperCamelCase__(       lowerCAmelCase	):
														__magic_name__    :				List[str]  		=	False







										class 	UpperCamelCase__(       lowerCAmelCase	):
														__magic_name__    :				int  		=	False







										class 	UpperCamelCase__(       lowerCAmelCase	):
														__magic_name__    :				Optional[int]  		=	"AutoFeatureExtractor"
														__magic_name__    :				List[Any]  		=	"AutoTokenizer"
														__magic_name__    :				Any  		=	False

										try:
																AutoConfig.register('''custom'''				,   lowerCAmelCase				)
																AutoFeatureExtractor.register(lowerCAmelCase				,   lowerCAmelCase				)
																AutoTokenizer.register(lowerCAmelCase				,   slow_tokenizer_class=lowerCAmelCase				)
																AutoProcessor.register(lowerCAmelCase				,   lowerCAmelCase				)
																# If remote code is not set, the default is to use local classes.
																UpperCAmelCase					    =     AutoProcessor.from_pretrained('''hf-internal-testing/test_dynamic_processor'''				)
																self.assertEqual(processor.__class__.__name__				,   '''NewProcessor'''				)
																self.assertFalse(processor.special_attribute_present				)
																self.assertFalse(processor.feature_extractor.special_attribute_present				)
																self.assertFalse(processor.tokenizer.special_attribute_present				)

																# If remote code is disabled, we load the local ones.
																UpperCAmelCase					    =     AutoProcessor.from_pretrained(
																    '''hf-internal-testing/test_dynamic_processor'''				,   trust_remote_code=lowerCAmelCase				)
																self.assertEqual(processor.__class__.__name__				,   '''NewProcessor'''				)
																self.assertFalse(processor.special_attribute_present				)
																self.assertFalse(processor.feature_extractor.special_attribute_present				)
																self.assertFalse(processor.tokenizer.special_attribute_present				)

																# If remote is enabled, we load from the Hub.
																UpperCAmelCase					    =     AutoProcessor.from_pretrained(
																    '''hf-internal-testing/test_dynamic_processor'''				,   trust_remote_code=lowerCAmelCase				)
																self.assertEqual(processor.__class__.__name__				,   '''NewProcessor'''				)
																self.assertTrue(processor.special_attribute_present				)
																self.assertTrue(processor.feature_extractor.special_attribute_present				)
																self.assertTrue(processor.tokenizer.special_attribute_present				)

										finally:
																if "custom" in CONFIG_MAPPING._extra_content:
																						del CONFIG_MAPPING._extra_content["custom"]
																if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
																						del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]
																if CustomConfig in TOKENIZER_MAPPING._extra_content:
																						del TOKENIZER_MAPPING._extra_content[CustomConfig]
																if CustomConfig in PROCESSOR_MAPPING._extra_content:
																						del PROCESSOR_MAPPING._extra_content[CustomConfig]




				def 					a__(		self      :			Optional[int]				)->					Dict:

										"""simple docstring"""




										UpperCAmelCase					    =     AutoProcessor.from_pretrained('''hf-internal-testing/tiny-random-bert'''				)
										self.assertEqual(processor.__class__.__name__				,   '''BertTokenizerFast'''				)





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

										"""simple docstring"""




										UpperCAmelCase					    =     AutoProcessor.from_pretrained('''hf-internal-testing/tiny-random-convnext'''				)
										self.assertEqual(processor.__class__.__name__				,   '''ConvNextImageProcessor'''				)







@is_staging_test
class 	UpperCamelCase__(       unittest.TestCase	):
				__magic_name__    :				str  		=	["[UNK]", "[CLS]", "[SEP]", "[PAD]", "[MASK]", "bla", "blou"]




				@classmethod
				def 					a__(		cls      :			Optional[int]				)->					Tuple:

										"""simple docstring"""




										UpperCAmelCase					    =     TOKEN
										HfFolder.save_token(lowerCAmelCase				)




				@classmethod
				def 					a__(		cls      :			List[Any]				)->					Union[str, Any]:

										"""simple docstring"""




										try:
																delete_repo(token=cls._token				,   repo_id='''test-processor'''				)
										except HTTPError:
																pass

										try:
																delete_repo(token=cls._token				,   repo_id='''valid_org/test-processor-org'''				)
										except HTTPError:
																pass

										try:
																delete_repo(token=cls._token				,   repo_id='''test-dynamic-processor'''				)
										except HTTPError:
																pass




				def 					a__(		self      :			str				)->					Tuple:

										"""simple docstring"""




										UpperCAmelCase					    =     WavaVecaProcessor.from_pretrained(lowerCAmelCase				)
										with tempfile.TemporaryDirectory() as tmp_dir:
																processor.save_pretrained(
																    os.path.join(lowerCAmelCase				,   '''test-processor'''				)				,   push_to_hub=lowerCAmelCase				,   use_auth_token=self._token				)

																UpperCAmelCase					    =     WavaVecaProcessor.from_pretrained(F"""{USER}/test-processor"""				)
																for k, v in processor.feature_extractor.__dict__.items():
																						self.assertEqual(lowerCAmelCase				,   getattr(new_processor.feature_extractor				,   lowerCAmelCase				)				)
																self.assertDictEqual(new_processor.tokenizer.get_vocab()				,   processor.tokenizer.get_vocab()				)




				def 					a__(		self      :			Optional[Any]				)->					Optional[int]:

										"""simple docstring"""




										UpperCAmelCase					    =     WavaVecaProcessor.from_pretrained(lowerCAmelCase				)

										with tempfile.TemporaryDirectory() as tmp_dir:
																processor.save_pretrained(
																    os.path.join(lowerCAmelCase				,   '''test-processor-org'''				)				,   push_to_hub=lowerCAmelCase				,   use_auth_token=self._token				,   organization='''valid_org'''				,   )

																UpperCAmelCase					    =     WavaVecaProcessor.from_pretrained('''valid_org/test-processor-org'''				)
																for k, v in processor.feature_extractor.__dict__.items():
																						self.assertEqual(lowerCAmelCase				,   getattr(new_processor.feature_extractor				,   lowerCAmelCase				)				)
																self.assertDictEqual(new_processor.tokenizer.get_vocab()				,   processor.tokenizer.get_vocab()				)





				def 					a__(		self      :			Dict				)->					Tuple:

										"""simple docstring"""




										CustomFeatureExtractor.register_for_auto_class()
										CustomTokenizer.register_for_auto_class()
										CustomProcessor.register_for_auto_class()

										UpperCAmelCase					    =     CustomFeatureExtractor.from_pretrained(lowerCAmelCase				)

										with tempfile.TemporaryDirectory() as tmp_dir:
																UpperCAmelCase					    =     os.path.join(lowerCAmelCase				,   '''vocab.txt'''				)
																with open(lowerCAmelCase				,   '''w'''				,   encoding='''utf-8'''				) as vocab_writer:
																						vocab_writer.write(''''''.join([x + '''\n''' for x in self.vocab_tokens]				)				)
																UpperCAmelCase					    =     CustomTokenizer(lowerCAmelCase				)

										UpperCAmelCase					    =     CustomProcessor(lowerCAmelCase				,   lowerCAmelCase				)

										with tempfile.TemporaryDirectory() as tmp_dir:
																create_repo(F"""{USER}/test-dynamic-processor"""				,   token=self._token				)
																UpperCAmelCase					    =     Repository(lowerCAmelCase				,   clone_from=F"""{USER}/test-dynamic-processor"""				,   token=self._token				)
																processor.save_pretrained(lowerCAmelCase				)

																# This has added the proper auto_map field to the feature extractor config
																self.assertDictEqual(
																    processor.feature_extractor.auto_map				,   {
																        '''AutoFeatureExtractor''': '''custom_feature_extraction.CustomFeatureExtractor''',
																        '''AutoProcessor''': '''custom_processing.CustomProcessor''',
																    }				,   )

																# This has added the proper auto_map field to the tokenizer config
																with open(os.path.join(lowerCAmelCase				,   '''tokenizer_config.json'''				)				) as f:
																						UpperCAmelCase					    =     json.load(lowerCAmelCase				)
																self.assertDictEqual(
																    tokenizer_config['''auto_map''']				,   {
																        '''AutoTokenizer''': ['''custom_tokenization.CustomTokenizer''', None],
																        '''AutoProcessor''': '''custom_processing.CustomProcessor''',
																    }				,   )

																# The code has been copied from fixtures
																self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase				,   '''custom_feature_extraction.py'''				)				)				)
																self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase				,   '''custom_tokenization.py'''				)				)				)
																self.assertTrue(os.path.isfile(os.path.join(lowerCAmelCase				,   '''custom_processing.py'''				)				)				)

																repo.push_to_hub()

										UpperCAmelCase					    =     AutoProcessor.from_pretrained(F"""{USER}/test-dynamic-processor"""				,   trust_remote_code=lowerCAmelCase				)
										# Can't make an isinstance check because the new_processor is from the CustomProcessor class of a dynamic module
										self.assertEqual(new_processor.__class__.__name__				,   '''CustomProcessor'''				)









'''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 	UpperCamelCase__(       nn.Module	):
				__magic_name__    :				int
				__magic_name__    :				int
				__magic_name__    :				float  		=	0.0
				__magic_name__    :				int  		=	1
				__magic_name__    :				int  		=	1
				__magic_name__    :				bool  		=	True
				__magic_name__    :				bool  		=	False
				__magic_name__    :				bool  		=	False
				__magic_name__    :				bool  		=	False
				__magic_name__    :				jnp.dtype  		=	jnp.floataa




				def 					a__(		self      :			str				)->					Dict:

										"""simple docstring"""




										UpperCAmelCase					    =     []
										UpperCAmelCase					    =     []

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

																UpperCAmelCase					    =     FlaxResnetBlockaD(
																    in_channels=lowerCAmelCase				,   out_channels=self.out_channels				,   dropout_prob=self.dropout				,   dtype=self.dtype				,   )
																resnets.append(lowerCAmelCase				)

																UpperCAmelCase					    =     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(lowerCAmelCase				)

										UpperCAmelCase					    =     resnets
										UpperCAmelCase					    =     attentions

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





				def __call__(		self      :			List[Any]				,   lowerCAmelCase      :			Any				,   lowerCAmelCase      :			int				,   lowerCAmelCase      :			Dict				,   lowerCAmelCase      :			List[str]=True				)->					Optional[int]:

										"""simple docstring"""




										UpperCAmelCase					    =     ()

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

										if self.add_downsample:
																UpperCAmelCase					    =     self.downsamplers_a(lowerCAmelCase				)
																output_states += (hidden_states,)

										return hidden_states, output_states







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




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

										"""simple docstring"""




										UpperCAmelCase					    =     []

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

																UpperCAmelCase					    =     FlaxResnetBlockaD(
																    in_channels=lowerCAmelCase				,   out_channels=self.out_channels				,   dropout_prob=self.dropout				,   dtype=self.dtype				,   )
																resnets.append(lowerCAmelCase				)
										UpperCAmelCase					    =     resnets

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





				def __call__(		self      :			List[str]				,   lowerCAmelCase      :			Tuple				,   lowerCAmelCase      :			Dict				,   lowerCAmelCase      :			List[str]=True				)->					Optional[int]:

										"""simple docstring"""




										UpperCAmelCase					    =     ()

										for resnet in self.resnets:
																UpperCAmelCase					    =     resnet(lowerCAmelCase				,   lowerCAmelCase				,   deterministic=lowerCAmelCase				)
																output_states += (hidden_states,)

										if self.add_downsample:
																UpperCAmelCase					    =     self.downsamplers_a(lowerCAmelCase				)
																output_states += (hidden_states,)

										return hidden_states, output_states







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




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

										"""simple docstring"""




										UpperCAmelCase					    =     []
										UpperCAmelCase					    =     []

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

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

																UpperCAmelCase					    =     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(lowerCAmelCase				)

										UpperCAmelCase					    =     resnets
										UpperCAmelCase					    =     attentions

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





				def __call__(		self      :			Union[str, Any]				,   lowerCAmelCase      :			List[str]				,   lowerCAmelCase      :			Optional[int]				,   lowerCAmelCase      :			List[str]				,   lowerCAmelCase      :			Any				,   lowerCAmelCase      :			Union[str, Any]=True				)->					Optional[int]:

										"""simple docstring"""




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

																UpperCAmelCase					    =     resnet(lowerCAmelCase				,   lowerCAmelCase				,   deterministic=lowerCAmelCase				)
																UpperCAmelCase					    =     attn(lowerCAmelCase				,   lowerCAmelCase				,   deterministic=lowerCAmelCase				)

										if self.add_upsample:
																UpperCAmelCase					    =     self.upsamplers_a(lowerCAmelCase				)

										return hidden_states







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




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

										"""simple docstring"""




										UpperCAmelCase					    =     []

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

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

										UpperCAmelCase					    =     resnets

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





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

										"""simple docstring"""




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

																UpperCAmelCase					    =     resnet(lowerCAmelCase				,   lowerCAmelCase				,   deterministic=lowerCAmelCase				)

										if self.add_upsample:
																UpperCAmelCase					    =     self.upsamplers_a(lowerCAmelCase				)

										return hidden_states







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




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

										"""simple docstring"""




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

										UpperCAmelCase					    =     []

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

																UpperCAmelCase					    =     FlaxResnetBlockaD(
																    in_channels=self.in_channels				,   out_channels=self.in_channels				,   dropout_prob=self.dropout				,   dtype=self.dtype				,   )
																resnets.append(lowerCAmelCase				)

										UpperCAmelCase					    =     resnets
										UpperCAmelCase					    =     attentions





				def __call__(		self      :			Dict				,   lowerCAmelCase      :			Optional[Any]				,   lowerCAmelCase      :			Tuple				,   lowerCAmelCase      :			Dict				,   lowerCAmelCase      :			Any=True				)->					Union[str, Any]:

										"""simple docstring"""




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

										return hidden_states





