






"""simple docstring"""



import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import MBartConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow
from transformers.utils import cached_property

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
 import os

 # The slow tests are often failing with OOM error on GPU
 # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
 # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
 a  :						List[str]					       =							'''platform'''

 import jax
 import jax.numpy as jnp

 from transformers import AutoTokenizer
 from transformers.models.mbart.modeling_flax_mbart import (
     FlaxMBartForConditionalGeneration,
     FlaxMBartForQuestionAnswering,
     FlaxMBartForSequenceClassification,
     FlaxMBartModel,
     shift_tokens_right,
 )




def       snake_case__		(   _SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE=None       ,	_SCREAMING_SNAKE_CASE=None       ,	_SCREAMING_SNAKE_CASE=None       ,	_SCREAMING_SNAKE_CASE=None       ,	_SCREAMING_SNAKE_CASE=None       ,	_SCREAMING_SNAKE_CASE=None       ,	)    ->Any:
 if attention_mask is None:
  UpperCAmelCase__ 						=				np.where(input_ids != config.pad_token_id       ,	1       ,	0			)
 if decoder_attention_mask is None:
  UpperCAmelCase__ 						=				np.where(decoder_input_ids != config.pad_token_id       ,	1       ,	0			)
 if head_mask is None:
  UpperCAmelCase__ 						=				np.ones((config.encoder_layers, config.encoder_attention_heads)			)
 if decoder_head_mask is None:
  UpperCAmelCase__ 						=				np.ones((config.decoder_layers, config.decoder_attention_heads)			)
 if cross_attn_head_mask is None:
  UpperCAmelCase__ 						=				np.ones((config.decoder_layers, config.decoder_attention_heads)			)
 return {
     "input_ids": input_ids,
     "decoder_input_ids": decoder_input_ids,
     "attention_mask": attention_mask,
     "decoder_attention_mask": decoder_attention_mask,
 }


