






import inspect
import unittest

import numpy as np

from transformers import ViTConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
							import jax

							from transformers.models.vit.modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel
class  _UpperCAmelCase   (			unittest.TestCase      ):

				"""simple docstring"""




				def __init__(	self		:       Union[str, Any],    lowerCamelCase		:       Optional[int],    lowerCamelCase		:       Optional[int]=13,    lowerCamelCase		:       List[str]=30,    lowerCamelCase		:       Dict=2,    lowerCamelCase		:       Dict=3,    lowerCamelCase		:       Tuple=True,    lowerCamelCase		:       List[Any]=True,    lowerCamelCase		:       List[str]=32,    lowerCamelCase		:       Optional[int]=5,    lowerCamelCase		:       int=4,    lowerCamelCase		:       Tuple=37,    lowerCamelCase		:       int="gelu",    lowerCamelCase		:       List[str]=0.1,    lowerCamelCase		:       Any=0.1,    lowerCamelCase		:       Tuple=10,    lowerCamelCase		:       str=0.02,    ):

								'''simple docstring'''
								lowercase__  					=     parent
								lowercase__  					=     batch_size
								lowercase__  					=     image_size
								lowercase__  					=     patch_size
								lowercase__  					=     num_channels
								lowercase__  					=     is_training
								lowercase__  					=     use_labels
								lowercase__  					=     hidden_size
								lowercase__  					=     num_hidden_layers
								lowercase__  					=     num_attention_heads
								lowercase__  					=     intermediate_size
								lowercase__  					=     hidden_act
								lowercase__  					=     hidden_dropout_prob
								lowercase__  					=     attention_probs_dropout_prob
								lowercase__  					=     type_sequence_label_size
								lowercase__  					=     initializer_range

								# in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
								lowercase__  					=     (image_size // patch_size) ** 2
								lowercase__  					=     num_patches + 1




				def 							lowercase__  (	self		:       Dict    ):

								'''simple docstring'''
								lowercase__  					=     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

								lowercase__  					=     ViTConfig(
								    image_size=self.image_size,    patch_size=self.patch_size,    num_channels=self.num_channels,    hidden_size=self.hidden_size,    num_hidden_layers=self.num_hidden_layers,    num_attention_heads=self.num_attention_heads,    intermediate_size=self.intermediate_size,    hidden_act=self.hidden_act,    hidden_dropout_prob=self.hidden_dropout_prob,    attention_probs_dropout_prob=self.attention_probs_dropout_prob,    is_decoder=lowerCamelCase,    initializer_range=self.initializer_range,    )

								return config, pixel_values




				def 							lowercase__  (	self		:       List[str],    lowerCamelCase		:       Dict,    lowerCamelCase		:       Tuple    ):

								'''simple docstring'''
								lowercase__  					=     FlaxViTModel(config=lowerCamelCase    )
								lowercase__  					=     model(lowerCamelCase    )
								# expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
								lowercase__  					=     (self.image_size, self.image_size)
								lowercase__  					=     (self.patch_size, self.patch_size)
								lowercase__  					=     (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
								self.parent.assertEqual(result.last_hidden_state.shape,    (self.batch_size, num_patches + 1, self.hidden_size)    )




				def 							lowercase__  (	self		:       Dict,    lowerCamelCase		:       str,    lowerCamelCase		:       int    ):

								'''simple docstring'''
								lowercase__  					=     self.type_sequence_label_size
								lowercase__  					=     FlaxViTForImageClassification(config=lowerCamelCase    )
								lowercase__  					=     model(lowerCamelCase    )
								self.parent.assertEqual(result.logits.shape,    (self.batch_size, self.type_sequence_label_size)    )

								# test greyscale images
								lowercase__  					=     1
								lowercase__  					=     FlaxViTForImageClassification(lowerCamelCase    )

								lowercase__  					=     floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
								lowercase__  					=     model(lowerCamelCase    )




				def 							lowercase__  (	self		:       Tuple    ):

								'''simple docstring'''
								lowercase__  					=     self.prepare_config_and_inputs()
								(
								    (
								    lowercase__
								)					,		(
								    lowercase__
								)					,		
								)  					=     config_and_inputs
								lowercase__  					=     {'''pixel_values''': pixel_values}
								return config, inputs_dict



@require_flax
class  _UpperCAmelCase   (			A__						,unittest.TestCase      ):

				"""simple docstring"""




				lowercase__											=			(FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()




				def 							lowercase__  (	self		:       List[Any]    ):

								'''simple docstring'''
								lowercase__  					=     FlaxViTModelTester(self    )
								lowercase__  					=     ConfigTester(self,    config_class=lowerCamelCase,    has_text_modality=lowerCamelCase,    hidden_size=37    )




				def 							lowercase__  (	self		:       List[Any]    ):

								'''simple docstring'''
								self.config_tester.run_common_tests()




				def 							lowercase__  (	self		:       Any    ):

								'''simple docstring'''
								lowercase__  					=     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*lowerCamelCase    )




				def 							lowercase__  (	self		:       Union[str, Any]    ):

								'''simple docstring'''
								lowercase__  					=     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*lowerCamelCase    )




				def 							lowercase__  (	self		:       Any    ):

								'''simple docstring'''
								lowercase__					,		lowercase__  					=     self.model_tester.prepare_config_and_inputs_for_common()

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

												lowercase__  					=     ['''pixel_values''']
												self.assertListEqual(arg_names[:1],    lowerCamelCase    )




				def 							lowercase__  (	self		:       Optional[int]    ):

								'''simple docstring'''
								lowercase__					,		lowercase__  					=     self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												with self.subTest(model_class.__name__    ):
																lowercase__  					=     self._prepare_for_class(lowerCamelCase,    lowerCamelCase    )
																lowercase__  					=     model_class(lowerCamelCase    )

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

																with self.subTest('''JIT Enabled'''    ):
																				lowercase__  					=     model_jitted(**lowerCamelCase    ).to_tuple()

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




				@slow
				def 							lowercase__  (	self		:       List[str]    ):

								'''simple docstring'''
								for model_class_name in self.all_model_classes:
												lowercase__  					=     model_class_name.from_pretrained('''google/vit-base-patch16-224'''    )
												lowercase__  					=     model(np.ones((1, 3, 224, 224)    )    )
												self.assertIsNotNone(lowerCamelCase    )












import numpy as np

import datasets


A__      :     int       =					'\nCompute the Mahalanobis Distance\n\nMahalonobis distance is the distance between a point and a distribution.\nAnd not between two distinct points. It is effectively a multivariate equivalent of the Euclidean distance.\nIt was introduced by Prof. P. C. Mahalanobis in 1936\nand has been used in various statistical applications ever since\n[source: https://www.machinelearningplus.com/statistics/mahalanobis-distance/]\n'

A__      :     Optional[int]       =					'\\n@article{de2000mahalanobis,\n  title={The mahalanobis distance},\n  author={De Maesschalck, Roy and Jouan-Rimbaud, Delphine and Massart, D{\'e}sir{\'e} L},\n  journal={Chemometrics and intelligent laboratory systems},\n  volume={50},\n  number={1},\n  pages={1--18},\n  year={2000},\n  publisher={Elsevier}\n}\n'

A__      :     Optional[int]       =					'\nArgs:\n    X: List of datapoints to be compared with the `reference_distribution`.\n    reference_distribution: List of datapoints from the reference distribution we want to compare to.\nReturns:\n    mahalanobis: The Mahalonobis distance for each datapoint in `X`.\nExamples:\n\n    >>> mahalanobis_metric = datasets.load_metric("mahalanobis")\n    >>> results = mahalanobis_metric.compute(reference_distribution=[[0, 1], [1, 0]], X=[[0, 1]])\n    >>> print(results)\n    {\'mahalanobis\': array([0.5])}\n'
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION						,_KWARGS_DESCRIPTION      )
class  _UpperCAmelCase   (			datasets.Metric      ):

				"""simple docstring"""




				def 							lowercase__  (	self		:       str    ):

								'''simple docstring'''
								return datasets.MetricInfo(
								    description=_DESCRIPTION,    citation=_CITATION,    inputs_description=_KWARGS_DESCRIPTION,    features=datasets.Features(
								        {
								            '''X''': datasets.Sequence(datasets.Value('''float''',    id='''sequence'''    ),    id='''X'''    ),
								        }    ),    )




				def 							lowercase__  (	self		:       Dict,    lowerCamelCase		:       int,    lowerCamelCase		:       Optional[int]    ):

								'''simple docstring'''
								# convert to numpy arrays
								lowercase__  					=     np.array(lowerCamelCase    )
								lowercase__  					=     np.array(lowerCamelCase    )

								# Assert that arrays are 2D
								if len(X.shape    ) != 2:
												raise ValueError('''Expected `X` to be a 2D vector'''    )
								if len(reference_distribution.shape    ) != 2:
												raise ValueError('''Expected `reference_distribution` to be a 2D vector'''    )
								if reference_distribution.shape[0] < 2:
												raise ValueError(
												    '''Expected `reference_distribution` to be a 2D vector with more than one element in the first dimension'''    )

								# Get mahalanobis distance for each prediction
								lowercase__  					=     X - np.mean(lowerCamelCase    )
								lowercase__  					=     np.cov(reference_distribution.T    )
								try:
												lowercase__  					=     np.linalg.inv(lowerCamelCase    )
								except np.linalg.LinAlgError:
												lowercase__  					=     np.linalg.pinv(lowerCamelCase    )
								lowercase__  					=     np.dot(lowerCamelCase,    lowerCamelCase    )
								lowercase__  					=     np.dot(lowerCamelCase,    X_minus_mu.T    ).diagonal()

								return {"mahalanobis": mahal_dist}




