





'''simple docstring'''


from __future__ import annotations

import inspect
import unittest
from typing import List, Tuple

from transformers import RegNetConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
							import tensorflow as tf

							from transformers import TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST, TFRegNetForImageClassification, TFRegNetModel


if is_vision_available():
							from PIL import Image

							from transformers import AutoImageProcessor

class    lowerCAmelCase_		:




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




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

					_lowerCAmelCase				=      None
					if self.use_labels:
									_lowerCAmelCase				=      ids_tensor([self.batch_size]			,     self.num_labels		)

					_lowerCAmelCase				=      self.get_config()
					return config, pixel_values, labels




	def 							_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			,     )




	def 							_snake_case   (	self			,     _lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase		)    ->							Optional[int]:
					_lowerCAmelCase				=      TFRegNetModel(config=_lowerCAmelCase		)
					_lowerCAmelCase				=      model(_lowerCAmelCase			,     training=_lowerCAmelCase		)
					# expected last hidden states: B, C, H // 32, W // 32
					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			,     _lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase		)    ->							Optional[Any]:
					_lowerCAmelCase				=      self.num_labels
					_lowerCAmelCase				=      TFRegNetForImageClassification(_lowerCAmelCase		)
					_lowerCAmelCase				=      model(_lowerCAmelCase			,     labels=_lowerCAmelCase			,     training=_lowerCAmelCase		)
					self.parent.assertEqual(result.logits.shape			,     (self.batch_size, self.num_labels)		)





	def 							_snake_case   (	self		)    ->							Optional[int]:
					_lowerCAmelCase				=      self.prepare_config_and_inputs()
					_lowerCAmelCase			,			_lowerCAmelCase			,			_lowerCAmelCase				=      config_and_inputs
					_lowerCAmelCase				=      {"pixel_values": pixel_values}
					return config, inputs_dict