class 	_UpperCamelCase ( unittest.TestCase							):



    '''simple docstring'''


    def __init__(      self   ,			__lowercase   ,			__lowercase=13   ,			__lowercase=7   ,			__lowercase=True   ,			__lowercase=False   ,			__lowercase=99   ,			__lowercase=16   ,			__lowercase=2   ,			__lowercase=4   ,			__lowercase=4   ,			__lowercase="gelu"   ,			__lowercase=0.1   ,			__lowercase=0.1   ,			__lowercase=32   ,			__lowercase=2   ,			__lowercase=1   ,			__lowercase=0   ,			__lowercase=2   ,			__lowercase=0.02   ,			):
     UpperCAmelCase__ 						=				parent
     UpperCAmelCase__ 						=				batch_size
     UpperCAmelCase__ 						=				seq_length
     UpperCAmelCase__ 						=				is_training
     UpperCAmelCase__ 						=				use_labels
     UpperCAmelCase__ 						=				vocab_size
     UpperCAmelCase__ 						=				hidden_size
     UpperCAmelCase__ 						=				num_hidden_layers
     UpperCAmelCase__ 						=				num_attention_heads
     UpperCAmelCase__ 						=				intermediate_size
     UpperCAmelCase__ 						=				hidden_act
     UpperCAmelCase__ 						=				hidden_dropout_prob
     UpperCAmelCase__ 						=				attention_probs_dropout_prob
     UpperCAmelCase__ 						=				max_position_embeddings
     UpperCAmelCase__ 						=				eos_token_id
     UpperCAmelCase__ 						=				pad_token_id
     UpperCAmelCase__ 						=				bos_token_id
     UpperCAmelCase__ 						=				decoder_start_token_id
     UpperCAmelCase__ 						=				initializer_range


    def 		A__      (      self			):
     UpperCAmelCase__ 						=				np.clip(ids_tensor([self.batch_size, self.seq_length - 1]   ,			self.vocab_size			)   ,			3   ,			self.vocab_size			)
     UpperCAmelCase__ 						=				np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)   ,			dtype=np.intaa			))   ,			-1			)

     UpperCAmelCase__ 						=				shift_tokens_right(_snake_case   ,			1			)

     UpperCAmelCase__ 						=				MBartConfig(
         vocab_size=self.vocab_size   ,			d_model=self.hidden_size   ,			encoder_layers=self.num_hidden_layers   ,			decoder_layers=self.num_hidden_layers   ,			encoder_attention_heads=self.num_attention_heads   ,			decoder_attention_heads=self.num_attention_heads   ,			encoder_ffn_dim=self.intermediate_size   ,			decoder_ffn_dim=self.intermediate_size   ,			dropout=self.hidden_dropout_prob   ,			attention_dropout=self.attention_probs_dropout_prob   ,			max_position_embeddings=self.max_position_embeddings   ,			eos_token_id=self.eos_token_id   ,			bos_token_id=self.bos_token_id   ,			pad_token_id=self.pad_token_id   ,			decoder_start_token_id=self.decoder_start_token_id   ,			initializer_range=self.initializer_range   ,			use_cache=_snake_case   ,			)
     UpperCAmelCase__ 						=				prepare_mbart_inputs_dict(_snake_case   ,			_snake_case   ,			_snake_case			)
     return config, inputs_dict


    def 		A__      (      self			):
     UpperCAmelCase__ 						=				self.prepare_config_and_inputs()
     return config, inputs_dict


    def 		A__      (      self   ,			__lowercase   ,			__lowercase   ,			__lowercase			):
     UpperCAmelCase__ 						=				20
     UpperCAmelCase__ 						=				model_class_name(_snake_case			)

     UpperCAmelCase__ 						=				model.encode(inputs_dict["""input_ids"""]			)

     UpperCAmelCase__ 						=				(
         inputs_dict['decoder_input_ids'],
         inputs_dict['decoder_attention_mask'],
     )

     UpperCAmelCase__ 						=				model.init_cache(decoder_input_ids.shape[0]   ,			_snake_case   ,			_snake_case			)
     UpperCAmelCase__ 						=				jnp.ones((decoder_input_ids.shape[0], max_decoder_length)   ,			dtype="""i4"""			)

     UpperCAmelCase__ 						=				jnp.broadcast_to(
         jnp.arange(decoder_input_ids.shape[-1] - 1			)[None, :]   ,			(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)   ,			)
     UpperCAmelCase__ 						=				model.decode(
         decoder_input_ids[:, :-1]   ,			_snake_case   ,			decoder_attention_mask=_snake_case   ,			past_key_values=_snake_case   ,			decoder_position_ids=_snake_case   ,			)

     UpperCAmelCase__ 						=				jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]   ,			dtype="""i4"""			)
     UpperCAmelCase__ 						=				model.decode(
         decoder_input_ids[:, -1:]   ,			_snake_case   ,			decoder_attention_mask=_snake_case   ,			past_key_values=outputs_cache.past_key_values   ,			decoder_position_ids=_snake_case   ,			)

     UpperCAmelCase__ 						=				model.decode(_snake_case   ,			_snake_case			)

     UpperCAmelCase__ 						=				np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])			)			)
     self.parent.assertTrue(diff < 1e-3   ,			msg=F'''Max diff is {diff}'''			)






    def 		A__      (      self   ,			__lowercase   ,			__lowercase   ,			__lowercase			):
     UpperCAmelCase__ 						=				20
     UpperCAmelCase__ 						=				model_class_name(_snake_case			)

     UpperCAmelCase__ 						=				model.encode(inputs_dict["""input_ids"""]			)

     UpperCAmelCase__ 						=				(
         inputs_dict['decoder_input_ids'],
         inputs_dict['decoder_attention_mask'],
     )

     UpperCAmelCase__ 						=				jnp.concatenate(
         [
             decoder_attention_mask,
             jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])			),
         ]   ,			axis=-1   ,			)

     UpperCAmelCase__ 						=				model.init_cache(decoder_input_ids.shape[0]   ,			_snake_case   ,			_snake_case			)
     UpperCAmelCase__ 						=				jnp.broadcast_to(
         jnp.arange(decoder_input_ids.shape[-1] - 1			)[None, :]   ,			(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)   ,			)

     UpperCAmelCase__ 						=				model.decode(
         decoder_input_ids[:, :-1]   ,			_snake_case   ,			decoder_attention_mask=_snake_case   ,			past_key_values=_snake_case   ,			decoder_position_ids=_snake_case   ,			)
     UpperCAmelCase__ 						=				jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]   ,			dtype="""i4"""			)
     UpperCAmelCase__ 						=				model.decode(
         decoder_input_ids[:, -1:]   ,			_snake_case   ,			past_key_values=outputs_cache.past_key_values   ,			decoder_attention_mask=_snake_case   ,			decoder_position_ids=_snake_case   ,			)

     UpperCAmelCase__ 						=				model.decode(_snake_case   ,			_snake_case   ,			decoder_attention_mask=_snake_case			)

     UpperCAmelCase__ 						=				np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])			)			)
     self.parent.assertTrue(diff < 1e-3   ,			msg=F'''Max diff is {diff}'''			)



