


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    __snake_case	(							unittest.TestCase      ):





       '''simple docstring'''







       def __init__(  self			:   Any							,    A			:   Any							,    A			:   Optional[int]=3							,    A			:   str=32							,    A			:   Any=3							,    A			:   Dict=10							,    A			:   Any=[10, 20, 30, 40]							,    A			:   Union[str, Any]=[1, 1, 2, 1]							,    A			:   Tuple=True							,    A			:   int=True							,    A			:   Dict="relu"							,    A			:   Union[str, Any]=3							,    A			:   str=None							,    ):
         __snake_case:						Dict			  =						parent
         __snake_case:						Any			  =						batch_size
         __snake_case:						Tuple			  =						image_size
         __snake_case:						int			  =						num_channels
         __snake_case:						Any			  =						embeddings_size
         __snake_case:						Any			  =						hidden_sizes
         __snake_case:						int			  =						depths
         __snake_case:						Any			  =						is_training
         __snake_case:						Optional[int]			  =						use_labels
         __snake_case:						List[Any]			  =						hidden_act
         __snake_case:						Tuple			  =						num_labels
         __snake_case:						Union[str, Any]			  =						scope
         __snake_case:						List[str]			  =						len(lowerCAmelCase__		)




       def 							UpperCAmelCase__			(  self			:   int		):
         __snake_case:						str			  =						floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

         __snake_case:						str			  =						self.get_config()

         return config, pixel_values




       def 							UpperCAmelCase__			(  self			:   List[Any]		):
         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 							UpperCAmelCase__			(  self			:   str							,    A			:   Union[str, Any]							,    A			:   Dict		):
         __snake_case:						Tuple			  =						FlaxRegNetModel(config=lowerCAmelCase__		)
         __snake_case:						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 // 32, self.image_size // 32)							,    )




       def 							UpperCAmelCase__			(  self			:   Any							,    A			:   Union[str, Any]							,    A			:   Dict		):
         __snake_case:						Union[str, Any]			  =						self.num_labels
         __snake_case:						Optional[Any]			  =						FlaxRegNetForImageClassification(config=lowerCAmelCase__		)
         __snake_case:						Union[str, Any]			  =						model(lowerCAmelCase__		)
         self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.num_labels)		)






       def 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         __snake_case:						List[str]			  =						self.prepare_config_and_inputs()
         __snake_case  ,       __snake_case:						Dict			  =						config_and_inputs
         __snake_case:						Any			  =						{"""pixel_values""": pixel_values}
         return config, inputs_dict