@require_tf
class    lowerCAmelCase_		(		__magic_name__						,__magic_name__						,unittest.TestCase		):
	__lowerCamelCase						:							Optional[Any]         =					(TFRegNetModel, TFRegNetForImageClassification) if is_tf_available() else ()
	__lowerCamelCase						:							Dict         =					(
	    {"feature-extraction": TFRegNetModel, "image-classification": TFRegNetForImageClassification}
	    if is_tf_available()
	    else {}
	)

	__lowerCamelCase						:							List[str]         =					False
	__lowerCamelCase						:							Dict         =					False
	__lowerCamelCase						:							Optional[Any]         =					False
	__lowerCamelCase						:							Optional[Any]         =					False
	__lowerCamelCase						:							Optional[int]         =					False




	def 							_snake_case   (	self		)    ->							List[Any]:
					_lowerCAmelCase				=      TFRegNetModelTester(self		)
					_lowerCAmelCase				=      ConfigTester(self			,     config_class=_lowerCAmelCase			,     has_text_modality=_lowerCAmelCase		)




	def 							_snake_case   (	self		)    ->							Tuple:
					return




	@unittest.skip(reason="RegNet does not use inputs_embeds"		)
	def 							_snake_case   (	self		)    ->							int:
					pass




	@unittest.skipIf(
	    not is_tf_available() or len(tf.config.list_physical_devices("GPU"		)		) == 0			,     reason="TF does not support backprop for grouped convolutions on CPU."			,     )
	@slow
	def 							_snake_case   (	self		)    ->							Union[str, Any]:
					super().test_keras_fit()




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




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

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

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




	def 							_snake_case   (	self		)    ->							List[str]:
					_lowerCAmelCase				=      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*_lowerCAmelCase		)




	def 							_snake_case   (	self		)    ->							List[Any]:
					def check_hidden_states_output(_lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase		):
									_lowerCAmelCase				=      model_class(_lowerCAmelCase		)
									_lowerCAmelCase				=      model(**self._prepare_for_class(_lowerCAmelCase			,     _lowerCAmelCase		)			,     training=_lowerCAmelCase		)

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

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

									# RegNet's feature maps are of shape (batch_size, num_channels, height, width)
									self.assertListEqual(
									    list(hidden_states[0].shape[-2:]		)			,     [self.model_tester.image_size // 2, self.model_tester.image_size // 2]			,     )

					_lowerCAmelCase			,			_lowerCAmelCase				=      self.model_tester.prepare_config_and_inputs_for_common()
					_lowerCAmelCase				=      ["basic", "bottleneck"]
					for model_class in self.all_model_classes:
									for layer_type in layers_type:
													_lowerCAmelCase				=      layer_type
													_lowerCAmelCase				=      True
													check_hidden_states_output(_lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase		)

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

													check_hidden_states_output(_lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase		)




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

					def check_equivalence(_lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase={}		):
									_lowerCAmelCase				=      model(_lowerCAmelCase			,     return_dict=_lowerCAmelCase			,     **_lowerCAmelCase		)
									_lowerCAmelCase				=      model(_lowerCAmelCase			,     return_dict=_lowerCAmelCase			,     **_lowerCAmelCase		).to_tuple()

									def recursive_check(_lowerCAmelCase			,     _lowerCAmelCase		):
													if isinstance(_lowerCAmelCase			,     (List, Tuple)		):
																	for tuple_iterable_value, dict_iterable_value in zip(_lowerCAmelCase			,     _lowerCAmelCase		):
																					recursive_check(_lowerCAmelCase			,     _lowerCAmelCase		)
													elif tuple_object is None:
																	return
													else:
																	self.assertTrue(
																	    all(tf.equal(_lowerCAmelCase			,     _lowerCAmelCase		)		)			,     msg=(
																	        "Tuple and dict output are not equal. Difference:"
																	        f''' {tf.math.reduce_max(tf.abs(tuple_object - dict_object		)		)}'''
																	    )			,     )

													recursive_check(_lowerCAmelCase			,     _lowerCAmelCase		)

					for model_class in self.all_model_classes:
									_lowerCAmelCase				=      model_class(_lowerCAmelCase		)

									_lowerCAmelCase				=      self._prepare_for_class(_lowerCAmelCase			,     _lowerCAmelCase		)
									_lowerCAmelCase				=      self._prepare_for_class(_lowerCAmelCase			,     _lowerCAmelCase		)
									check_equivalence(_lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase		)

									_lowerCAmelCase				=      self._prepare_for_class(_lowerCAmelCase			,     _lowerCAmelCase			,     return_labels=_lowerCAmelCase		)
									_lowerCAmelCase				=      self._prepare_for_class(_lowerCAmelCase			,     _lowerCAmelCase			,     return_labels=_lowerCAmelCase		)
									check_equivalence(_lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase		)

									_lowerCAmelCase				=      self._prepare_for_class(_lowerCAmelCase			,     _lowerCAmelCase		)
									_lowerCAmelCase				=      self._prepare_for_class(_lowerCAmelCase			,     _lowerCAmelCase		)
									check_equivalence(_lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase			,     {"output_hidden_states": True}		)

									_lowerCAmelCase				=      self._prepare_for_class(_lowerCAmelCase			,     _lowerCAmelCase			,     return_labels=_lowerCAmelCase		)
									_lowerCAmelCase				=      self._prepare_for_class(_lowerCAmelCase			,     _lowerCAmelCase			,     return_labels=_lowerCAmelCase		)
									check_equivalence(_lowerCAmelCase			,     _lowerCAmelCase			,     _lowerCAmelCase			,     {"output_hidden_states": True}		)




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





	@slow
	def 							_snake_case   (	self		)    ->							int:
					for model_name in TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									_lowerCAmelCase				=      TFRegNetModel.from_pretrained(_lowerCAmelCase		)
									self.assertIsNotNone(_lowerCAmelCase		)







def 	__a():



				'''simple docstring'''





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





@require_tf
@require_vision
class    lowerCAmelCase_		(		unittest.TestCase		):




	@cached_property
	def 							_snake_case   (	self		)    ->							Optional[int]:
					return (
					    AutoImageProcessor.from_pretrained(TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]		)
					    if is_vision_available()
					    else None
					)





	@slow
	def 							_snake_case   (	self		)    ->							List[Any]:
					_lowerCAmelCase				=      TFRegNetForImageClassification.from_pretrained(TF_REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]		)

					_lowerCAmelCase				=      self.default_image_processor
					_lowerCAmelCase				=      prepare_img()
					_lowerCAmelCase				=      image_processor(images=_lowerCAmelCase			,     return_tensors="tf"		)

					# forward pass
					_lowerCAmelCase				=      model(**_lowerCAmelCase			,     training=_lowerCAmelCase		)

					# verify the logits
					_lowerCAmelCase				=      tf.TensorShape((1, 1000)		)
					self.assertEqual(outputs.logits.shape			,     _lowerCAmelCase		)

					_lowerCAmelCase				=      tf.constant([-0.4180, -1.5051, -3.4836]		)

					tf.debugging.assert_near(outputs.logits[0, :3]			,     _lowerCAmelCase			,     atol=1E-4		)









'''simple docstring'''


import jax.numpy as jnp

from ...utils import logging
from ..ta.modeling_flax_ta import FlaxTaEncoderModel, FlaxTaForConditionalGeneration, FlaxTaModel
from .configuration_mta import MTaConfig


_SCREAMING_SNAKE_CASE	 =							logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE	 =							"T5Config"


def 	__a(SCREAMING_SNAKE_CASE_  :					jnp.array		,       SCREAMING_SNAKE_CASE_  :					int		,       SCREAMING_SNAKE_CASE_  :					int    ):



    '''simple docstring'''





    _lowerCAmelCase				=      jnp.zeros_like(SCREAMING_SNAKE_CASE_    )
    _lowerCAmelCase				=      shifted_input_ids.at[:, 1:].set(input_ids[:, :-1]    )
    _lowerCAmelCase				=      shifted_input_ids.at[:, 0].set(SCREAMING_SNAKE_CASE_    )

    _lowerCAmelCase				=      jnp.where(shifted_input_ids == -100		,       SCREAMING_SNAKE_CASE_		,       SCREAMING_SNAKE_CASE_    )
    return shifted_input_ids

class    lowerCAmelCase_		(		__magic_name__		):
 __lowerCamelCase						:							List[str]         =					"mt5"
 __lowerCamelCase						:							Any         =					MTaConfig







class    lowerCAmelCase_		(		__magic_name__		):
 __lowerCamelCase						:							List[str]         =					"mt5"
 __lowerCamelCase						:							Dict         =					MTaConfig







class    lowerCAmelCase_		(		__magic_name__		):
 __lowerCamelCase						:							Optional[Any]         =					"mt5"
 __lowerCamelCase						:							str         =					MTaConfig