@require_flax
class 	_UpperCamelCase ( unittest.TestCase							):



    '''simple docstring'''


    __lowercase    :   Union[str, Any]											=						99


    def 		A__      (      self			):
     UpperCAmelCase__ 						=				np.array(
         [
             [71, 82, 18, 33, 46, 91, 2],
             [68, 34, 26, 58, 30, 82, 2],
             [5, 97, 17, 39, 94, 40, 2],
             [76, 83, 94, 25, 70, 78, 2],
             [87, 59, 41, 35, 48, 66, 2],
             [55, 13, 16, 58, 5, 2, 1],  # note padding
             [64, 27, 31, 51, 12, 75, 2],
             [52, 64, 86, 17, 83, 39, 2],
             [48, 61, 9, 24, 71, 82, 2],
             [26, 1, 60, 48, 22, 13, 2],
             [21, 5, 62, 28, 14, 76, 2],
             [45, 98, 37, 86, 59, 48, 2],
             [70, 70, 50, 9, 28, 0, 2],
         ]   ,			dtype=np.intaa   ,			)

     UpperCAmelCase__ 						=				input_ids.shape[0]
     UpperCAmelCase__ 						=				MBartConfig(
         vocab_size=self.vocab_size   ,			d_model=24   ,			encoder_layers=2   ,			decoder_layers=2   ,			encoder_attention_heads=2   ,			decoder_attention_heads=2   ,			encoder_ffn_dim=32   ,			decoder_ffn_dim=32   ,			max_position_embeddings=48   ,			eos_token_id=2   ,			pad_token_id=1   ,			bos_token_id=0   ,			)
     return config, input_ids, batch_size


    def 		A__      (      self			):
     UpperCAmelCase__ 						=				self._get_config_and_data()
     UpperCAmelCase__ 						=				FlaxMBartForSequenceClassification(_snake_case			)
     UpperCAmelCase__ 						=				model(input_ids=_snake_case   ,			decoder_input_ids=_snake_case			)
     UpperCAmelCase__ 						=				(batch_size, config.num_labels)
     self.assertEqual(outputs["""logits"""].shape   ,			_snake_case			)


    def 		A__      (      self			):
     UpperCAmelCase__ 						=				self._get_config_and_data()
     UpperCAmelCase__ 						=				FlaxMBartForQuestionAnswering(_snake_case			)
     UpperCAmelCase__ 						=				model(input_ids=_snake_case			)

     self.assertEqual(outputs["""start_logits"""].shape   ,			input_ids.shape			)
     self.assertEqual(outputs["""end_logits"""].shape   ,			input_ids.shape			)


    def 		A__      (      self			):
     UpperCAmelCase__ 						=				self._get_config_and_data()
     UpperCAmelCase__ 						=				FlaxMBartForConditionalGeneration(_snake_case			)
     UpperCAmelCase__ 						=				lm_model(input_ids=_snake_case			)
     UpperCAmelCase__ 						=				(batch_size, input_ids.shape[1], config.vocab_size)
     self.assertEqual(outputs["""logits"""].shape   ,			_snake_case			)


    def 		A__      (      self			):
     UpperCAmelCase__ 						=				MBartConfig(
         vocab_size=self.vocab_size   ,			d_model=14   ,			encoder_layers=2   ,			decoder_layers=2   ,			encoder_attention_heads=2   ,			decoder_attention_heads=2   ,			encoder_ffn_dim=8   ,			decoder_ffn_dim=8   ,			max_position_embeddings=48   ,			)
     UpperCAmelCase__ 						=				FlaxMBartForConditionalGeneration(_snake_case			)
     UpperCAmelCase__ 						=				np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]   ,			dtype=np.intaa			)
     UpperCAmelCase__ 						=				np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]   ,			dtype=np.intaa			)
     UpperCAmelCase__ 						=				lm_model(input_ids=_snake_case   ,			decoder_input_ids=_snake_case			)
     UpperCAmelCase__ 						=				(*summary.shape, config.vocab_size)
     self.assertEqual(outputs["""logits"""].shape   ,			_snake_case			)






    def 		A__      (      self			):
     UpperCAmelCase__ 						=				np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]   ,			dtype=np.intaa			)
     UpperCAmelCase__ 						=				shift_tokens_right(_snake_case   ,			1			)
     UpperCAmelCase__ 						=				np.equal(_snake_case   ,			1			).astype(np.floataa			).sum()
     UpperCAmelCase__ 						=				np.equal(_snake_case   ,			1			).astype(np.floataa			).sum()
     self.assertEqual(shifted.shape   ,			input_ids.shape			)
     self.assertEqual(_snake_case   ,			n_pad_before - 1			)
     self.assertTrue(np.equal(shifted[:, 0]   ,			2			).all()			)



