
import unittest

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

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
							import jax

							from transformers.models.big_bird.modeling_flax_big_bird import (
							    FlaxBigBirdForCausalLM,
							    FlaxBigBirdForMaskedLM,
							    FlaxBigBirdForMultipleChoice,
							    FlaxBigBirdForPreTraining,
							    FlaxBigBirdForQuestionAnswering,
							    FlaxBigBirdForSequenceClassification,
							    FlaxBigBirdForTokenClassification,
							    FlaxBigBirdModel,
							)




class   __lowercase						(						unittest.TestCase       ):






	def __init__(self				: str      ,	snake_case				: Tuple      ,	snake_case				: Optional[Any]=2      ,	snake_case				: Dict=56      ,	snake_case				: Any=True      ,	snake_case				: Dict=True      ,	snake_case				: Union[str, Any]=True      ,	snake_case				: int=True      ,	snake_case				: Dict=99      ,	snake_case				: List[str]=32      ,	snake_case				: int=2      ,	snake_case				: Dict=2      ,	snake_case				: Dict=7      ,	snake_case				: int="gelu_new"      ,	snake_case				: int=0.1      ,	snake_case				: List[str]=0.1      ,	snake_case				: Any=512      ,	snake_case				: List[Any]=16      ,	snake_case				: Union[str, Any]=2      ,	snake_case				: Dict=0.02      ,	snake_case				: int=4      ,	snake_case				: List[str]="block_sparse"      ,	snake_case				: List[Any]=True      ,	snake_case				: Any=False      ,	snake_case				: Any=2      ,	snake_case				: List[str]=3      ,	) ->			str:
							_lowercase       :	List[Any]           =   parent
							_lowercase       :	Union[str, Any]           =   batch_size
							_lowercase       :	str           =   seq_length
							_lowercase       :	int           =   is_training
							_lowercase       :	Tuple           =   use_attention_mask
							_lowercase       :	Optional[int]           =   use_token_type_ids
							_lowercase       :	Dict           =   use_labels
							_lowercase       :	Any           =   vocab_size
							_lowercase       :	str           =   hidden_size
							_lowercase       :	Union[str, Any]           =   num_hidden_layers
							_lowercase       :	int           =   num_attention_heads
							_lowercase       :	Union[str, Any]           =   intermediate_size
							_lowercase       :	List[str]           =   hidden_act
							_lowercase       :	List[str]           =   hidden_dropout_prob
							_lowercase       :	Union[str, Any]           =   attention_probs_dropout_prob
							_lowercase       :	Dict           =   max_position_embeddings
							_lowercase       :	Optional[Any]           =   type_vocab_size
							_lowercase       :	List[Any]           =   type_sequence_label_size
							_lowercase       :	Union[str, Any]           =   initializer_range
							_lowercase       :	Optional[int]           =   num_choices

							_lowercase       :	Dict           =   rescale_embeddings
							_lowercase       :	Union[str, Any]           =   attention_type
							_lowercase       :	List[Any]           =   use_bias
							_lowercase       :	str           =   block_size
							_lowercase       :	List[Any]           =   num_random_blocks






	def       _a(self				: List[str]    ) ->			str:
							_lowercase       :	Any           =   ids_tensor([self.batch_size, self.seq_length]      ,	self.vocab_size    )

							_lowercase       :	Union[str, Any]           =   None
							if self.use_attention_mask:
													_lowercase       :	Optional[Any]           =   random_attention_mask([self.batch_size, self.seq_length]    )

							_lowercase       :	Tuple           =   None
							if self.use_token_type_ids:
													_lowercase       :	List[Any]           =   ids_tensor([self.batch_size, self.seq_length]      ,	self.type_vocab_size    )

							_lowercase       :	Any           =   BigBirdConfig(
							    vocab_size=self.vocab_size      ,	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      ,	max_position_embeddings=self.max_position_embeddings      ,	type_vocab_size=self.type_vocab_size      ,	is_decoder=snake_case      ,	initializer_range=self.initializer_range      ,	attention_type=self.attention_type      ,	block_size=self.block_size      ,	num_random_blocks=self.num_random_blocks      ,	use_bias=self.use_bias      ,	rescale_embeddings=self.rescale_embeddings      ,	)

							return config, input_ids, token_type_ids, attention_mask






	def       _a(self				: str    ) ->			Union[str, Any]:
							_lowercase       :	Tuple           =   self.prepare_config_and_inputs()
							_lowercase      ,    _lowercase      ,    _lowercase      ,    _lowercase       :	str           =   config_and_inputs
							_lowercase       :	Union[str, Any]           =   {
							    "input_ids": input_ids,
							    "token_type_ids": token_type_ids,
							    "attention_mask": attention_mask,
							}
							return config, inputs_dict




