






from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


UpperCamelCase			    =     logging.get_logger(__name__)


if is_vision_available():
		import PIL



class 				__UpperCAmelCase		(_UpperCAmelCase   ):
			__snake_case    :				List[Any]									=							["pixel_values"]







			def __init__(     self:  Optional[Any]				,     UpperCAmelCase_:  bool = True				,     UpperCAmelCase_:  Dict[str, int] = None				,     UpperCAmelCase_:  PILImageResampling = PILImageResampling.BICUBIC				,     UpperCAmelCase_:  bool = True				,     UpperCAmelCase_:  Dict[str, int] = None				,     UpperCAmelCase_:  bool = True				,     UpperCAmelCase_:  Union[int, float] = 1 / 255				,     UpperCAmelCase_:  bool = True				,     UpperCAmelCase_:  Optional[Union[float, List[float]]] = None				,     UpperCAmelCase_:  Optional[Union[float, List[float]]] = None				,     UpperCAmelCase_:  bool = True				,     **UpperCAmelCase_:  Any				,     ):







									'''simple docstring'''




									super().__init__(**UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  size if size is not None else {"""shortest_edge""": 224}
									_SCREAMING_SNAKE_CASE     				=  get_size_dict(UpperCAmelCase_				,     default_to_square=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  crop_size if crop_size is not None else {"""height""": 224, """width""": 224}
									_SCREAMING_SNAKE_CASE     				=  get_size_dict(UpperCAmelCase_				,     default_to_square=UpperCAmelCase_				,     param_name="""crop_size"""					)

									_SCREAMING_SNAKE_CASE     				=  do_resize
									_SCREAMING_SNAKE_CASE     				=  size
									_SCREAMING_SNAKE_CASE     				=  resample
									_SCREAMING_SNAKE_CASE     				=  do_center_crop
									_SCREAMING_SNAKE_CASE     				=  crop_size
									_SCREAMING_SNAKE_CASE     				=  do_rescale
									_SCREAMING_SNAKE_CASE     				=  rescale_factor
									_SCREAMING_SNAKE_CASE     				=  do_normalize
									_SCREAMING_SNAKE_CASE     				=  image_mean if image_mean is not None else OPENAI_CLIP_MEAN
									_SCREAMING_SNAKE_CASE     				=  image_std if image_std is not None else OPENAI_CLIP_STD
									_SCREAMING_SNAKE_CASE     				=  do_convert_rgb







			def 							UpperCamelCase				(     self:  List[Any]				,     UpperCAmelCase_:  np.ndarray				,     UpperCAmelCase_:  Dict[str, int]				,     UpperCAmelCase_:  PILImageResampling = PILImageResampling.BICUBIC				,     UpperCAmelCase_:  Optional[Union[str, ChannelDimension]] = None				,     **UpperCAmelCase_:  Tuple				,     ):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  get_size_dict(UpperCAmelCase_				,     default_to_square=UpperCAmelCase_					)
									if "shortest_edge" not in size:
															raise ValueError(F'The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}'					)
									_SCREAMING_SNAKE_CASE     				=  get_resize_output_image_size(UpperCAmelCase_				,     size=size["""shortest_edge"""]				,     default_to_square=UpperCAmelCase_					)
									return resize(UpperCAmelCase_				,     size=UpperCAmelCase_				,     resample=UpperCAmelCase_				,     data_format=UpperCAmelCase_				,     **UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Dict				,     UpperCAmelCase_:  np.ndarray				,     UpperCAmelCase_:  Dict[str, int]				,     UpperCAmelCase_:  Optional[Union[str, ChannelDimension]] = None				,     **UpperCAmelCase_:  List[str]				,     ):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  get_size_dict(UpperCAmelCase_					)
									if "height" not in size or "width" not in size:
															raise ValueError(F'The `size` parameter must contain the keys (height, width). Got {size.keys()}'					)
									return center_crop(UpperCAmelCase_				,     size=(size["""height"""], size["""width"""])				,     data_format=UpperCAmelCase_				,     **UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Union[str, Any]				,     UpperCAmelCase_:  np.ndarray				,     UpperCAmelCase_:  Union[int, float]				,     UpperCAmelCase_:  Optional[Union[str, ChannelDimension]] = None				,     **UpperCAmelCase_:  Union[str, Any]				,     ):







									'''simple docstring'''




									return rescale(UpperCAmelCase_				,     scale=UpperCAmelCase_				,     data_format=UpperCAmelCase_				,     **UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  int				,     UpperCAmelCase_:  np.ndarray				,     UpperCAmelCase_:  Union[float, List[float]]				,     UpperCAmelCase_:  Union[float, List[float]]				,     UpperCAmelCase_:  Optional[Union[str, ChannelDimension]] = None				,     **UpperCAmelCase_:  Union[str, Any]				,     ):







									'''simple docstring'''




									return normalize(UpperCAmelCase_				,     mean=UpperCAmelCase_				,     std=UpperCAmelCase_				,     data_format=UpperCAmelCase_				,     **UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Union[str, Any]				,     UpperCAmelCase_:  ImageInput				,     UpperCAmelCase_:  bool = None				,     UpperCAmelCase_:  Dict[str, int] = None				,     UpperCAmelCase_:  PILImageResampling = None				,     UpperCAmelCase_:  bool = None				,     UpperCAmelCase_:  int = None				,     UpperCAmelCase_:  bool = None				,     UpperCAmelCase_:  float = None				,     UpperCAmelCase_:  bool = None				,     UpperCAmelCase_:  Optional[Union[float, List[float]]] = None				,     UpperCAmelCase_:  Optional[Union[float, List[float]]] = None				,     UpperCAmelCase_:  bool = None				,     UpperCAmelCase_:  Optional[Union[str, TensorType]] = None				,     UpperCAmelCase_:  Optional[ChannelDimension] = ChannelDimension.FIRST				,     **UpperCAmelCase_:  List[Any]				,     ):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  do_resize if do_resize is not None else self.do_resize
									_SCREAMING_SNAKE_CASE     				=  size if size is not None else self.size
									_SCREAMING_SNAKE_CASE     				=  get_size_dict(UpperCAmelCase_				,     param_name="""size"""				,     default_to_square=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  resample if resample is not None else self.resample
									_SCREAMING_SNAKE_CASE     				=  do_center_crop if do_center_crop is not None else self.do_center_crop
									_SCREAMING_SNAKE_CASE     				=  crop_size if crop_size is not None else self.crop_size
									_SCREAMING_SNAKE_CASE     				=  get_size_dict(UpperCAmelCase_				,     param_name="""crop_size"""				,     default_to_square=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  do_rescale if do_rescale is not None else self.do_rescale
									_SCREAMING_SNAKE_CASE     				=  rescale_factor if rescale_factor is not None else self.rescale_factor
									_SCREAMING_SNAKE_CASE     				=  do_normalize if do_normalize is not None else self.do_normalize
									_SCREAMING_SNAKE_CASE     				=  image_mean if image_mean is not None else self.image_mean
									_SCREAMING_SNAKE_CASE     				=  image_std if image_std is not None else self.image_std
									_SCREAMING_SNAKE_CASE     				=  do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

									_SCREAMING_SNAKE_CASE     				=  make_list_of_images(UpperCAmelCase_					)

									if not valid_images(UpperCAmelCase_					):
															raise ValueError(
															    """Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, """
															    """torch.Tensor, tf.Tensor or jax.ndarray."""					)

									if do_resize and size is None:
															raise ValueError("""Size must be specified if do_resize is True."""					)

									if do_center_crop and crop_size is None:
															raise ValueError("""Crop size must be specified if do_center_crop is True."""					)

									if do_rescale and rescale_factor is None:
															raise ValueError("""Rescale factor must be specified if do_rescale is True."""					)

									if do_normalize and (image_mean is None or image_std is None):
															raise ValueError("""Image mean and std must be specified if do_normalize is True."""					)

									# PIL RGBA images are converted to RGB
									if do_convert_rgb:
															_SCREAMING_SNAKE_CASE     				=  [convert_to_rgb(UpperCAmelCase_					) for image in images]

									# All transformations expect numpy arrays.
									_SCREAMING_SNAKE_CASE     				=  [to_numpy_array(UpperCAmelCase_					) for image in images]

									if do_resize:
															_SCREAMING_SNAKE_CASE     				=  [self.resize(image=UpperCAmelCase_				,     size=UpperCAmelCase_				,     resample=UpperCAmelCase_					) for image in images]

									if do_center_crop:
															_SCREAMING_SNAKE_CASE     				=  [self.center_crop(image=UpperCAmelCase_				,     size=UpperCAmelCase_					) for image in images]

									if do_rescale:
															_SCREAMING_SNAKE_CASE     				=  [self.rescale(image=UpperCAmelCase_				,     scale=UpperCAmelCase_					) for image in images]

									if do_normalize:
															_SCREAMING_SNAKE_CASE     				=  [self.normalize(image=UpperCAmelCase_				,     mean=UpperCAmelCase_				,     std=UpperCAmelCase_					) for image in images]

									_SCREAMING_SNAKE_CASE     				=  [to_channel_dimension_format(UpperCAmelCase_				,     UpperCAmelCase_					) for image in images]

									_SCREAMING_SNAKE_CASE     				=  {"""pixel_values""": images}
									return BatchFeature(data=UpperCAmelCase_				,     tensor_type=UpperCAmelCase_					)









import tempfile
import unittest

import numpy as np
from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_flax


if is_flax_available():
		import os

		from flax.core.frozen_dict import unfreeze
		from flax.traverse_util import flatten_dict

		from transformers import FlaxBertModel

		UpperCamelCase			    =     '''0.12'''  # assumed parallelism: 8



@require_flax
@is_staging_test
class 				__UpperCAmelCase		(unittest.TestCase   ):







			@classmethod
			def 							UpperCamelCase				(     cls:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TOKEN
									HfFolder.save_token(UpperCAmelCase_					)







			@classmethod
			def 							UpperCamelCase				(     cls:  Union[str, Any]					):







									'''simple docstring'''




									try:
															delete_repo(token=cls._token				,     repo_id="""test-model-flax"""					)
									except HTTPError:
															pass

									try:
															delete_repo(token=cls._token				,     repo_id="""valid_org/test-model-flax-org"""					)
									except HTTPError:
															pass







			def 							UpperCamelCase				(     self:  str					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  BertConfig(
									    vocab_size=99				,     hidden_size=32				,     num_hidden_layers=5				,     num_attention_heads=4				,     intermediate_size=37					)
									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel(UpperCAmelCase_					)
									model.push_to_hub("""test-model-flax"""				,     use_auth_token=self._token					)

									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(F'{USER}/test-model-flax'					)

									_SCREAMING_SNAKE_CASE     				=  flatten_dict(unfreeze(model.params					)					)
									_SCREAMING_SNAKE_CASE     				=  flatten_dict(unfreeze(new_model.params					)					)

									for key in base_params.keys():
															_SCREAMING_SNAKE_CASE     				=  (base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(UpperCAmelCase_				,     1E-3				,     msg=F'{key} not identical'					)

									# Reset repo
									delete_repo(token=self._token				,     repo_id="""test-model-flax"""					)

									# Push to hub via save_pretrained
									with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(UpperCAmelCase_				,     repo_id="""test-model-flax"""				,     push_to_hub=UpperCAmelCase_				,     use_auth_token=self._token					)

									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(F'{USER}/test-model-flax'					)

									_SCREAMING_SNAKE_CASE     				=  flatten_dict(unfreeze(model.params					)					)
									_SCREAMING_SNAKE_CASE     				=  flatten_dict(unfreeze(new_model.params					)					)

									for key in base_params.keys():
															_SCREAMING_SNAKE_CASE     				=  (base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(UpperCAmelCase_				,     1E-3				,     msg=F'{key} not identical'					)







			def 							UpperCamelCase				(     self:  int					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  BertConfig(
									    vocab_size=99				,     hidden_size=32				,     num_hidden_layers=5				,     num_attention_heads=4				,     intermediate_size=37					)
									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel(UpperCAmelCase_					)
									model.push_to_hub("""valid_org/test-model-flax-org"""				,     use_auth_token=self._token					)

									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained("""valid_org/test-model-flax-org"""					)

									_SCREAMING_SNAKE_CASE     				=  flatten_dict(unfreeze(model.params					)					)
									_SCREAMING_SNAKE_CASE     				=  flatten_dict(unfreeze(new_model.params					)					)

									for key in base_params.keys():
															_SCREAMING_SNAKE_CASE     				=  (base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(UpperCAmelCase_				,     1E-3				,     msg=F'{key} not identical'					)

									# Reset repo
									delete_repo(token=self._token				,     repo_id="""valid_org/test-model-flax-org"""					)

									# Push to hub via save_pretrained
									with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(
															    UpperCAmelCase_				,     repo_id="""valid_org/test-model-flax-org"""				,     push_to_hub=UpperCAmelCase_				,     use_auth_token=self._token					)

									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained("""valid_org/test-model-flax-org"""					)

									_SCREAMING_SNAKE_CASE     				=  flatten_dict(unfreeze(model.params					)					)
									_SCREAMING_SNAKE_CASE     				=  flatten_dict(unfreeze(new_model.params					)					)

									for key in base_params.keys():
															_SCREAMING_SNAKE_CASE     				=  (base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(UpperCAmelCase_				,     1E-3				,     msg=F'{key} not identical'					)

def 	__lowerCamelCase      (		snake_case__	,snake_case__		)							->  Tuple:






						"""simple docstring"""





						_SCREAMING_SNAKE_CASE     				=  True
						_SCREAMING_SNAKE_CASE     				=  flatten_dict(modela.params		)
						_SCREAMING_SNAKE_CASE     				=  flatten_dict(modela.params		)
						for key in flat_params_a.keys():
												if np.sum(np.abs(flat_params_a[key] - flat_params_a[key]		)		) > 1e-4:
																		_SCREAMING_SNAKE_CASE     				=  False

						return models_are_equal




@require_flax
class 				__UpperCAmelCase		(unittest.TestCase   ):







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  BertConfig.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""					)
									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  """bert"""
									with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(os.path.join(UpperCAmelCase_				,     UpperCAmelCase_					)					)

															with self.assertRaises(UpperCAmelCase_					):
																					_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_					)

															_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_				,     subfolder=UpperCAmelCase_					)

									self.assertTrue(check_models_equal(UpperCAmelCase_				,     UpperCAmelCase_					)					)







			def 							UpperCamelCase				(     self:  List[str]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  BertConfig.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""					)
									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  """bert"""
									with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(os.path.join(UpperCAmelCase_				,     UpperCAmelCase_					)				,     max_shard_size="""10KB"""					)

															with self.assertRaises(UpperCAmelCase_					):
																					_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_					)

															_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_				,     subfolder=UpperCAmelCase_					)

									self.assertTrue(check_models_equal(UpperCAmelCase_				,     UpperCAmelCase_					)					)







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  """bert"""
									_SCREAMING_SNAKE_CASE     				=  """hf-internal-testing/tiny-random-bert-subfolder"""

									with self.assertRaises(UpperCAmelCase_					):
															_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_				,     subfolder=UpperCAmelCase_					)

									self.assertIsNotNone(UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Optional[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  """bert"""
									_SCREAMING_SNAKE_CASE     				=  """hf-internal-testing/tiny-random-bert-sharded-subfolder"""
									with self.assertRaises(UpperCAmelCase_					):
															_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_				,     subfolder=UpperCAmelCase_					)

									self.assertIsNotNone(UpperCAmelCase_					)