@require_flax
class 	_UpperCamelCase ( a_		,     unittest.TestCase		,     a_							):



    '''simple docstring'''


    __lowercase    :   Union[str, Any]											=						True
    __lowercase    :   int											=						(
        (
            FlaxMBartModel,
            FlaxMBartForConditionalGeneration,
            FlaxMBartForSequenceClassification,
            FlaxMBartForQuestionAnswering,
        )
        if is_flax_available()
        else ()
    )
    __lowercase    :   Tuple											=						(FlaxMBartForConditionalGeneration,) if is_flax_available() else ()


    def 		A__      (      self			):
     UpperCAmelCase__ 						=				FlaxMBartModelTester(self			)


    def 		A__      (      self			):
     UpperCAmelCase__ 						=				self.model_tester.prepare_config_and_inputs()
     for model_class in self.all_model_classes:
      self.model_tester.check_use_cache_forward(_snake_case   ,			_snake_case   ,			_snake_case			)


    def 		A__      (      self			):
     UpperCAmelCase__ 						=				self.model_tester.prepare_config_and_inputs()
     for model_class in self.all_model_classes:
      self.model_tester.check_use_cache_forward_with_attn_mask(_snake_case   ,			_snake_case   ,			_snake_case			)


    def 		A__      (      self			):
     UpperCAmelCase__ 						=				self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
      with self.subTest(model_class.__name__			):
       UpperCAmelCase__ 						=				self._prepare_for_class(_snake_case   ,			_snake_case			)
       UpperCAmelCase__ 						=				model_class(_snake_case			)

       @jax.jit
       def encode_jitted(__lowercase   ,			__lowercase=None   ,			**__lowercase			):
        return model.encode(input_ids=_snake_case   ,			attention_mask=_snake_case			)

       with self.subTest("""JIT Enabled"""			):
        UpperCAmelCase__ 						=				encode_jitted(**_snake_case			).to_tuple()

       with self.subTest("""JIT Disabled"""			):
        with jax.disable_jit():
         UpperCAmelCase__ 						=				encode_jitted(**_snake_case			).to_tuple()

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


    def 		A__      (      self			):
     UpperCAmelCase__ 						=				self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
      with self.subTest(model_class.__name__			):
       UpperCAmelCase__ 						=				model_class(_snake_case			)
       UpperCAmelCase__ 						=				model.encode(inputs_dict["""input_ids"""]   ,			inputs_dict["""attention_mask"""]			)

       UpperCAmelCase__ 						=				{
           'decoder_input_ids': inputs_dict['decoder_input_ids'],
           'decoder_attention_mask': inputs_dict['decoder_attention_mask'],
           'encoder_outputs': encoder_outputs,
       }

       @jax.jit
       def decode_jitted(__lowercase   ,			__lowercase   ,			__lowercase			):
        return model.decode(
            decoder_input_ids=_snake_case   ,			decoder_attention_mask=_snake_case   ,			encoder_outputs=_snake_case   ,			)

       with self.subTest("""JIT Enabled"""			):
        UpperCAmelCase__ 						=				decode_jitted(**_snake_case			).to_tuple()

       with self.subTest("""JIT Disabled"""			):
        with jax.disable_jit():
         UpperCAmelCase__ 						=				decode_jitted(**_snake_case			).to_tuple()

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






    @slow
    def 		A__      (      self			):
     for model_class_name in self.all_model_classes:
      UpperCAmelCase__ 						=				model_class_name.from_pretrained("""facebook/mbart-large-cc25"""   ,			from_pt=_snake_case			)
      # FlaxMBartForSequenceClassification expects eos token in input_ids
      UpperCAmelCase__ 						=				np.ones((1, 1)			) * model.config.eos_token_id
      UpperCAmelCase__ 						=				model(_snake_case			)
      self.assertIsNotNone(_snake_case			)