@require_flax
class   __lowercase						(						__snake_case   , unittest.TestCase       ):
	_A       =   (
	    (
	        FlaxBigBirdForCausalLM,
	        FlaxBigBirdModel,
	        FlaxBigBirdForPreTraining,
	        FlaxBigBirdForMaskedLM,
	        FlaxBigBirdForMultipleChoice,
	        FlaxBigBirdForQuestionAnswering,
	        FlaxBigBirdForSequenceClassification,
	        FlaxBigBirdForTokenClassification,
	    )
	    if is_flax_available()
	    else ()
	)

	_A       =   False
	_A       =   False






	def       _a(self				: Tuple    ) ->			Union[str, Any]:
							_lowercase       :	List[str]           =   FlaxBigBirdModelTester(self    )






	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def       _a(self				: Any    ) ->			List[Any]:
							super().test_from_pretrained_save_pretrained()






	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def       _a(self				: Optional[int]    ) ->			int:
							super().test_from_pretrained_with_no_automatic_init()






	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def       _a(self				: List[str]    ) ->			List[Any]:
							super().test_no_automatic_init()






	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def       _a(self				: Optional[Any]    ) ->			Optional[Any]:
							super().test_hidden_states_output()






	@slow
	def       _a(self				: Tuple    ) ->			List[Any]:
							for model_class_name in self.all_model_classes:
													_lowercase       :	Union[str, Any]           =   model_class_name.from_pretrained("google/bigbird-roberta-base"    )
													self.assertIsNotNone(snake_case    )






	def       _a(self				: Any    ) ->			Any:
							if self.test_attn_probs:
													super().test_attention_outputs()






	@slow
	# copied from `test_modeling_flax_common` because it takes much longer than other models
	def       _a(self				: Optional[Any]    ) ->			List[str]:
							_lowercase      ,    _lowercase       :	Any           =   self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
													with self.subTest(model_class.__name__    ):
																			_lowercase       :	int           =   self._prepare_for_class(snake_case      ,	snake_case    )
																			_lowercase       :	Any           =   model_class(snake_case    )

																			@jax.jit
																			def model_jitted(snake_case				: Optional[Any]      ,	snake_case				: List[str]=None      ,	**snake_case				: Dict    ):
																									return model(input_ids=snake_case      ,	attention_mask=snake_case      ,	**snake_case    )

																			with self.subTest("JIT Enabled"    ):
																									_lowercase       :	int           =   model_jitted(**snake_case    ).to_tuple()

																			with self.subTest("JIT Disabled"    ):
																									with jax.disable_jit():
																															_lowercase       :	int           =   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       _a(self				: Dict      ,	snake_case				: str      ,	snake_case				: str      ,	snake_case				: List[str]      ,	snake_case				: List[Any]=1e-5      ,	snake_case				: int="outputs"      ,	snake_case				: List[Any]=None    ) ->			Dict:
							# `bigbird_block_sparse_attention` in `FlaxBigBird` returns `attention_probs = None`, while in PyTorch version,
							# an effort was done to return `attention_probs` (yet to be verified).
							if name.startswith("outputs.attentions"    ):
													return
							else:
													super().check_pt_flax_outputs(snake_case      ,	snake_case      ,	snake_case      ,	snake_case      ,	snake_case      ,	snake_case    )






from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, logging


if is_torch_available():
							import torch


UpperCamelCase_						:     Tuple	       =							logging.get_logger(__name__)




