




import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD







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






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : str					  =      []
							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_   : str					  =      FlaxResnetBlockaD(
										    in_channels=snake_case__      ,out_channels=self.out_channels      ,dropout_prob=self.dropout      ,dtype=self.dtype      ,)
										resnets.append(snake_case__ )

										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(snake_case__ )

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

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






				def __call__(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__=True ):
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      ()

							for resnet, attn in zip(self.resnets      ,self.attentions ):
										SCREAMING_SNAKE_CASE_   : Tuple					  =      resnet(snake_case__      ,snake_case__      ,deterministic=snake_case__ )
										SCREAMING_SNAKE_CASE_   : str					  =      attn(snake_case__      ,snake_case__      ,deterministic=snake_case__ )
										output_states += (hidden_states,)

							if self.add_downsample:
										SCREAMING_SNAKE_CASE_   : List[Any]					  =      self.downsamplers_a(snake_case__ )
										output_states += (hidden_states,)

							return hidden_states, output_states







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






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      []

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

										SCREAMING_SNAKE_CASE_   : List[str]					  =      FlaxResnetBlockaD(
										    in_channels=snake_case__      ,out_channels=self.out_channels      ,dropout_prob=self.dropout      ,dtype=self.dtype      ,)
										resnets.append(snake_case__ )
							SCREAMING_SNAKE_CASE_   : Dict					  =      resnets

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






				def __call__(			self      ,snake_case__      ,snake_case__      ,snake_case__=True ):
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      ()

							for resnet in self.resnets:
										SCREAMING_SNAKE_CASE_   : List[Any]					  =      resnet(snake_case__      ,snake_case__      ,deterministic=snake_case__ )
										output_states += (hidden_states,)

							if self.add_downsample:
										SCREAMING_SNAKE_CASE_   : List[str]					  =      self.downsamplers_a(snake_case__ )
										output_states += (hidden_states,)

							return hidden_states, output_states







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






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Any					  =      []
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      []

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

										SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      FlaxResnetBlockaD(
										    in_channels=resnet_in_channels + res_skip_channels      ,out_channels=self.out_channels      ,dropout_prob=self.dropout      ,dtype=self.dtype      ,)
										resnets.append(snake_case__ )

										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(snake_case__ )

							SCREAMING_SNAKE_CASE_   : str					  =      resnets
							SCREAMING_SNAKE_CASE_   : Any					  =      attentions

							if self.add_upsample:
										SCREAMING_SNAKE_CASE_   : List[str]					  =      FlaxUpsampleaD(self.out_channels      ,dtype=self.dtype )






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

										SCREAMING_SNAKE_CASE_   : Optional[int]					  =      resnet(snake_case__      ,snake_case__      ,deterministic=snake_case__ )
										SCREAMING_SNAKE_CASE_   : Dict					  =      attn(snake_case__      ,snake_case__      ,deterministic=snake_case__ )

							if self.add_upsample:
										SCREAMING_SNAKE_CASE_   : Tuple					  =      self.upsamplers_a(snake_case__ )

							return hidden_states







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






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : int					  =      []

							for i in range(self.num_layers ):
										SCREAMING_SNAKE_CASE_   : int					  =      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_   : List[str]					  =      FlaxResnetBlockaD(
										    in_channels=resnet_in_channels + res_skip_channels      ,out_channels=self.out_channels      ,dropout_prob=self.dropout      ,dtype=self.dtype      ,)
										resnets.append(snake_case__ )

							SCREAMING_SNAKE_CASE_   : List[Any]					  =      resnets

							if self.add_upsample:
										SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      FlaxUpsampleaD(self.out_channels      ,dtype=self.dtype )






				def __call__(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__=True ):
							for resnet in self.resnets:
										# pop res hidden states
										SCREAMING_SNAKE_CASE_   : List[str]					  =      res_hidden_states_tuple[-1]
										SCREAMING_SNAKE_CASE_   : str					  =      res_hidden_states_tuple[:-1]
										SCREAMING_SNAKE_CASE_   : List[Any]					  =      jnp.concatenate((hidden_states, res_hidden_states)      ,axis=-1 )

										SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      resnet(snake_case__      ,snake_case__      ,deterministic=snake_case__ )

							if self.add_upsample:
										SCREAMING_SNAKE_CASE_   : str					  =      self.upsamplers_a(snake_case__ )

							return hidden_states







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






				def 		snake_case	(			self ):
							# there is always at least one resnet
							SCREAMING_SNAKE_CASE_   : Any					  =      [
							    FlaxResnetBlockaD(
							        in_channels=self.in_channels      ,out_channels=self.in_channels      ,dropout_prob=self.dropout      ,dtype=self.dtype      ,)
							]

							SCREAMING_SNAKE_CASE_   : Tuple					  =      []

							for _ in range(self.num_layers ):
										SCREAMING_SNAKE_CASE_   : Optional[int]					  =      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(snake_case__ )

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

							SCREAMING_SNAKE_CASE_   : Tuple					  =      resnets
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      attentions






				def __call__(			self      ,snake_case__      ,snake_case__      ,snake_case__      ,snake_case__=True ):
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      self.resnets[0](snake_case__      ,snake_case__ )
							for attn, resnet in zip(self.attentions      ,self.resnets[1:] ):
										SCREAMING_SNAKE_CASE_   : Tuple					  =      attn(snake_case__      ,snake_case__      ,deterministic=snake_case__ )
										SCREAMING_SNAKE_CASE_   : Optional[int]					  =      resnet(snake_case__      ,snake_case__      ,deterministic=snake_case__ )

							return hidden_states









