




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					):

    '''simple docstring'''






    def __init__(      self				:     Dict ,			UpperCAmelCase__				:     Optional[int] ,			UpperCAmelCase__				:     Tuple=3 ,			UpperCAmelCase__				:     List[Any]=32 ,			UpperCAmelCase__				:     Dict=3 ,			UpperCAmelCase__				:     Dict=10 ,			UpperCAmelCase__				:     int=[10, 20, 30, 40] ,			UpperCAmelCase__				:     Union[str, Any]=[1, 1, 2, 1] ,			UpperCAmelCase__				:     int=True ,			UpperCAmelCase__				:     Optional[Any]=True ,			UpperCAmelCase__				:     List[str]="relu" ,			UpperCAmelCase__				:     Dict=3 ,			UpperCAmelCase__				:     Any=None ,			)						->Tuple:





           '''simple docstring'''

           A__							=   parent
           A__							=   batch_size
           A__							=   image_size
           A__							=   num_channels
           A__							=   embeddings_size
           A__							=   hidden_sizes
           A__							=   depths
           A__							=   is_training
           A__							=   use_labels
           A__							=   hidden_act
           A__							=   num_labels
           A__							=   scope
           A__							=   len(_SCREAMING_SNAKE_CASE)


    def 	SCREAMING_SNAKE_CASE							(      self				:     str)						->Dict:





           '''simple docstring'''

           A__							=   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

           A__							=   self.get_config()

           return config, pixel_values


    def 	SCREAMING_SNAKE_CASE							(      self				:     Tuple)						->Union[str, Any]:





           '''simple docstring'''

           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				:     Optional[Any] ,			UpperCAmelCase__				:     List[Any] ,			UpperCAmelCase__				:     str)						->List[str]:





           '''simple docstring'''

           A__							=   FlaxRegNetModel(config=_SCREAMING_SNAKE_CASE)
           A__							=   model(_SCREAMING_SNAKE_CASE)

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


    def 	SCREAMING_SNAKE_CASE							(      self				:     Any ,			UpperCAmelCase__				:     int ,			UpperCAmelCase__				:     Optional[Any])						->Any:





           '''simple docstring'''

           A__							=   self.num_labels
           A__							=   FlaxRegNetForImageClassification(config=_SCREAMING_SNAKE_CASE)
           A__							=   model(_SCREAMING_SNAKE_CASE)
           self.parent.assertEqual(result.logits.shape ,			(self.batch_size, self.num_labels))





    def 	SCREAMING_SNAKE_CASE							(      self				:     Dict)						->List[Any]:





           '''simple docstring'''

           A__							=   self.prepare_config_and_inputs()
           A__							=   config_and_inputs
           A__							=   {"pixel_values": pixel_values}
           return config, inputs_dict