@require_flax
class    __snake_case	(							lowerCamelCase_       ,      unittest.TestCase      ):





       '''simple docstring'''







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

       lowerCAmelCase__						=					False
       lowerCAmelCase__						=					False
       lowerCAmelCase__						=					False




       def 							UpperCAmelCase__			(  self			:   Dict		):
         __snake_case:						Tuple			  =						FlaxRegNetModelTester(self		)
         __snake_case:						str			  =						ConfigTester(self							,    config_class=lowerCAmelCase__							,    has_text_modality=lowerCAmelCase__		)




       def 							UpperCAmelCase__			(  self			:   Any		):
         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 							UpperCAmelCase__			(  self			:   Any		):
         return




       def 							UpperCAmelCase__			(  self			:   Dict		):
         __snake_case:						Optional[int]			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*lowerCAmelCase__		)




       def 							UpperCAmelCase__			(  self			:   List[str]		):
         __snake_case:						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 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         pass




       @unittest.skip(reason="""RegNet does not support input and output embeddings"""		)
       def 							UpperCAmelCase__			(  self			:   List[str]		):
         pass




       def 							UpperCAmelCase__			(  self			:   Optional[Any]		):
         __snake_case  ,       __snake_case:						List[Any]			  =						self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
           __snake_case:						Optional[int]			  =						model_class(lowerCAmelCase__		)
           __snake_case:						List[Any]			  =						inspect.signature(model.__call__		)
           # signature.parameters is an OrderedDict => so arg_names order is deterministic
           __snake_case:						int			  =						[*signature.parameters.keys()]

           __snake_case:						Optional[int]			  =						["""pixel_values"""]
           self.assertListEqual(arg_names[:1]							,    lowerCAmelCase__		)




       def 							UpperCAmelCase__			(  self			:   Optional[int]		):
         def check_hidden_states_output(A			:   Any							,    A			:   Union[str, Any]							,    A			:   str		):
           __snake_case:						int			  =						model_class(lowerCAmelCase__		)

           __snake_case:						Tuple			  =						model(**self._prepare_for_class(lowerCAmelCase__							,    lowerCAmelCase__		)		)

           __snake_case:						Union[str, Any]			  =						outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

           __snake_case:						Optional[int]			  =						self.model_tester.num_stages
           self.assertEqual(len(lowerCAmelCase__		)							,    expected_num_stages + 1		)

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

         for model_class in self.all_model_classes:
           __snake_case:						str			  =						True
           check_hidden_states_output(lowerCAmelCase__							,    lowerCAmelCase__							,    lowerCAmelCase__		)

           # check that output_hidden_states also work using config
           del inputs_dict["output_hidden_states"]
           __snake_case:						Union[str, Any]			  =						True

           check_hidden_states_output(lowerCAmelCase__							,    lowerCAmelCase__							,    lowerCAmelCase__		)






       def 							UpperCAmelCase__			(  self			:   Optional[Any]		):
         __snake_case  ,       __snake_case:						str			  =						self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
           with self.subTest(model_class.__name__		):
             __snake_case:						Tuple			  =						self._prepare_for_class(lowerCAmelCase__							,    lowerCAmelCase__		)
             __snake_case:						Dict			  =						model_class(lowerCAmelCase__		)

             @jax.jit
             def model_jitted(A			:   Optional[int]							,    **A			:   int		):
               return model(pixel_values=lowerCAmelCase__							,    **lowerCAmelCase__		)

             with self.subTest("""JIT Enabled"""		):
               __snake_case:						Optional[int]			  =						model_jitted(**lowerCAmelCase__		).to_tuple()

             with self.subTest("""JIT Disabled"""		):
               with jax.disable_jit():
                 __snake_case:						Any			  =						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__							(   )       ->   str:
  __snake_case:						Optional[int]			  =						Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png""")
  return image







@require_flax
class    __snake_case	(							unittest.TestCase      ):





       '''simple docstring'''







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






       @slow
       def 							UpperCAmelCase__			(  self			:   str		):
         __snake_case:						Optional[int]			  =						FlaxRegNetForImageClassification.from_pretrained("""facebook/regnet-y-040"""		)

         __snake_case:						Dict			  =						self.default_image_processor
         __snake_case:						Optional[int]			  =						prepare_img()
         __snake_case:						Optional[int]			  =						image_processor(images=lowerCAmelCase__							,    return_tensors="""np"""		)

         __snake_case:						List[str]			  =						model(**lowerCAmelCase__		)

         # verify the logits
         __snake_case:						Dict			  =						(1, 1_000)
         self.assertEqual(outputs.logits.shape							,    lowerCAmelCase__		)

         __snake_case:						Any			  =						jnp.array([-0.4180, -1.5051, -3.4836]		)

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













from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__UpperCAmelCase	:       List[str]    	=							{
    "configuration_roberta": ["ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP", "RobertaConfig", "RobertaOnnxConfig"],
    "tokenization_roberta": ["RobertaTokenizer"],
}

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase	:       Optional[Any]    	=							["RobertaTokenizerFast"]

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase	:       Tuple    	=							[
          "ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST",
          "RobertaForCausalLM",
          "RobertaForMaskedLM",
          "RobertaForMultipleChoice",
          "RobertaForQuestionAnswering",
          "RobertaForSequenceClassification",
          "RobertaForTokenClassification",
          "RobertaModel",
          "RobertaPreTrainedModel",
      ]

try:
      if not is_tf_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase	:       Optional[int]    	=							[
          "TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST",
          "TFRobertaForCausalLM",
          "TFRobertaForMaskedLM",
          "TFRobertaForMultipleChoice",
          "TFRobertaForQuestionAnswering",
          "TFRobertaForSequenceClassification",
          "TFRobertaForTokenClassification",
          "TFRobertaMainLayer",
          "TFRobertaModel",
          "TFRobertaPreTrainedModel",
      ]

try:
      if not is_flax_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase	:       List[Any]    	=							[
          "FlaxRobertaForCausalLM",
          "FlaxRobertaForMaskedLM",
          "FlaxRobertaForMultipleChoice",
          "FlaxRobertaForQuestionAnswering",
          "FlaxRobertaForSequenceClassification",
          "FlaxRobertaForTokenClassification",
          "FlaxRobertaModel",
          "FlaxRobertaPreTrainedModel",
      ]


if TYPE_CHECKING:
      from .configuration_roberta import ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, RobertaConfig, RobertaOnnxConfig
      from .tokenization_roberta import RobertaTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_roberta_fast import RobertaTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_roberta import (
                ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
                RobertaForCausalLM,
                RobertaForMaskedLM,
                RobertaForMultipleChoice,
                RobertaForQuestionAnswering,
                RobertaForSequenceClassification,
                RobertaForTokenClassification,
                RobertaModel,
                RobertaPreTrainedModel,
            )

      try:
            if not is_tf_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_tf_roberta import (
                TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
                TFRobertaForCausalLM,
                TFRobertaForMaskedLM,
                TFRobertaForMultipleChoice,
                TFRobertaForQuestionAnswering,
                TFRobertaForSequenceClassification,
                TFRobertaForTokenClassification,
                TFRobertaMainLayer,
                TFRobertaModel,
                TFRobertaPreTrainedModel,
            )

      try:
            if not is_flax_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_flax_roberta import (
                FlaxRobertaForCausalLM,
                FlaxRobertaForMaskedLM,
                FlaxRobertaForMultipleChoice,
                FlaxRobertaForQuestionAnswering,
                FlaxRobertaForSequenceClassification,
                FlaxRobertaForTokenClassification,
                FlaxRobertaModel,
                FlaxRobertaPreTrainedModel,
            )

else:
      import sys

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






