


'''simple docstring'''





import inspect
import unittest

from transformers import RegNetConfig, is_flax_available
from transformers.testing_utils import require_flax, slow
from transformers.utils import cached_property, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
    import jax
    import jax.numpy as jnp

    from transformers.models.regnet.modeling_flax_regnet import FlaxRegNetForImageClassification, FlaxRegNetModel

if is_vision_available():
    from PIL import Image

    from transformers import AutoImageProcessor




class 				UpperCAmelCase			(			unittest.TestCase						):


 def __init__(      self							:	str			,       __lowerCamelCase							:	Tuple			,       __lowerCamelCase							:	Union[str, Any]=3			,       __lowerCamelCase							:	Union[str, Any]=3_2			,       __lowerCamelCase							:	Union[str, Any]=3			,       __lowerCamelCase							:	Optional[int]=1_0			,       __lowerCamelCase							:	Tuple=[1_0, 2_0, 3_0, 4_0]			,       __lowerCamelCase							:	str=[1, 1, 2, 1]			,       __lowerCamelCase							:	List[Any]=True			,       __lowerCamelCase							:	List[Any]=True			,       __lowerCamelCase							:	List[str]="relu"			,       __lowerCamelCase							:	Optional[Any]=3			,       __lowerCamelCase							:	int=None			,       ):
    UpperCAmelCase__	:Optional[Any]      =      parent
    UpperCAmelCase__	:str      =      batch_size
    UpperCAmelCase__	:str      =      image_size
    UpperCAmelCase__	:Dict      =      num_channels
    UpperCAmelCase__	:int      =      embeddings_size
    UpperCAmelCase__	:Optional[int]      =      hidden_sizes
    UpperCAmelCase__	:Optional[Any]      =      depths
    UpperCAmelCase__	:str      =      is_training
    UpperCAmelCase__	:Tuple      =      use_labels
    UpperCAmelCase__	:Union[str, Any]      =      hidden_act
    UpperCAmelCase__	:int      =      num_labels
    UpperCAmelCase__	:Tuple      =      scope
    UpperCAmelCase__	:int      =      len(__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Union[str, Any] ):
    UpperCAmelCase__	:List[Any]      =      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

    UpperCAmelCase__	:int      =      self.get_config()

    return config, pixel_values


 def        __SCREAMING_SNAKE_CASE					(      self							:	int ):
    return RegNetConfig(
        num_channels=self.num_channels			,       embeddings_size=self.embeddings_size			,       hidden_sizes=self.hidden_sizes			,       depths=self.depths			,       hidden_act=self.hidden_act			,       num_labels=self.num_labels			,       image_size=self.image_size			,       )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Dict			,       __lowerCamelCase							:	Optional[Any]			,       __lowerCamelCase							:	Union[str, Any] ):
    UpperCAmelCase__	:List[Any]      =      FlaxRegNetModel(config=__lowerCamelCase )
    UpperCAmelCase__	:Optional[Any]      =      model(__lowerCamelCase )

    # Output shape (b, c, h, w)
    self.parent.assertEqual(
        result.last_hidden_state.shape			,       (self.batch_size, self.hidden_sizes[-1], self.image_size // 3_2, self.image_size // 3_2)			,       )


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str]			,       __lowerCamelCase							:	Optional[Any]			,       __lowerCamelCase							:	List[Any] ):
    UpperCAmelCase__	:Optional[Any]      =      self.num_labels
    UpperCAmelCase__	:Any      =      FlaxRegNetForImageClassification(config=__lowerCamelCase )
    UpperCAmelCase__	:Union[str, Any]      =      model(__lowerCamelCase )
    self.parent.assertEqual(result.logits.shape			,       (self.batch_size, self.num_labels) )





 def        __SCREAMING_SNAKE_CASE					(      self							:	Dict ):
    UpperCAmelCase__	:List[str]      =      self.prepare_config_and_inputs()
    UpperCAmelCase__   ,						UpperCAmelCase__	:Optional[int]      =      config_and_inputs
    UpperCAmelCase__	:List[Any]      =      {'''pixel_values''': pixel_values}
    return config, inputs_dict




