

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



    '''simple docstring'''




    def __init__( self ,lowerCamelCase_ ,lowerCamelCase_=3 ,lowerCamelCase_=3_2 ,lowerCamelCase_=3 ,lowerCamelCase_=1_0 ,lowerCamelCase_=[1_0, 2_0, 3_0, 4_0] ,lowerCamelCase_=[1, 1, 2, 1] ,lowerCamelCase_=True ,lowerCamelCase_=True ,lowerCamelCase_="relu" ,lowerCamelCase_=3 ,lowerCamelCase_=None ,)    ->      Optional[int]:
          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(lowerCamelCase_    )




    def UpperCamelCase__		( self    )    ->      int:
          A				  =		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

          A				  =		self.get_config()

          return config, pixel_values




    def UpperCamelCase__		( self    )    ->      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 ,lowerCamelCase_ ,lowerCamelCase_    )    ->      List[str]:
          A				  =		FlaxRegNetModel(config=lowerCamelCase_    )
          A				  =		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 UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_    )    ->      Any:
          A				  =		self.num_labels
          A				  =		FlaxRegNetForImageClassification(config=lowerCamelCase_    )
          A				  =		model(lowerCamelCase_    )
          self.parent.assertEqual(result.logits.shape ,(self.batch_size, self.num_labels)    )




    def UpperCamelCase__		( self    )    ->      List[str]:
          A				  =		self.prepare_config_and_inputs()
          A  ,			A				  =		config_and_inputs
          A				  =		{"""pixel_values""": pixel_values}
          return config, inputs_dict





@require_flax
class lowerCamelCase__   (				SCREAMING_SNAKE_CASE     ,   unittest.TestCase    ):



    '''simple docstring'''




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

    _lowerCamelCase					=   False
    _lowerCamelCase					=   False
    _lowerCamelCase					=   False




    def UpperCamelCase__		( self    )    ->      None:
          A				  =		FlaxRegNetModelTester(self    )
          A				  =		ConfigTester(self ,config_class=lowerCamelCase_ ,has_text_modality=lowerCamelCase_    )




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




    def UpperCamelCase__		( self    )    ->      List[str]:
          A				  =		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*lowerCamelCase_    )




    def UpperCamelCase__		( self    )    ->      int:
          A				  =		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    )    ->      Optional[int]:
          pass




    @unittest.skip(reason="""RegNet does not support input and output embeddings"""    )
    def UpperCamelCase__		( self    )    ->      Any:
          pass




    def UpperCamelCase__		( self    )    ->      Union[str, Any]:
          A  ,			A				  =		self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
                A				  =		model_class(lowerCamelCase_    )
                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] ,lowerCamelCase_    )




    def UpperCamelCase__		( self    )    ->      int:
          def check_hidden_states_output(lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    ):
                A				  =		model_class(lowerCamelCase_    )

                A				  =		model(**self._prepare_for_class(lowerCamelCase_ ,lowerCamelCase_    )    )

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

                A				  =		self.model_tester.num_stages
                self.assertEqual(len(lowerCamelCase_    ) ,expected_num_stages + 1    )

          A  ,			A				  =		self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
                A				  =		True
                check_hidden_states_output(lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )

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

                check_hidden_states_output(lowerCamelCase_ ,lowerCamelCase_ ,lowerCamelCase_    )




    def UpperCamelCase__		( self    )    ->      Any:
          A  ,			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(lowerCamelCase_ ,lowerCamelCase_    )
                      A				  =		model_class(lowerCamelCase_    )

                      @jax.jit
                      def model_jitted(lowerCamelCase_ ,**lowerCamelCase_    ):
                            return model(pixel_values=lowerCamelCase_ ,**lowerCamelCase_    )

                      with self.subTest("""JIT Enabled"""    ):
                            A				  =		model_jitted(**lowerCamelCase_    ).to_tuple()

                      with self.subTest("""JIT Disabled"""    ):
                            with jax.disable_jit():
                                  A				  =		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					(	):



      """simple docstring"""






      A				  =		Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""				)
      return image





@require_flax
class lowerCamelCase__   (				unittest.TestCase    ):



    '''simple docstring'''




    @cached_property
    def UpperCamelCase__		( self    )    ->      Optional[int]:
          return AutoImageProcessor.from_pretrained("""facebook/regnet-y-040"""    ) if is_vision_available() else None




    @slow
    def UpperCamelCase__		( self    )    ->      Any:
          A				  =		FlaxRegNetForImageClassification.from_pretrained("""facebook/regnet-y-040"""    )

          A				  =		self.default_image_processor
          A				  =		prepare_img()
          A				  =		image_processor(images=lowerCamelCase_ ,return_tensors="""np"""    )

          A				  =		model(**lowerCamelCase_    )

          # verify the logits
          A				  =		(1, 1_0_0_0)
          self.assertEqual(outputs.logits.shape ,lowerCamelCase_    )

          A				  =		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"""


import functools
import gc
import inspect

import torch

from .imports import is_npu_available, is_xpu_available



def   _A					(	*_a	:						int				):



      """simple docstring"""






      if not isinstance(_a     ,   _a				):
            A				  =		list(_a				)
      for i in range(len(_a				)				):
            A				  =		None
      gc.collect()
      if is_xpu_available():
            torch.xpu.empty_cache()
      elif is_npu_available():
            torch.npu.empty_cache()
      else:
            torch.cuda.empty_cache()
      return objects



def   _A					(	_a	:						Exception				):



      """simple docstring"""






      A				  =		[
          """CUDA out of memory.""",  # CUDA OOM
          """cuDNN error: CUDNN_STATUS_NOT_SUPPORTED.""",  # CUDNN SNAFU
          """DefaultCPUAllocator: can't allocate memory""",  # CPU OOM
      ]
      if isinstance(_a     ,   _a				) and len(exception.args				) == 1:
            return any(err in exception.args[0] for err in _statements				)
      return False





def   _A					(	_a	:						callable = None     ,   _a	:						int = 1_2_8				):



      """simple docstring"""






      if function is None:
            return functools.partial(_a     ,   starting_batch_size=_a				)

      A				  =		starting_batch_size

      def decorator(*_a	:						Union[str, Any]     ,   **_a	:						List[Any]				):
            nonlocal batch_size
            gc.collect()
            if is_xpu_available():
                  torch.xpu.empty_cache()
            elif is_npu_available():
                  torch.npu.empty_cache()
            else:
                  torch.cuda.empty_cache()
            A				  =		list(inspect.signature(_a				).parameters.keys()				)
            # Guard against user error
            if len(_a				) < (len(_a				) + 1):
                  A				  =		""", """.join([f'{arg}={value}' for arg, value in zip(params[1:]     ,   args[1:]				)]				)
                  raise TypeError(
                      f'Batch size was passed into `{function.__name__}` as the first argument when called.'
                      f'Remove this as the decorator already does so: `{function.__name__}({arg_str})`'				)
            while True:
                  if batch_size == 0:
                        raise RuntimeError("""No executable batch size found, reached zero."""				)
                  try:
                        return function(_a     ,   *_a     ,   **_a				)
                  except Exception as e:
                        if should_reduce_batch_size(_a				):
                              gc.collect()
                              if is_xpu_available():
                                    torch.xpu.empty_cache()
                              elif is_npu_available():
                                    torch.npu.empty_cache()
                              else:
                                    torch.cuda.empty_cache()
                              batch_size //= 2
                        else:
                              raise

      return decorator