import os
import sys
import unittest


UpperCamelCase__ :   Optional[Any]						=							os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, '''utils'''))

import check_dummies  # noqa: E402
from check_dummies import create_dummy_files, create_dummy_object, find_backend, read_init  # noqa: E402


# Align TRANSFORMERS_PATH in check_dummies with the current path
UpperCamelCase__ :   Tuple						=							os.path.join(git_repo_path, '''src''', '''transformers''')

UpperCamelCase__ :   List[Any]						=							'''
{0} = None
'''

UpperCamelCase__ :   str						=							'''
class {0}(metaclass=DummyObject):
    _backends = {1}

    def __init__(self, *args, **kwargs):
        requires_backends(self, {1})
'''


UpperCamelCase__ :   List[Any]						=							'''
def {0}(*args, **kwargs):
    requires_backends({0}, {1})
'''







class 				lowerCAmelCase_       (					unittest.TestCase							):






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : int					  =      find_backend('    _import_structure["models.albert"].append("AlbertTokenizerFast")' )
							self.assertIsNone(snake_case__ )

							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      find_backend('    if not is_tokenizers_available():' )
							self.assertEqual(snake_case__      ,'tokenizers' )

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      find_backend('    if not is_tensorflow_text_available():' )
							self.assertEqual(snake_case__      ,'tensorflow_text' )

							SCREAMING_SNAKE_CASE_   : int					  =      find_backend('    if not (is_sentencepiece_available() and is_tokenizers_available()):' )
							self.assertEqual(snake_case__      ,'sentencepiece_and_tokenizers' )

							SCREAMING_SNAKE_CASE_   : str					  =      find_backend(
							    '    if not (is_sentencepiece_available() and is_tensorflow_text_available()):' )
							self.assertEqual(snake_case__      ,'sentencepiece_and_tensorflow_text' )

							SCREAMING_SNAKE_CASE_   : Tuple					  =      find_backend(
							    '    if not (is_sentencepiece_available() and is_tokenizers_available() and is_vision_available()):' )
							self.assertEqual(snake_case__      ,'sentencepiece_and_tokenizers_and_vision' )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : str					  =      read_init()
							# We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects
							self.assertIn('torch'      ,snake_case__ )
							self.assertIn('tensorflow_text'      ,snake_case__ )
							self.assertIn('sentencepiece_and_tokenizers'      ,snake_case__ )

							# Likewise, we can't assert on the exact content of a key
							self.assertIn('BertModel'      ,objects['torch'] )
							self.assertIn('TFBertModel'      ,objects['tf'] )
							self.assertIn('FlaxBertModel'      ,objects['flax'] )
							self.assertIn('BertModel'      ,objects['torch'] )
							self.assertIn('TFBertTokenizer'      ,objects['tensorflow_text'] )
							self.assertIn('convert_slow_tokenizer'      ,objects['sentencepiece_and_tokenizers'] )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : str					  =      create_dummy_object('CONSTANT'      ,'\'torch\'' )
							self.assertEqual(snake_case__      ,'\nCONSTANT = None\n' )

							SCREAMING_SNAKE_CASE_   : Dict					  =      create_dummy_object('function'      ,'\'torch\'' )
							self.assertEqual(
							    snake_case__      ,'\ndef function(*args, **kwargs):\n    requires_backends(function, \'torch\')\n' )

							SCREAMING_SNAKE_CASE_   : List[str]					  =      '\nclass FakeClass(metaclass=DummyObject):\n    _backends = \'torch\'\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, \'torch\')\n'
							SCREAMING_SNAKE_CASE_   : Dict					  =      create_dummy_object('FakeClass'      ,'\'torch\'' )
							self.assertEqual(snake_case__      ,snake_case__ )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      '# This file is autogenerated by the command `make fix-copies`, do not edit.\nfrom ..utils import DummyObject, requires_backends\n\n\nCONSTANT = None\n\n\ndef function(*args, **kwargs):\n    requires_backends(function, ["torch"])\n\n\nclass FakeClass(metaclass=DummyObject):\n    _backends = ["torch"]\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, ["torch"])\n'
							SCREAMING_SNAKE_CASE_   : Dict					  =      create_dummy_files({'torch': ['CONSTANT', 'function', 'FakeClass']} )
							self.assertEqual(dummy_files['torch']      ,snake_case__ )