@require_flax
class 				UpperCAmelCase			(			_snake_case						, unittest.TestCase						):
 UpperCAmelCase				=					(FlaxRegNetModel, FlaxRegNetForImageClassification) if is_flax_available() else ()

 UpperCAmelCase				=					False
 UpperCAmelCase				=					False
 UpperCAmelCase				=					False


 def        __SCREAMING_SNAKE_CASE					(      self							:	Union[str, Any] ):
    UpperCAmelCase__	:Union[str, Any]      =      FlaxRegNetModelTester(self )
    UpperCAmelCase__	:Any      =      ConfigTester(self			,       config_class=__lowerCamelCase			,       has_text_modality=__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[int] ):
    self.create_and_test_config_common_properties()
    self.config_tester.create_and_test_config_to_json_string()
    self.config_tester.create_and_test_config_to_json_file()
    self.config_tester.create_and_test_config_from_and_save_pretrained()
    self.config_tester.create_and_test_config_with_num_labels()
    self.config_tester.check_config_can_be_init_without_params()
    self.config_tester.check_config_arguments_init()


 def        __SCREAMING_SNAKE_CASE					(      self							:	Tuple ):
    return


 def        __SCREAMING_SNAKE_CASE					(      self							:	Any ):
    UpperCAmelCase__	:List[str]      =      self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*__lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[Any] ):
    UpperCAmelCase__	:Optional[int]      =      self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_image_classification(*__lowerCamelCase )


 @unittest.skip(reason='''RegNet does not use inputs_embeds''' )
 def        __SCREAMING_SNAKE_CASE					(      self							:	str ):
    pass


 @unittest.skip(reason='''RegNet does not support input and output embeddings''' )
 def        __SCREAMING_SNAKE_CASE					(      self							:	Any ):
    pass


 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str] ):
    UpperCAmelCase__   ,						UpperCAmelCase__	:Dict      =      self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
       UpperCAmelCase__	:Optional[int]      =      model_class(__lowerCamelCase )
       UpperCAmelCase__	:str      =      inspect.signature(model.__call__ )
       # signature.parameters is an OrderedDict => so arg_names order is deterministic
       UpperCAmelCase__	:Tuple      =      [*signature.parameters.keys()]

       UpperCAmelCase__	:Optional[Any]      =      ['''pixel_values''']
       self.assertListEqual(arg_names[:1]			,       __lowerCamelCase )


 def        __SCREAMING_SNAKE_CASE					(      self							:	Tuple ):
    def check_hidden_states_output(__lowerCamelCase							:	Tuple			,       __lowerCamelCase							:	List[Any]			,       __lowerCamelCase							:	Optional[Any] ):
       UpperCAmelCase__	:Optional[int]      =      model_class(__lowerCamelCase )

       UpperCAmelCase__	:Dict      =      model(**self._prepare_for_class(__lowerCamelCase			,       __lowerCamelCase ) )

       UpperCAmelCase__	:str      =      outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

       UpperCAmelCase__	:Tuple      =      self.model_tester.num_stages
       self.assertEqual(len(__lowerCamelCase )			,       expected_num_stages + 1 )

    UpperCAmelCase__   ,						UpperCAmelCase__	:Union[str, Any]      =      self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
       UpperCAmelCase__	:Optional[int]      =      True
       check_hidden_states_output(__lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase )

       # check that output_hidden_states also work using config
       del inputs_dict["output_hidden_states"]
       UpperCAmelCase__	:str      =      True

       check_hidden_states_output(__lowerCamelCase			,       __lowerCamelCase			,       __lowerCamelCase )





 def        __SCREAMING_SNAKE_CASE					(      self							:	Dict ):
    UpperCAmelCase__   ,						UpperCAmelCase__	:Optional[int]      =      self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
       with self.subTest(model_class.__name__ ):
          UpperCAmelCase__	:Optional[Any]      =      self._prepare_for_class(__lowerCamelCase			,       __lowerCamelCase )
          UpperCAmelCase__	:int      =      model_class(__lowerCamelCase )

          @jax.jit
          def model_jitted(__lowerCamelCase							:	List[str]			,       **__lowerCamelCase							:	Dict ):
             return model(pixel_values=__lowerCamelCase			,       **__lowerCamelCase )

          with self.subTest('''JIT Enabled''' ):
             UpperCAmelCase__	:Union[str, Any]      =      model_jitted(**__lowerCamelCase ).to_tuple()

          with self.subTest('''JIT Disabled''' ):
             with jax.disable_jit():
                UpperCAmelCase__	:Optional[int]      =      model_jitted(**__lowerCamelCase ).to_tuple()

          self.assertEqual(len(__lowerCamelCase )			,       len(__lowerCamelCase ) )
          for jitted_output, output in zip(__lowerCamelCase			,       __lowerCamelCase ):
             self.assertEqual(jitted_output.shape			,       output.shape )


def        a__     (					):
   UpperCAmelCase__	:List[Any]      =      Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''						)
   return image