@require_flax
@require_sentencepiece
@require_tokenizers
class 	_UpperCamelCase ( unittest.TestCase							):



    '''simple docstring'''


    __lowercase    :   Union[str, Any]											=						[
        " UN Chief Says There Is No Military Solution in Syria",
    ]
    __lowercase    :   Union[str, Any]											=						[
        "Şeful ONU declară că nu există o soluţie militară în Siria",
    ]
    __lowercase    :   int											=						"facebook/mbart-large-en-ro"


    @cached_property
    def 		A__      (      self			):
     return AutoTokenizer.from_pretrained(self.model_name			)


    @cached_property
    def 		A__      (      self			):
     UpperCAmelCase__ 						=				FlaxMBartForConditionalGeneration.from_pretrained(self.model_name   ,			from_pt=_snake_case			)
     return model


    def 		A__      (      self   ,			**__lowercase			):
     UpperCAmelCase__ 						=				self.translate_src_text(**_snake_case			)
     self.assertListEqual(self.expected_text   ,			_snake_case			)


    def 		A__      (      self   ,			**__lowercase			):
     UpperCAmelCase__ 						=				self.tokenizer(self.src_text   ,			**_snake_case   ,			return_tensors="""np"""			)
     UpperCAmelCase__ 						=				self.model.generate(
         model_inputs.input_ids   ,			attention_mask=model_inputs.attention_mask   ,			decoder_start_token_id=self.tokenizer.lang_code_to_id["""ro_RO"""]   ,			early_stopping=_snake_case   ,			num_beams=2   ,			).sequences
     UpperCAmelCase__ 						=				self.tokenizer.batch_decode(_snake_case   ,			skip_special_tokens=_snake_case			)
     return generated_words






    @slow
    def 		A__      (      self			):
     self._assert_generated_batch_equal_expected()






'''simple docstring'''


import inspect
import unittest

from transformers import ResNetConfig, 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.resnet.modeling_flax_resnet import FlaxResNetForImageClassification, FlaxResNetModel

if is_vision_available():
						from PIL import Image

						from transformers import AutoImageProcessor