class   __lowercase						(						__snake_case       ):
	_A       =   ["pixel_values"]






	def __init__(self				: List[str]      ,	snake_case				: bool = True      ,	snake_case				: Optional[Dict[str, int]] = None      ,	snake_case				: PILImageResampling = PILImageResampling.BILINEAR      ,	snake_case				: bool = True      ,	snake_case				: Dict[str, int] = None      ,	snake_case				: bool = True      ,	snake_case				: Union[int, float] = 1 / 255      ,	snake_case				: bool = True      ,	snake_case				: Optional[Union[float, List[float]]] = None      ,	snake_case				: Optional[Union[float, List[float]]] = None      ,	**snake_case				: List[Any]      ,	) ->			None:
							super().__init__(**snake_case    )
							_lowercase       :	List[str]           =   size if size is not None else {"shortest_edge": 256}
							_lowercase       :	Union[str, Any]           =   get_size_dict(snake_case      ,	default_to_square=snake_case    )
							_lowercase       :	List[str]           =   crop_size if crop_size is not None else {"height": 224, "width": 224}
							_lowercase       :	Dict           =   get_size_dict(snake_case      ,	param_name="crop_size"    )
							_lowercase       :	List[Any]           =   do_resize
							_lowercase       :	Optional[Any]           =   size
							_lowercase       :	Tuple           =   resample
							_lowercase       :	Tuple           =   do_center_crop
							_lowercase       :	Any           =   crop_size
							_lowercase       :	str           =   do_rescale
							_lowercase       :	int           =   rescale_factor
							_lowercase       :	List[Any]           =   do_normalize
							_lowercase       :	List[str]           =   image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
							_lowercase       :	Optional[Any]           =   image_std if image_std is not None else IMAGENET_STANDARD_STD






	def       _a(self				: Tuple      ,	snake_case				: np.ndarray      ,	snake_case				: Dict[str, int]      ,	snake_case				: PILImageResampling = PILImageResampling.BICUBIC      ,	snake_case				: Optional[Union[str, ChannelDimension]] = None      ,	**snake_case				: int      ,	) ->			np.ndarray:
							_lowercase       :	Union[str, Any]           =   get_size_dict(snake_case      ,	default_to_square=snake_case    )
							if "shortest_edge" not in size:
													raise ValueError(F"""The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}"""    )
							_lowercase       :	int           =   get_resize_output_image_size(snake_case      ,	size=size["shortest_edge"]      ,	default_to_square=snake_case    )
							return resize(snake_case      ,	size=snake_case      ,	resample=snake_case      ,	data_format=snake_case      ,	**snake_case    )






	def       _a(self				: str      ,	snake_case				: np.ndarray      ,	snake_case				: Dict[str, int]      ,	snake_case				: Optional[Union[str, ChannelDimension]] = None      ,	**snake_case				: Union[str, Any]      ,	) ->			np.ndarray:
							_lowercase       :	Union[str, Any]           =   get_size_dict(snake_case    )
							if "height" not in size or "width" not in size:
													raise ValueError(F"""The `size` parameter must contain the keys `height` and `width`. Got {size.keys()}"""    )
							return center_crop(snake_case      ,	size=(size["height"], size["width"])      ,	data_format=snake_case      ,	**snake_case    )






	def       _a(self				: Union[str, Any]      ,	snake_case				: np.ndarray      ,	snake_case				: float      ,	snake_case				: Optional[Union[str, ChannelDimension]] = None      ,	**snake_case				: List[Any]    ) ->			np.ndarray:
							return rescale(snake_case      ,	scale=snake_case      ,	data_format=snake_case      ,	**snake_case    )






	def       _a(self				: int      ,	snake_case				: np.ndarray      ,	snake_case				: Union[float, List[float]]      ,	snake_case				: Union[float, List[float]]      ,	snake_case				: Optional[Union[str, ChannelDimension]] = None      ,	**snake_case				: Tuple      ,	) ->			np.ndarray:
							return normalize(snake_case      ,	mean=snake_case      ,	std=snake_case      ,	data_format=snake_case      ,	**snake_case    )






	def       _a(self				: Optional[int]      ,	snake_case				: ImageInput      ,	snake_case				: Optional[bool] = None      ,	snake_case				: Dict[str, int] = None      ,	snake_case				: PILImageResampling = None      ,	snake_case				: bool = None      ,	snake_case				: Dict[str, int] = None      ,	snake_case				: Optional[bool] = None      ,	snake_case				: Optional[float] = None      ,	snake_case				: Optional[bool] = None      ,	snake_case				: Optional[Union[float, List[float]]] = None      ,	snake_case				: Optional[Union[float, List[float]]] = None      ,	snake_case				: Optional[Union[str, TensorType]] = None      ,	snake_case				: Union[str, ChannelDimension] = ChannelDimension.FIRST      ,	**snake_case				: Tuple      ,	) ->			Union[str, Any]:
							_lowercase       :	Any           =   do_resize if do_resize is not None else self.do_resize
							_lowercase       :	List[str]           =   size if size is not None else self.size
							_lowercase       :	Optional[int]           =   get_size_dict(snake_case      ,	default_to_square=snake_case    )
							_lowercase       :	Tuple           =   resample if resample is not None else self.resample
							_lowercase       :	List[Any]           =   do_center_crop if do_center_crop is not None else self.do_center_crop
							_lowercase       :	int           =   crop_size if crop_size is not None else self.crop_size
							_lowercase       :	Dict           =   get_size_dict(snake_case      ,	param_name="crop_size"    )
							_lowercase       :	Dict           =   do_rescale if do_rescale is not None else self.do_rescale
							_lowercase       :	List[Any]           =   rescale_factor if rescale_factor is not None else self.rescale_factor
							_lowercase       :	Union[str, Any]           =   do_normalize if do_normalize is not None else self.do_normalize
							_lowercase       :	Dict           =   image_mean if image_mean is not None else self.image_mean
							_lowercase       :	Optional[Any]           =   image_std if image_std is not None else self.image_std

							_lowercase       :	int           =   make_list_of_images(snake_case    )

							if not valid_images(snake_case    ):
													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."    )

							# All transformations expect numpy arrays.
							_lowercase       :	Tuple           =   [to_numpy_array(snake_case    ) for image in images]

							if do_resize:
													_lowercase       :	List[str]           =   [self.resize(image=snake_case      ,	size=snake_case      ,	resample=snake_case    ) for image in images]

							if do_center_crop:
													_lowercase       :	List[str]           =   [self.center_crop(image=snake_case      ,	size=snake_case    ) for image in images]

							if do_rescale:
													_lowercase       :	Any           =   [self.rescale(image=snake_case      ,	scale=snake_case    ) for image in images]

							if do_normalize:
													_lowercase       :	Optional[Any]           =   [self.normalize(image=snake_case      ,	mean=snake_case      ,	std=snake_case    ) for image in images]

							_lowercase       :	Any           =   [to_channel_dimension_format(snake_case      ,	snake_case    ) for image in images]

							_lowercase       :	Optional[Any]           =   {"pixel_values": images}
							return BatchFeature(data=snake_case      ,	tensor_type=snake_case    )






	def       _a(self				: Dict      ,	snake_case				: List[str]      ,	snake_case				: List[Tuple] = None    ) ->			Optional[Any]:
							_lowercase       :	Dict           =   outputs.logits

							# Resize logits and compute semantic segmentation maps
							if target_sizes is not None:
													if len(snake_case    ) != len(snake_case    ):
																			raise ValueError(
																			    "Make sure that you pass in as many target sizes as the batch dimension of the logits"    )

													if is_torch_tensor(snake_case    ):
																			_lowercase       :	Dict           =   target_sizes.numpy()

													_lowercase       :	Tuple           =   []

													for idx in range(len(snake_case    )    ):
																			_lowercase       :	List[str]           =   torch.nn.functional.interpolate(
																			    logits[idx].unsqueeze(dim=0    )      ,	size=target_sizes[idx]      ,	mode="bilinear"      ,	align_corners=snake_case    )
																			_lowercase       :	int           =   resized_logits[0].argmax(dim=0    )
																			semantic_segmentation.append(snake_case    )
							else:
													_lowercase       :	Optional[int]           =   logits.argmax(dim=1    )
													_lowercase       :	int           =   [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]    )]

							return semantic_segmentation