@require_flax
class 				UpperCAmelCase			(			unittest.TestCase						):


 @cached_property
 def        __SCREAMING_SNAKE_CASE					(      self							:	List[str] ):
    return AutoImageProcessor.from_pretrained('''facebook/regnet-y-040''' ) if is_vision_available() else None





 @slow
 def        __SCREAMING_SNAKE_CASE					(      self							:	Optional[Any] ):
    UpperCAmelCase__	:int      =      FlaxRegNetForImageClassification.from_pretrained('''facebook/regnet-y-040''' )

    UpperCAmelCase__	:Optional[Any]      =      self.default_image_processor
    UpperCAmelCase__	:Tuple      =      prepare_img()
    UpperCAmelCase__	:Union[str, Any]      =      image_processor(images=__lowerCamelCase			,       return_tensors='''np''' )

    UpperCAmelCase__	:Union[str, Any]      =      model(**__lowerCamelCase )

    # verify the logits
    UpperCAmelCase__	:str      =      (1, 1_0_0_0)
    self.assertEqual(outputs.logits.shape			,       __lowerCamelCase )

    UpperCAmelCase__	:Dict      =      jnp.array([-0.41_80, -1.50_51, -3.48_36] )

    self.assertTrue(jnp.allclose(outputs.logits[0, :3]			,       __lowerCamelCase			,       atol=1e-4 ) )







'''simple docstring'''





from typing import Optional

from torch import nn

from .transformer_ad import TransformeraDModel, TransformeraDModelOutput




class 				UpperCAmelCase			(			nn.Module						):


 def __init__(      self							:	int			,       __lowerCamelCase							:	int = 1_6			,       __lowerCamelCase							:	int = 8_8			,       __lowerCamelCase							:	Optional[int] = None			,       __lowerCamelCase							:	int = 1			,       __lowerCamelCase							:	float = 0.0			,       __lowerCamelCase							:	int = 3_2			,       __lowerCamelCase							:	Optional[int] = None			,       __lowerCamelCase							:	bool = False			,       __lowerCamelCase							:	Optional[int] = None			,       __lowerCamelCase							:	Optional[int] = None			,       __lowerCamelCase							:	str = "geglu"			,       __lowerCamelCase							:	Optional[int] = None			,       ):
    super().__init__()
    UpperCAmelCase__	:Optional[int]      =      nn.ModuleList(
        [
            TransformeraDModel(
                num_attention_heads=__lowerCamelCase			,       attention_head_dim=__lowerCamelCase			,       in_channels=__lowerCamelCase			,       num_layers=__lowerCamelCase			,       dropout=__lowerCamelCase			,       norm_num_groups=__lowerCamelCase			,       cross_attention_dim=__lowerCamelCase			,       attention_bias=__lowerCamelCase			,       sample_size=__lowerCamelCase			,       num_vector_embeds=__lowerCamelCase			,       activation_fn=__lowerCamelCase			,       num_embeds_ada_norm=__lowerCamelCase			,       )
            for _ in range(2 )
        ] )

    # Variables that can be set by a pipeline:

    # The ratio of transformer1 to transformer2's output states to be combined during inference
    UpperCAmelCase__	:Tuple      =      0.5

    # The shape of `encoder_hidden_states` is expected to be
    # `(batch_size, condition_lengths[0]+condition_lengths[1], num_features)`
    UpperCAmelCase__	:int      =      [7_7, 2_5_7]

    # Which transformer to use to encode which condition.
    # E.g. `(1, 0)` means that we'll use `transformers[1](conditions[0])` and `transformers[0](conditions[1])`
    UpperCAmelCase__	:Any      =      [1, 0]





 def        __SCREAMING_SNAKE_CASE					(      self							:	str			,       __lowerCamelCase							:	Optional[int]			,       __lowerCamelCase							:	Any			,       __lowerCamelCase							:	int=None			,       __lowerCamelCase							:	Dict=None			,       __lowerCamelCase							:	Dict=None			,       __lowerCamelCase							:	bool = True			,       ):
    UpperCAmelCase__	:Optional[int]      =      hidden_states

    UpperCAmelCase__	:Dict      =      []
    UpperCAmelCase__	:Tuple      =      0
    # attention_mask is not used yet
    for i in range(2 ):
       # for each of the two transformers, pass the corresponding condition tokens
       UpperCAmelCase__	:List[Any]      =      encoder_hidden_states[:, tokens_start : tokens_start + self.condition_lengths[i]]
       UpperCAmelCase__	:Dict      =      self.transformer_index_for_condition[i]
       UpperCAmelCase__	:Tuple      =      self.transformers[transformer_index](
           __lowerCamelCase			,       encoder_hidden_states=__lowerCamelCase			,       timestep=__lowerCamelCase			,       cross_attention_kwargs=__lowerCamelCase			,       return_dict=__lowerCamelCase			,       )[0]
       encoded_states.append(encoded_state - input_states )
       tokens_start += self.condition_lengths[i]

    UpperCAmelCase__	:Optional[int]      =      encoded_states[0] * self.mix_ratio + encoded_states[1] * (1 - self.mix_ratio)
    UpperCAmelCase__	:Optional[Any]      =      output_states + input_states

    if not return_dict:
       return (output_states,)

    return TransformeraDModelOutput(sample=__lowerCamelCase )