class        lowercase    (	unittest.TestCase			):


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


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

								UpperCAmelCase_		:   Optional[int]													=       self.get_config()

								return config, pixel_values


							def     _snake_case				(      self)				->					Union[str, Any]:
								return ResNetConfig(
								    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     _snake_case				(      self     ,      _snake_case     ,      _snake_case)				->					str:
								UpperCAmelCase_		:   Any													=       FlaxResNetModel(config=_snake_case)
								UpperCAmelCase_		:   Optional[Any]													=       model(_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     _snake_case				(      self     ,      _snake_case     ,      _snake_case)				->					List[str]:
								UpperCAmelCase_		:   Union[str, Any]													=       self.num_labels
								UpperCAmelCase_		:   Optional[Any]													=       FlaxResNetForImageClassification(config=_snake_case)
								UpperCAmelCase_		:   Optional[Any]													=       model(_snake_case)
								self.parent.assertEqual(result.logits.shape     ,      (self.batch_size, self.num_labels))





							def     _snake_case				(      self)				->					str:
								UpperCAmelCase_		:   List[Any]													=       self.prepare_config_and_inputs()
								UpperCAmelCase_						,     UpperCAmelCase_		:   List[Any]													=       config_and_inputs
								UpperCAmelCase_		:   str													=       {'pixel_values': pixel_values}
								return config, inputs_dict



@require_flax
class        lowercase    (	a_,       unittest.TestCase			):
							_lowerCamelCase		:   Dict= (FlaxResNetModel, FlaxResNetForImageClassification) if is_flax_available() else ()

							_lowerCamelCase		:   Any= False
							_lowerCamelCase		:   List[Any]= False
							_lowerCamelCase		:   str= False


							def     _snake_case				(      self)				->					None:
								UpperCAmelCase_		:   Any													=       FlaxResNetModelTester(self)
								UpperCAmelCase_		:   Optional[int]													=       ConfigTester(self     ,      config_class=_snake_case     ,      has_text_modality=_snake_case)


							def     _snake_case				(      self)				->					Tuple:
								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     _snake_case				(      self)				->					int:
								return


							def     _snake_case				(      self)				->					Optional[int]:
								UpperCAmelCase_		:   Optional[Any]													=       self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_snake_case)


							def     _snake_case				(      self)				->					Dict:
								UpperCAmelCase_		:   Union[str, Any]													=       self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*_snake_case)


							@unittest.skip(reason='ResNet does not use inputs_embeds')
							def     _snake_case				(      self)				->					Optional[Any]:
								pass


							@unittest.skip(reason='ResNet does not support input and output embeddings')
							def     _snake_case				(      self)				->					Union[str, Any]:
								pass


							def     _snake_case				(      self)				->					Dict:
								UpperCAmelCase_						,     UpperCAmelCase_		:   List[Any]													=       self.model_tester.prepare_config_and_inputs_for_common()

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

									UpperCAmelCase_		:   List[str]													=       ['pixel_values']
									self.assertListEqual(arg_names[:1]     ,      _snake_case)


							def     _snake_case				(      self)				->					List[Any]:
								def check_hidden_states_output(_snake_case     ,      _snake_case     ,      _snake_case):
									UpperCAmelCase_		:   Union[str, Any]													=       model_class(_snake_case)

									UpperCAmelCase_		:   Any													=       model(**self._prepare_for_class(_snake_case     ,      _snake_case))

									UpperCAmelCase_		:   List[Any]													=       outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

									UpperCAmelCase_		:   int													=       self.model_tester.num_stages
									self.assertEqual(len(_snake_case)     ,      expected_num_stages + 1)


							@unittest.skip(reason='ResNet does not use feedforward chunking')
							def     _snake_case				(      self)				->					List[Any]:
								pass





							def     _snake_case				(      self)				->					Tuple:
								UpperCAmelCase_						,     UpperCAmelCase_		:   Any													=       self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
									with self.subTest(model_class.__name__):
										UpperCAmelCase_		:   str													=       self._prepare_for_class(_snake_case     ,      _snake_case)
										UpperCAmelCase_		:   str													=       model_class(_snake_case)

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

										with self.subTest('JIT Enabled'):
											UpperCAmelCase_		:   Optional[Any]													=       model_jitted(**_snake_case).to_tuple()

										with self.subTest('JIT Disabled'):
											with jax.disable_jit():
												UpperCAmelCase_		:   Union[str, Any]													=       model_jitted(**_snake_case).to_tuple()

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







def     SCREAMING_SNAKE_CASE(	)     ->     Dict:
	UpperCAmelCase_		:   Tuple													=       Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'	)
	return image







@require_flax
class        lowercase    (	unittest.TestCase			):


							@cached_property
							def     _snake_case				(      self)				->					List[str]:
								return AutoImageProcessor.from_pretrained('microsoft/resnet-50') if is_vision_available() else None





							@slow
							def     _snake_case				(      self)				->					Union[str, Any]:
								UpperCAmelCase_		:   Union[str, Any]													=       FlaxResNetForImageClassification.from_pretrained('microsoft/resnet-50')

								UpperCAmelCase_		:   List[str]													=       self.default_image_processor
								UpperCAmelCase_		:   Union[str, Any]													=       prepare_img()
								UpperCAmelCase_		:   int													=       image_processor(images=_snake_case     ,      return_tensors='np')

								UpperCAmelCase_		:   Any													=       model(**_snake_case)

								# verify the logits
								UpperCAmelCase_		:   Optional[int]													=       (1, 1000)
								self.assertEqual(outputs.logits.shape     ,      _snake_case)

								UpperCAmelCase_		:   Tuple													=       jnp.array([-11.1_069, -9.7_877, -8.3_777])

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






