


import copy
import json
import os
import tempfile

from transformers import is_torch_available

from .test_configuration_utils import config_common_kwargs




class A  (   UpperCamelCase_     ):






      def __init__(					self						:     List[Any]			,   lowercase_						:     str			,   lowercase_						:     str=None			,   lowercase_						:     int=True			,   lowercase_						:     List[Any]=None			,   **lowercase_						:     int )					->						int:
             """simple docstring"""




             _lowerCamelCase       :		Optional[Any]		=parent
             _lowerCamelCase       :		Optional[int]		=config_class
             _lowerCamelCase       :		Optional[Any]		=has_text_modality
             _lowerCamelCase       :		Optional[Any]		=kwargs
             _lowerCamelCase       :		Any		=common_properties






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




             _lowerCamelCase       :		Any		=self.config_class(**self.inputs_dict )
             _lowerCamelCase       :		List[str]		=(
                 ['hidden_size', 'num_attention_heads', 'num_hidden_layers']
                 if self.common_properties is None
                 else self.common_properties
             )

             # Add common fields for text models
             if self.has_text_modality:
                    common_properties.extend(['vocab_size'] )

             # Test that config has the common properties as getters
             for prop in common_properties:
                    self.parent.assertTrue(hasattr(lowercase_			,   lowercase_ )			,   msg=F'''`{prop}` does not exist''' )

             # Test that config has the common properties as setter
             for idx, name in enumerate(lowercase_ ):
                    try:
                           setattr(lowercase_			,   lowercase_			,   lowercase_ )
                           self.parent.assertEqual(
                               getattr(lowercase_			,   lowercase_ )			,   lowercase_			,   msg=F'''`{name} value {idx} expected, but was {getattr(lowercase_			,   lowercase_ )}''' )
                    except NotImplementedError:
                           # Some models might not be able to implement setters for common_properties
                           # In that case, a NotImplementedError is raised
                           pass

        # Test if config class can be called with Config(prop_name=..)
             for idx, name in enumerate(lowercase_ ):
                    try:
                           _lowerCamelCase       :		Optional[int]		=self.config_class(**{name: idx} )
                           self.parent.assertEqual(
                               getattr(lowercase_			,   lowercase_ )			,   lowercase_			,   msg=F'''`{name} value {idx} expected, but was {getattr(lowercase_			,   lowercase_ )}''' )
                    except NotImplementedError:
                           # Some models might not be able to implement setters for common_properties
                           # In that case, a NotImplementedError is raised
                           pass






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




             _lowerCamelCase       :		int		=self.config_class(**self.inputs_dict )
             _lowerCamelCase       :		Any		=json.loads(config.to_json_string() )
             for key, value in self.inputs_dict.items():
                    self.parent.assertEqual(obj[key]			,   lowercase_ )






      def        lowerCamelCase       (					self						:     Dict )					->						Optional[int]:
             """simple docstring"""




             _lowerCamelCase       :		List[str]		=self.config_class(**self.inputs_dict )

             with tempfile.TemporaryDirectory() as tmpdirname:
                    _lowerCamelCase       :		List[str]		=os.path.join(lowercase_			,   'config.json' )
                    config_first.to_json_file(lowercase_ )
                    _lowerCamelCase       :		Optional[Any]		=self.config_class.from_json_file(lowercase_ )

             self.parent.assertEqual(config_second.to_dict()			,   config_first.to_dict() )






      def        lowerCamelCase       (					self						:     Any )					->						int:
             """simple docstring"""




             _lowerCamelCase       :		int		=self.config_class(**self.inputs_dict )

             with tempfile.TemporaryDirectory() as tmpdirname:
                    config_first.save_pretrained(lowercase_ )
                    _lowerCamelCase       :		int		=self.config_class.from_pretrained(lowercase_ )

             self.parent.assertEqual(config_second.to_dict()			,   config_first.to_dict() )






      def        lowerCamelCase       (					self						:     str )					->						str:
             """simple docstring"""




             _lowerCamelCase       :		Optional[Any]		=self.config_class(**self.inputs_dict )

             _lowerCamelCase       :		List[Any]		='test'
             with tempfile.TemporaryDirectory() as tmpdirname:
                    _lowerCamelCase       :		Any		=os.path.join(lowercase_			,   lowercase_ )
                    config_first.save_pretrained(lowercase_ )
                    _lowerCamelCase       :		Union[str, Any]		=self.config_class.from_pretrained(lowercase_			,   subfolder=lowercase_ )

             self.parent.assertEqual(config_second.to_dict()			,   config_first.to_dict() )






      def        lowerCamelCase       (					self						:     List[Any] )					->						int:
             """simple docstring"""




             _lowerCamelCase       :		Optional[Any]		=self.config_class(**self.inputs_dict			,   num_labels=5 )
             self.parent.assertEqual(len(config.idalabel )			,   5 )
             self.parent.assertEqual(len(config.labelaid )			,   5 )

             _lowerCamelCase       :		Dict		=3
             self.parent.assertEqual(len(config.idalabel )			,   3 )
             self.parent.assertEqual(len(config.labelaid )			,   3 )






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




             if self.config_class.is_composition:
                    return
             _lowerCamelCase       :		Any		=self.config_class()
             self.parent.assertIsNotNone(lowercase_ )






      def        lowerCamelCase       (					self						:     Optional[int] )					->						Tuple:
             """simple docstring"""




             _lowerCamelCase       :		Optional[int]		=copy.deepcopy(lowercase_ )
             _lowerCamelCase       :		Union[str, Any]		=self.config_class(**lowercase_ )
             _lowerCamelCase       :		Union[str, Any]		=[]
             for key, value in config_common_kwargs.items():
                    if key == "torch_dtype":
                           if not is_torch_available():
                                  continue
                           else:
                                  import torch

                                  if config.torch_dtype != torch.floataa:
                                         wrong_values.append(('torch_dtype', config.torch_dtype, torch.floataa) )
                    elif getattr(lowercase_			,   lowercase_ ) != value:
                           wrong_values.append((key, getattr(lowercase_			,   lowercase_ ), value) )

             if len(lowercase_ ) > 0:
                    _lowerCamelCase       :		Dict		='\n'.join([F'''- {v[0]}: got {v[1]} instead of {v[2]}''' for v in wrong_values] )
                    raise ValueError(F'''The following keys were not properly set in the config:\n{errors}''' )






      def        lowerCamelCase       (					self						:     List[Any] )					->						int:
             """simple docstring"""




             self.create_and_test_config_common_properties()
             self.create_and_test_config_to_json_string()
             self.create_and_test_config_to_json_file()
             self.create_and_test_config_from_and_save_pretrained()
             self.create_and_test_config_from_and_save_pretrained_subfolder()
             self.create_and_test_config_with_num_labels()
             self.check_config_can_be_init_without_params()
             self.check_config_arguments_init()






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






      def        lowerCamelCase       (					self						:     Any )					->						Any:
             """simple docstring"""




             _lowerCamelCase       :		str		=[]
             _lowerCamelCase       :		Dict		=[]

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

                    _lowerCamelCase       :		Any		=FlaxResnetBlockaD(
                        in_channels=lowercase_			,   out_channels=self.out_channels			,   dropout_prob=self.dropout			,   dtype=self.dtype			,   )
                    resnets.append(lowercase_ )

                    _lowerCamelCase       :		Dict		=FlaxTransformeraDModel(
                        in_channels=self.out_channels			,   n_heads=self.num_attention_heads			,   d_head=self.out_channels // self.num_attention_heads			,   depth=1			,   use_linear_projection=self.use_linear_projection			,   only_cross_attention=self.only_cross_attention			,   use_memory_efficient_attention=self.use_memory_efficient_attention			,   dtype=self.dtype			,   )
                    attentions.append(lowercase_ )

             _lowerCamelCase       :		Optional[int]		=resnets
             _lowerCamelCase       :		Dict		=attentions

             if self.add_downsample:
                    _lowerCamelCase       :		Union[str, Any]		=FlaxDownsampleaD(self.out_channels			,   dtype=self.dtype )






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




             _lowerCamelCase       :		Tuple		=()

             for resnet, attn in zip(self.resnets			,   self.attentions ):
                    _lowerCamelCase       :		Union[str, Any]		=resnet(lowercase_			,   lowercase_			,   deterministic=lowercase_ )
                    _lowerCamelCase       :		Union[str, Any]		=attn(lowercase_			,   lowercase_			,   deterministic=lowercase_ )
                    output_states += (hidden_states,)

             if self.add_downsample:
                    _lowerCamelCase       :		Optional[int]		=self.downsamplers_a(lowercase_ )
                    output_states += (hidden_states,)

             return hidden_states, output_states




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






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




             _lowerCamelCase       :		str		=[]

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

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

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






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




             _lowerCamelCase       :		Optional[int]		=()

             for resnet in self.resnets:
                    _lowerCamelCase       :		Tuple		=resnet(lowercase_			,   lowercase_			,   deterministic=lowercase_ )
                    output_states += (hidden_states,)

             if self.add_downsample:
                    _lowerCamelCase       :		Tuple		=self.downsamplers_a(lowercase_ )
                    output_states += (hidden_states,)

             return hidden_states, output_states




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






      def        lowerCamelCase       (					self						:     Dict )					->						Tuple:
             """simple docstring"""




             _lowerCamelCase       :		str		=[]
             _lowerCamelCase       :		List[str]		=[]

             for i in range(self.num_layers ):
                    _lowerCamelCase       :		List[str]		=self.in_channels if (i == self.num_layers - 1) else self.out_channels
                    _lowerCamelCase       :		Tuple		=self.prev_output_channel if i == 0 else self.out_channels

                    _lowerCamelCase       :		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(lowercase_ )

                    _lowerCamelCase       :		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(lowercase_ )

             _lowerCamelCase       :		int		=resnets
             _lowerCamelCase       :		Dict		=attentions

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






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




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

                    _lowerCamelCase       :		Optional[Any]		=resnet(lowercase_			,   lowercase_			,   deterministic=lowercase_ )
                    _lowerCamelCase       :		List[Any]		=attn(lowercase_			,   lowercase_			,   deterministic=lowercase_ )

             if self.add_upsample:
                    _lowerCamelCase       :		Optional[Any]		=self.upsamplers_a(lowercase_ )

             return hidden_states




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






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




             _lowerCamelCase       :		List[str]		=[]

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

                    _lowerCamelCase       :		str		=FlaxResnetBlockaD(
                        in_channels=resnet_in_channels + res_skip_channels			,   out_channels=self.out_channels			,   dropout_prob=self.dropout			,   dtype=self.dtype			,   )
                    resnets.append(lowercase_ )

             _lowerCamelCase       :		str		=resnets

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






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




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

                    _lowerCamelCase       :		Optional[Any]		=resnet(lowercase_			,   lowercase_			,   deterministic=lowercase_ )

             if self.add_upsample:
                    _lowerCamelCase       :		Union[str, Any]		=self.upsamplers_a(lowercase_ )

             return hidden_states




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






      def        lowerCamelCase       (					self						:     int )					->						Tuple:
             """simple docstring"""




             _lowerCamelCase       :		Optional[Any]		=[
                 FlaxResnetBlockaD(
                     in_channels=self.in_channels			,   out_channels=self.in_channels			,   dropout_prob=self.dropout			,   dtype=self.dtype			,   )
             ]

             _lowerCamelCase       :		Any		=[]

             for _ in range(self.num_layers ):
                    _lowerCamelCase       :		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(lowercase_ )

                    _lowerCamelCase       :		Tuple		=FlaxResnetBlockaD(
                        in_channels=self.in_channels			,   out_channels=self.in_channels			,   dropout_prob=self.dropout			,   dtype=self.dtype			,   )
                    resnets.append(lowercase_ )

             _lowerCamelCase       :		List[Any]		=resnets
             _lowerCamelCase       :		List[str]		=attentions






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




             _lowerCamelCase       :		Dict		=self.resnets[0](lowercase_			,   lowercase_ )
             for attn, resnet in zip(self.attentions			,   self.resnets[1:] ):
                    _lowerCamelCase       :		Tuple		=attn(lowercase_			,   lowercase_			,   deterministic=lowercase_ )
                    _lowerCamelCase       :		List[str]		=resnet(lowercase_			,   lowercase_			,   deterministic=lowercase_ )

             return hidden_states