@require_flax
class     UpperCamelCase_   (			SCREAMING_SNAKE_CASE__		, unittest.TestCase					):

    '''simple docstring'''






    UpperCAmelCase__							  =				(FlaxRegNetModel, FlaxRegNetForImageClassification) if is_flax_available() else ()

    UpperCAmelCase__							  =				False
    UpperCAmelCase__							  =				False
    UpperCAmelCase__							  =				False


    def 	SCREAMING_SNAKE_CASE							(      self				:     Optional[int])						->None:





           '''simple docstring'''

           A__							=   FlaxRegNetModelTester(self)
           A__							=   ConfigTester(self ,			config_class=_SCREAMING_SNAKE_CASE ,			has_text_modality=_SCREAMING_SNAKE_CASE)


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





           '''simple docstring'''

           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				:     Union[str, Any])						->str:





           '''simple docstring'''

           return


    def 	SCREAMING_SNAKE_CASE							(      self				:     Tuple)						->Any:





           '''simple docstring'''

           A__							=   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE)


    def 	SCREAMING_SNAKE_CASE							(      self				:     int)						->Union[str, Any]:





           '''simple docstring'''

           A__							=   self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*_SCREAMING_SNAKE_CASE)


    @unittest.skip(reason='''RegNet does not use inputs_embeds''')
    def 	SCREAMING_SNAKE_CASE							(      self				:     List[str])						->Optional[int]:





           '''simple docstring'''

           pass


    @unittest.skip(reason='''RegNet does not support input and output embeddings''')
    def 	SCREAMING_SNAKE_CASE							(      self				:     int)						->List[Any]:





           '''simple docstring'''

           pass


    def 	SCREAMING_SNAKE_CASE							(      self				:     Dict)						->str:





           '''simple docstring'''

           A__							=   self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  A__							=   model_class(_SCREAMING_SNAKE_CASE)
                  A__							=   inspect.signature(model.__call__)
                  # signature.parameters is an OrderedDict => so arg_names order is deterministic
                  A__							=   [*signature.parameters.keys()]

                  A__							=   ["pixel_values"]
                  self.assertListEqual(arg_names[:1] ,			_SCREAMING_SNAKE_CASE)


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





           '''simple docstring'''

           def check_hidden_states_output(UpperCAmelCase__				:     Dict ,			UpperCAmelCase__				:     Tuple ,			UpperCAmelCase__				:     Optional[int]):
                  A__							=   model_class(_SCREAMING_SNAKE_CASE)

                  A__							=   model(**self._prepare_for_class(_SCREAMING_SNAKE_CASE ,			_SCREAMING_SNAKE_CASE))

                  A__							=   outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                  A__							=   self.model_tester.num_stages
                  self.assertEqual(len(_SCREAMING_SNAKE_CASE) ,			expected_num_stages + 1)

           A__							=   self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  A__							=   True
                  check_hidden_states_output(_SCREAMING_SNAKE_CASE ,			_SCREAMING_SNAKE_CASE ,			_SCREAMING_SNAKE_CASE)

                  # check that output_hidden_states also work using config
                  del inputs_dict["output_hidden_states"]
                  A__							=   True

                  check_hidden_states_output(_SCREAMING_SNAKE_CASE ,			_SCREAMING_SNAKE_CASE ,			_SCREAMING_SNAKE_CASE)





    def 	SCREAMING_SNAKE_CASE							(      self				:     int)						->int:





           '''simple docstring'''

           A__							=   self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  with self.subTest(model_class.__name__):
                         A__							=   self._prepare_for_class(_SCREAMING_SNAKE_CASE ,			_SCREAMING_SNAKE_CASE)
                         A__							=   model_class(_SCREAMING_SNAKE_CASE)

                         @jax.jit
                         def model_jitted(UpperCAmelCase__				:     List[Any] ,			**UpperCAmelCase__				:     Optional[int]):
                                return model(pixel_values=_SCREAMING_SNAKE_CASE ,			**_SCREAMING_SNAKE_CASE)

                         with self.subTest('''JIT Enabled'''):
                                A__							=   model_jitted(**_SCREAMING_SNAKE_CASE).to_tuple()

                         with self.subTest('''JIT Disabled'''):
                                with jax.disable_jit():
                                       A__							=   model_jitted(**_SCREAMING_SNAKE_CASE).to_tuple()

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







def SCREAMING_SNAKE_CASE    (	)       ->	str:



       """simple docstring"""

       A__							=   Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''						)
       return image


@require_flax
class     UpperCamelCase_   (			unittest.TestCase					):

    '''simple docstring'''






    @cached_property
    def 	SCREAMING_SNAKE_CASE							(      self				:     Any)						->Optional[Any]:





           '''simple docstring'''

           return AutoImageProcessor.from_pretrained('''facebook/regnet-y-040''') if is_vision_available() else None





    @slow
    def 	SCREAMING_SNAKE_CASE							(      self				:     Dict)						->str:





           '''simple docstring'''

           A__							=   FlaxRegNetForImageClassification.from_pretrained('''facebook/regnet-y-040''')

           A__							=   self.default_image_processor
           A__							=   prepare_img()
           A__							=   image_processor(images=_SCREAMING_SNAKE_CASE ,			return_tensors='''np''')

           A__							=   model(**_SCREAMING_SNAKE_CASE)

           # verify the logits
           A__							=   (1, 1_000)
           self.assertEqual(outputs.logits.shape ,			_SCREAMING_SNAKE_CASE)

           A__							=   jnp.array([-0.4180, -1.5051, -3.4836])

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








from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_speech_available,
    is_torch_available,
)


_lowerCamelCase     :		Optional[Any]								=       {
    """configuration_trocr""": ["""TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP""", """TrOCRConfig"""],
    """processing_trocr""": ["""TrOCRProcessor"""],
}


try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCamelCase     :		Any								=       [
      """TROCR_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """TrOCRForCausalLM""",
      """TrOCRPreTrainedModel""",
  ]


if TYPE_CHECKING:
  from .configuration_trocr import TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP, TrOCRConfig
  from .processing_trocr import TrOCRProcessor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_trocr import TROCR_PRETRAINED_MODEL_ARCHIVE_LIST, TrOCRForCausalLM, TrOCRPreTrainedModel

else:
  import sys

  _lowerCamelCase     :		List[str]								=       _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


