



import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import VideoMAEConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch
 from torch import nn

 from transformers import (
     MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING,
     VideoMAEForPreTraining,
     VideoMAEForVideoClassification,
     VideoMAEModel,
 )
 from transformers.models.videomae.modeling_videomae import VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from transformers import VideoMAEImageProcessor




class 					snake_case__						:






     """simple docstring"""


     def __init__(       self  :  List[Any], _snake_case  :  str, _snake_case  :  Optional[int]=1_3, _snake_case  :  Optional[int]=1_0, _snake_case  :  Tuple=3, _snake_case  :  Tuple=2, _snake_case  :  Union[str, Any]=2, _snake_case  :  Any=2, _snake_case  :  Optional[Any]=True, _snake_case  :  Union[str, Any]=True, _snake_case  :  Dict=3_2, _snake_case  :  Any=5, _snake_case  :  Any=4, _snake_case  :  Optional[Any]=3_7, _snake_case  :  str="gelu", _snake_case  :  Tuple=0.1, _snake_case  :  List[str]=0.1, _snake_case  :  Optional[Any]=1_0, _snake_case  :  Tuple=0.0_2, _snake_case  :  str=0.9, _snake_case  :  int=None, )   ->List[Any]:
            snake_case__							:	str   	=	parent
            snake_case__							:	int   	=	batch_size
            snake_case__							:	Dict   	=	image_size
            snake_case__							:	Any   	=	num_channels
            snake_case__							:	int   	=	patch_size
            snake_case__							:	int   	=	tubelet_size
            snake_case__							:	Optional[int]   	=	num_frames
            snake_case__							:	int   	=	is_training
            snake_case__							:	Dict   	=	use_labels
            snake_case__							:	Union[str, Any]   	=	hidden_size
            snake_case__							:	List[Any]   	=	num_hidden_layers
            snake_case__							:	Optional[Any]   	=	num_attention_heads
            snake_case__							:	str   	=	intermediate_size
            snake_case__							:	str   	=	hidden_act
            snake_case__							:	List[Any]   	=	hidden_dropout_prob
            snake_case__							:	Optional[Any]   	=	attention_probs_dropout_prob
            snake_case__							:	Optional[int]   	=	type_sequence_label_size
            snake_case__							:	Dict   	=	initializer_range
            snake_case__							:	Any   	=	mask_ratio
            snake_case__							:	Dict   	=	scope

            # in VideoMAE, the number of tokens equals num_frames/tubelet_size * num_patches per frame
            snake_case__							:	Any   	=	(image_size // patch_size) ** 2
            snake_case__							:	Dict   	=	(num_frames // tubelet_size) * self.num_patches_per_frame

            # use this variable to define bool_masked_pos
            snake_case__							:	Dict   	=	int(mask_ratio * self.seq_length					)






     def    lowercase_   (       self  :  Optional[int]					)   ->int:
            snake_case__							:	Union[str, Any]   	=	floats_tensor(
                [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]					)

            snake_case__							:	Optional[Any]   	=	None
            if self.use_labels:
                   snake_case__							:	List[str]   	=	ids_tensor([self.batch_size], self.type_sequence_label_size					)

            snake_case__							:	Tuple   	=	self.get_config()

            return config, pixel_values, labels






     def    lowercase_   (       self  :  int					)   ->Optional[Any]:
            return VideoMAEConfig(
                image_size=self.image_size, patch_size=self.patch_size, num_channels=self.num_channels, num_frames=self.num_frames, tubelet_size=self.tubelet_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, is_decoder=lowerCAmelCase__, initializer_range=self.initializer_range, )






     def    lowercase_   (       self  :  str, _snake_case  :  Union[str, Any], _snake_case  :  str, _snake_case  :  Optional[Any]					)   ->Tuple:
            snake_case__							:	Optional[Any]   	=	VideoMAEModel(config=lowerCAmelCase__					)
            model.to(lowerCAmelCase__					)
            model.eval()
            snake_case__							:	Union[str, Any]   	=	model(lowerCAmelCase__					)
            self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size)					)






     def    lowercase_   (       self  :  str, _snake_case  :  Tuple, _snake_case  :  Any, _snake_case  :  Union[str, Any]					)   ->List[Any]:
            snake_case__							:	List[Any]   	=	VideoMAEForPreTraining(lowerCAmelCase__					)
            model.to(lowerCAmelCase__					)
            model.eval()
            # important: each video needs to have the same number of masked patches
            # hence we define a single mask, which we then repeat for each example in the batch
            snake_case__							:	Dict   	=	torch.ones((self.num_masks,)					)
            snake_case__							:	Union[str, Any]   	=	torch.cat([mask, torch.zeros(self.seq_length - mask.size(0					)					)]					)
            snake_case__							:	List[str]   	=	mask.expand(self.batch_size, -1					).bool()

            snake_case__							:	List[Any]   	=	model(lowerCAmelCase__, lowerCAmelCase__					)
            # model only returns predictions for masked patches
            snake_case__							:	Optional[int]   	=	mask.sum().item()
            snake_case__							:	str   	=	3 * self.tubelet_size * self.patch_size**2
            self.parent.assertEqual(result.logits.shape, (self.batch_size, num_masked_patches, decoder_num_labels)					)






     def    lowercase_   (       self  :  str					)   ->Tuple:
            snake_case__							:	Dict   	=	self.prepare_config_and_inputs()
            snake_case__							:	Any   	=	config_and_inputs
            snake_case__							:	int   	=	{"pixel_values": pixel_values}
            return config, inputs_dict




@require_torch
class 					snake_case__						(       UpperCAmelCase_				,     UpperCAmelCase_				,     unittest.TestCase     ):






     """simple docstring"""


     _SCREAMING_SNAKE_CASE						=	(
         (VideoMAEModel, VideoMAEForPreTraining, VideoMAEForVideoClassification) if is_torch_available() else ()
     )
     _SCREAMING_SNAKE_CASE						=	(
         {'''feature-extraction''': VideoMAEModel, '''video-classification''': VideoMAEForVideoClassification}
         if is_torch_available()
         else {}
     )

     _SCREAMING_SNAKE_CASE						=	False
     _SCREAMING_SNAKE_CASE						=	False
     _SCREAMING_SNAKE_CASE						=	False
     _SCREAMING_SNAKE_CASE						=	False






     def    lowercase_   (       self  :  Tuple					)   ->List[Any]:
            snake_case__							:	List[Any]   	=	VideoMAEModelTester(self					)
            snake_case__							:	Union[str, Any]   	=	ConfigTester(self, config_class=lowerCAmelCase__, has_text_modality=lowerCAmelCase__, hidden_size=3_7					)






     def    lowercase_   (       self  :  Tuple, _snake_case  :  int, _snake_case  :  Tuple, _snake_case  :  Dict=False					)   ->Optional[int]:
            snake_case__							:	Any   	=	copy.deepcopy(lowerCAmelCase__					)

            if model_class == VideoMAEForPreTraining:
                   # important: each video needs to have the same number of masked patches
                   # hence we define a single mask, which we then repeat for each example in the batch
                   snake_case__							:	List[str]   	=	torch.ones((self.model_tester.num_masks,)					)
                   snake_case__							:	List[str]   	=	torch.cat([mask, torch.zeros(self.model_tester.seq_length - mask.size(0					)					)]					)
                   snake_case__							:	Dict   	=	mask.expand(self.model_tester.batch_size, -1					).bool()
                   snake_case__							:	Union[str, Any]   	=	bool_masked_pos.to(lowerCAmelCase__					)

            if return_labels:
                   if model_class in [
                       *get_values(lowerCAmelCase__					),
                   ]:
                          snake_case__							:	Optional[int]   	=	torch.zeros(
                              self.model_tester.batch_size, dtype=torch.long, device=lowerCAmelCase__					)

            return inputs_dict






     def    lowercase_   (       self  :  int					)   ->str:
            self.config_tester.run_common_tests()






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






     def    lowercase_   (       self  :  Optional[Any]					)   ->Dict:
            snake_case__							:	Optional[Any]   	=	self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                   snake_case__							:	Tuple   	=	model_class(lowerCAmelCase__					)
                   self.assertIsInstance(model.get_input_embeddings(), (nn.Module)					)
                   snake_case__							:	Tuple   	=	model.get_output_embeddings()
                   self.assertTrue(x is None or isinstance(lowerCAmelCase__, nn.Linear					)					)






     def    lowercase_   (       self  :  List[Any]					)   ->Optional[int]:
            snake_case__							:	Any   	=	self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                   snake_case__							:	int   	=	model_class(lowerCAmelCase__					)
                   snake_case__							:	str   	=	inspect.signature(model.forward					)
                   # signature.parameters is an OrderedDict => so arg_names order is deterministic
                   snake_case__							:	Optional[int]   	=	[*signature.parameters.keys()]

                   snake_case__							:	int   	=	["pixel_values"]
                   self.assertListEqual(arg_names[:1], lowerCAmelCase__					)






     def    lowercase_   (       self  :  List[str]					)   ->Union[str, Any]:
            snake_case__							:	Dict   	=	self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*lowerCAmelCase__					)






     def    lowercase_   (       self  :  Tuple					)   ->int:
            snake_case__							:	Optional[int]   	=	self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_pretraining(*lowerCAmelCase__					)






     @slow
     def    lowercase_   (       self  :  List[Any]					)   ->Any:
            for model_name in VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                   snake_case__							:	Dict   	=	VideoMAEModel.from_pretrained(lowerCAmelCase__					)
                   self.assertIsNotNone(lowerCAmelCase__					)






     def    lowercase_   (       self  :  Tuple					)   ->Dict:
            if not self.has_attentions:
                   pass

            else:
                   snake_case__							:	str   	=	self.model_tester.prepare_config_and_inputs_for_common()
                   snake_case__							:	Union[str, Any]   	=	True

                   for model_class in self.all_model_classes:
                          snake_case__							:	int   	=	self.model_tester.seq_length - self.model_tester.num_masks
                          snake_case__							:	Optional[Any]   	=	(
                              num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length
                          )

                          snake_case__							:	Union[str, Any]   	=	True
                          snake_case__							:	Optional[int]   	=	False
                          snake_case__							:	Optional[Any]   	=	True
                          snake_case__							:	Optional[int]   	=	model_class(lowerCAmelCase__					)
                          model.to(lowerCAmelCase__					)
                          model.eval()
                          with torch.no_grad():
                                 snake_case__							:	Any   	=	model(**self._prepare_for_class(lowerCAmelCase__, lowerCAmelCase__					)					)
                          snake_case__							:	List[str]   	=	outputs.attentions
                          self.assertEqual(len(lowerCAmelCase__					), self.model_tester.num_hidden_layers					)

                          # check that output_attentions also work using config
                          del inputs_dict["output_attentions"]
                          snake_case__							:	Dict   	=	True
                          snake_case__							:	Union[str, Any]   	=	model_class(lowerCAmelCase__					)
                          model.to(lowerCAmelCase__					)
                          model.eval()
                          with torch.no_grad():
                                 snake_case__							:	Optional[Any]   	=	model(**self._prepare_for_class(lowerCAmelCase__, lowerCAmelCase__					)					)
                          snake_case__							:	Union[str, Any]   	=	outputs.attentions
                          self.assertEqual(len(lowerCAmelCase__					), self.model_tester.num_hidden_layers					)

                          self.assertListEqual(
                              list(attentions[0].shape[-3:]					), [self.model_tester.num_attention_heads, seq_len, seq_len], )
                          snake_case__							:	Union[str, Any]   	=	len(lowerCAmelCase__					)

                          # Check attention is always last and order is fine
                          snake_case__							:	Any   	=	True
                          snake_case__							:	Tuple   	=	True
                          snake_case__							:	List[str]   	=	model_class(lowerCAmelCase__					)
                          model.to(lowerCAmelCase__					)
                          model.eval()
                          with torch.no_grad():
                                 snake_case__							:	Dict   	=	model(**self._prepare_for_class(lowerCAmelCase__, lowerCAmelCase__					)					)

                          self.assertEqual(out_len + 1, len(lowerCAmelCase__					)					)

                          snake_case__							:	str   	=	outputs.attentions

                          self.assertEqual(len(lowerCAmelCase__					), self.model_tester.num_hidden_layers					)
                          self.assertListEqual(
                              list(self_attentions[0].shape[-3:]					), [self.model_tester.num_attention_heads, seq_len, seq_len], )






     def    lowercase_   (       self  :  Any					)   ->str:
            def check_hidden_states_output(_snake_case  :  Optional[Any], _snake_case  :  Union[str, Any], _snake_case  :  Any					):
                   snake_case__							:	Optional[int]   	=	model_class(lowerCAmelCase__					)
                   model.to(lowerCAmelCase__					)
                   model.eval()

                   with torch.no_grad():
                          snake_case__							:	Optional[Any]   	=	model(**self._prepare_for_class(lowerCAmelCase__, lowerCAmelCase__					)					)

                   snake_case__							:	Dict   	=	outputs.hidden_states
                   snake_case__							:	Dict   	=	self.model_tester.num_hidden_layers + 1
                   self.assertEqual(len(lowerCAmelCase__					), lowerCAmelCase__					)

                   snake_case__							:	Tuple   	=	self.model_tester.seq_length - self.model_tester.num_masks
                   snake_case__							:	Union[str, Any]   	=	num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length

                   self.assertListEqual(
                       list(hidden_states[0].shape[-2:]					), [seq_length, self.model_tester.hidden_size], )

            snake_case__							:	str   	=	self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                   snake_case__							:	List[str]   	=	True
                   check_hidden_states_output(lowerCAmelCase__, lowerCAmelCase__, lowerCAmelCase__					)

                   # check that output_hidden_states also work using config
                   del inputs_dict["output_hidden_states"]
                   snake_case__							:	Dict   	=	True

                   check_hidden_states_output(lowerCAmelCase__, lowerCAmelCase__, lowerCAmelCase__					)






     @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'					)
     def    lowercase_   (       self  :  List[str]					)   ->Union[str, Any]:
            pass


def        lowercase_							():
       snake_case__							:	Dict   	=	hf_hub_download(
           repo_id='hf-internal-testing/spaghetti-video'		,					filename='eating_spaghetti.npy'		,					repo_type='dataset'      )
       snake_case__							:	Optional[Any]   	=	np.load(_UpperCAmelCase      )
       return list(_UpperCAmelCase      )





@require_torch
@require_vision
class 					snake_case__						(       unittest.TestCase     ):






     """simple docstring"""


     @cached_property
     def    lowercase_   (       self  :  Tuple					)   ->Optional[Any]:
            # logits were tested with a different mean and std, so we use the same here
            return (
                VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5], image_std=[0.5, 0.5, 0.5]					)
                if is_vision_available()
                else None
            )






     @slow
     def    lowercase_   (       self  :  Any					)   ->str:
            snake_case__							:	Optional[int]   	=	VideoMAEForVideoClassification.from_pretrained('MCG-NJU/videomae-base-finetuned-kinetics'					).to(
                lowerCAmelCase__					)

            snake_case__							:	Optional[int]   	=	self.default_image_processor
            snake_case__							:	int   	=	prepare_video()
            snake_case__							:	Tuple   	=	image_processor(lowerCAmelCase__, return_tensors='pt'					).to(lowerCAmelCase__					)

            # forward pass
            with torch.no_grad():
                   snake_case__							:	Union[str, Any]   	=	model(**lowerCAmelCase__					)

            # verify the logits
            snake_case__							:	str   	=	torch.Size((1, 4_0_0)					)
            self.assertEqual(outputs.logits.shape, lowerCAmelCase__					)

            snake_case__							:	List[Any]   	=	torch.tensor([0.3_6_6_9, -0.0_6_8_8, -0.2_4_2_1]					).to(lowerCAmelCase__					)

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






     @slow
     def    lowercase_   (       self  :  Tuple					)   ->Any:
            snake_case__							:	Union[str, Any]   	=	VideoMAEForPreTraining.from_pretrained('MCG-NJU/videomae-base-short'					).to(lowerCAmelCase__					)

            snake_case__							:	Any   	=	self.default_image_processor
            snake_case__							:	Optional[int]   	=	prepare_video()
            snake_case__							:	Optional[int]   	=	image_processor(lowerCAmelCase__, return_tensors='pt'					).to(lowerCAmelCase__					)

            # add boolean mask, indicating which patches to mask
            snake_case__							:	List[Any]   	=	hf_hub_download(repo_id='hf-internal-testing/bool-masked-pos', filename='bool_masked_pos.pt'					)
            snake_case__							:	str   	=	torch.load(lowerCAmelCase__					)

            # forward pass
            with torch.no_grad():
                   snake_case__							:	Optional[int]   	=	model(**lowerCAmelCase__					)

            # verify the logits
            snake_case__							:	Optional[int]   	=	torch.Size([1, 1_4_0_8, 1_5_3_6]					)
            snake_case__							:	Optional[int]   	=	torch.tensor(
                [[0.7_9_9_4, 0.9_6_1_2, 0.8_5_0_8], [0.7_4_0_1, 0.8_9_5_8, 0.8_3_0_2], [0.5_8_6_2, 0.7_4_6_8, 0.7_3_2_5]], device=lowerCAmelCase__					)
            self.assertEqual(outputs.logits.shape, lowerCAmelCase__					)
            self.assertTrue(torch.allclose(outputs.logits[0, :3, :3], lowerCAmelCase__, atol=1e-4					)					)

            # verify the loss (`config.norm_pix_loss` = `True`)
            snake_case__							:	str   	=	torch.tensor([0.5_1_4_2], device=lowerCAmelCase__					)
            self.assertTrue(torch.allclose(outputs.loss, lowerCAmelCase__, atol=1e-4					)					)

            # verify the loss (`config.norm_pix_loss` = `False`)
            snake_case__							:	int   	=	VideoMAEForPreTraining.from_pretrained('MCG-NJU/videomae-base-short', norm_pix_loss=lowerCAmelCase__					).to(
                lowerCAmelCase__					)

            with torch.no_grad():
                   snake_case__							:	Any   	=	model(**lowerCAmelCase__					)

            snake_case__							:	Union[str, Any]   	=	torch.tensor(torch.tensor([0.6_4_6_9]					), device=lowerCAmelCase__					)
            self.assertTrue(torch.allclose(outputs.loss, lowerCAmelCase__, atol=1e-4					)					)









import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
						import jax
						import jax.numpy as jnp

						from transformers.generation import (
						    FlaxForcedBOSTokenLogitsProcessor,
						    FlaxForcedEOSTokenLogitsProcessor,
						    FlaxLogitsProcessorList,
						    FlaxMinLengthLogitsProcessor,
						    FlaxTemperatureLogitsWarper,
						    FlaxTopKLogitsWarper,
						    FlaxTopPLogitsWarper,
						)



@require_flax
class   __lowercase				( unittest.TestCase	):




			"""simple docstring"""







			def        _SCREAMING_SNAKE_CASE			( self  :     Optional[int]   ,						lowerCAmelCase__  :     int   ,						lowerCAmelCase__  :     int):
									SCREAMING_SNAKE_CASE_: str  		=       jnp.ones((batch_size, length)) / length
									return scores



			def        _SCREAMING_SNAKE_CASE			( self  :     Union[str, Any]):
									SCREAMING_SNAKE_CASE_: Dict  		=       None
									SCREAMING_SNAKE_CASE_: str  		=       20

									SCREAMING_SNAKE_CASE_: List[Any]  		=       self._get_uniform_logits(batch_size=2   ,						length=lowerCAmelCase__)

									# tweak scores to not be uniform anymore
									SCREAMING_SNAKE_CASE_: List[str]  		=       scores.at[1, 5].set((1 / length) + 0.1)  # peak, 1st batch
									SCREAMING_SNAKE_CASE_: Any  		=       scores.at[1, 10].set((1 / length) - 0.4)  # valley, 1st batch

									# compute softmax
									SCREAMING_SNAKE_CASE_: Dict  		=       jax.nn.softmax(lowerCAmelCase__   ,						axis=-1)

									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       FlaxTemperatureLogitsWarper(temperature=0.5)
									SCREAMING_SNAKE_CASE_: List[str]  		=       FlaxTemperatureLogitsWarper(temperature=1.3)

									SCREAMING_SNAKE_CASE_: str  		=       jax.nn.softmax(temp_dist_warper_sharper(lowerCAmelCase__   ,						scores.copy()   ,						cur_len=lowerCAmelCase__)   ,						axis=-1)
									SCREAMING_SNAKE_CASE_: int  		=       jax.nn.softmax(temp_dist_warper_smoother(lowerCAmelCase__   ,						scores.copy()   ,						cur_len=lowerCAmelCase__)   ,						axis=-1)

									# uniform distribution stays uniform
									self.assertTrue(jnp.allclose(probs[0, :]   ,						warped_prob_sharp[0, :]   ,						atol=1E-3))
									self.assertTrue(jnp.allclose(probs[0, :]   ,						warped_prob_smooth[0, :]   ,						atol=1E-3))

									# sharp peaks get higher, valleys get lower
									self.assertLess(probs[1, :].max()   ,						warped_prob_sharp[1, :].max())
									self.assertGreater(probs[1, :].min()   ,						warped_prob_sharp[1, :].min())

									# smooth peaks get lower, valleys get higher
									self.assertGreater(probs[1, :].max()   ,						warped_prob_smooth[1, :].max())
									self.assertLess(probs[1, :].min()   ,						warped_prob_smooth[1, :].min())



			def        _SCREAMING_SNAKE_CASE			( self  :     Optional[int]):
									SCREAMING_SNAKE_CASE_: List[str]  		=       None
									SCREAMING_SNAKE_CASE_: str  		=       10
									SCREAMING_SNAKE_CASE_: Tuple  		=       2

									# create ramp distribution
									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       np.broadcast_to(np.arange(lowerCAmelCase__)[None, :]   ,						(batch_size, vocab_size)).copy()
									SCREAMING_SNAKE_CASE_: Dict  		=       ramp_logits[1:, : vocab_size // 2] + vocab_size

									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       FlaxTopKLogitsWarper(3)

									SCREAMING_SNAKE_CASE_: Dict  		=       top_k_warp(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)

									# check that correct tokens are filtered
									self.assertListEqual(jnp.isinf(scores[0]).tolist()   ,						7 * [True] + 3 * [False])
									self.assertListEqual(jnp.isinf(scores[1]).tolist()   ,						2 * [True] + 3 * [False] + 5 * [True])

									# check special case
									SCREAMING_SNAKE_CASE_: Any  		=       5
									SCREAMING_SNAKE_CASE_: str  		=       FlaxTopKLogitsWarper(top_k=1   ,						filter_value=0.0   ,						min_tokens_to_keep=3)

									SCREAMING_SNAKE_CASE_: Any  		=       np.broadcast_to(np.arange(lowerCAmelCase__)[None, :]   ,						(batch_size, length)).copy()
									SCREAMING_SNAKE_CASE_: Any  		=       top_k_warp_safety_check(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)

									# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
									self.assertListEqual((scores == 0.0).sum(axis=-1).tolist()   ,						[2, 2])



			def        _SCREAMING_SNAKE_CASE			( self  :     int):
									SCREAMING_SNAKE_CASE_: Tuple  		=       None
									SCREAMING_SNAKE_CASE_: Dict  		=       10
									SCREAMING_SNAKE_CASE_: Dict  		=       2

									# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
									SCREAMING_SNAKE_CASE_: Tuple  		=       np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]))

									SCREAMING_SNAKE_CASE_: int  		=       FlaxTopPLogitsWarper(0.8)
									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       np.exp(top_p_warp(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__))

									# dist should be filtered to keep min num values so that sum is >= top_p
									# exp (-inf) => 0
									SCREAMING_SNAKE_CASE_: Dict  		=       np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]])
									self.assertTrue(np.allclose(lowerCAmelCase__   ,						lowerCAmelCase__   ,						atol=1E-3))

									# check edge cases with negative and extreme logits
									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       np.broadcast_to(np.arange(lowerCAmelCase__)[None, :]   ,						(batch_size, vocab_size)).copy() - (
									    vocab_size // 2
									)

									# make ramp_logits more extreme
									SCREAMING_SNAKE_CASE_: Dict  		=       ramp_logits[1] * 100.0

									# make sure at least 2 tokens are kept
									SCREAMING_SNAKE_CASE_: str  		=       FlaxTopPLogitsWarper(0.9   ,						min_tokens_to_keep=2   ,						filter_value=0.0)
									SCREAMING_SNAKE_CASE_: Any  		=       top_p_warp(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)

									# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
									self.assertListEqual((filtered_dist != 0.0).sum(axis=-1).tolist()   ,						[3, 2])



			def        _SCREAMING_SNAKE_CASE			( self  :     Tuple):
									SCREAMING_SNAKE_CASE_: Tuple  		=       20
									SCREAMING_SNAKE_CASE_: List[str]  		=       4
									SCREAMING_SNAKE_CASE_: Optional[int]  		=       0

									SCREAMING_SNAKE_CASE_: str  		=       FlaxMinLengthLogitsProcessor(min_length=10   ,						eos_token_id=lowerCAmelCase__)

									# check that min length is applied at length 5
									SCREAMING_SNAKE_CASE_: str  		=       ids_tensor((batch_size, 20)   ,						vocab_size=20)
									SCREAMING_SNAKE_CASE_: int  		=       5
									SCREAMING_SNAKE_CASE_: List[Any]  		=       self._get_uniform_logits(lowerCAmelCase__   ,						lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: str  		=       min_dist_processor(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
									self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()   ,						4 * [-float("inf")])

									# check that min length is not applied anymore at length 15
									SCREAMING_SNAKE_CASE_: List[str]  		=       self._get_uniform_logits(lowerCAmelCase__   ,						lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       15
									SCREAMING_SNAKE_CASE_: Any  		=       min_dist_processor(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
									self.assertFalse(jnp.isinf(lowerCAmelCase__).any())



			def        _SCREAMING_SNAKE_CASE			( self  :     Union[str, Any]):
									SCREAMING_SNAKE_CASE_: int  		=       20
									SCREAMING_SNAKE_CASE_: str  		=       4
									SCREAMING_SNAKE_CASE_: List[Any]  		=       0

									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase__)

									# check that all scores are -inf except the bos_token_id score
									SCREAMING_SNAKE_CASE_: int  		=       ids_tensor((batch_size, 1)   ,						vocab_size=20)
									SCREAMING_SNAKE_CASE_: List[str]  		=       1
									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       self._get_uniform_logits(lowerCAmelCase__   ,						lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: int  		=       logits_processor(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
									self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]).all())
									self.assertListEqual(scores[:, bos_token_id].tolist()   ,						4 * [0])  # score for bos_token_id shold be zero

									# check that bos_token_id is not forced if current length is greater than 1
									SCREAMING_SNAKE_CASE_: List[Any]  		=       3
									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       self._get_uniform_logits(lowerCAmelCase__   ,						lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: int  		=       logits_processor(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
									self.assertFalse(jnp.isinf(lowerCAmelCase__).any())



			def        _SCREAMING_SNAKE_CASE			( self  :     List[str]):
									SCREAMING_SNAKE_CASE_: Any  		=       20
									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       4
									SCREAMING_SNAKE_CASE_: Dict  		=       0
									SCREAMING_SNAKE_CASE_: List[Any]  		=       5

									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase__   ,						eos_token_id=lowerCAmelCase__)

									# check that all scores are -inf except the eos_token_id when max_length is reached
									SCREAMING_SNAKE_CASE_: List[Any]  		=       ids_tensor((batch_size, 4)   ,						vocab_size=20)
									SCREAMING_SNAKE_CASE_: Optional[int]  		=       4
									SCREAMING_SNAKE_CASE_: Dict  		=       self._get_uniform_logits(lowerCAmelCase__   ,						lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       logits_processor(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
									self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]).all())
									self.assertListEqual(scores[:, eos_token_id].tolist()   ,						4 * [0])  # score for eos_token_id should be zero

									# check that eos_token_id is not forced if max_length is not reached
									SCREAMING_SNAKE_CASE_: List[str]  		=       3
									SCREAMING_SNAKE_CASE_: str  		=       self._get_uniform_logits(lowerCAmelCase__   ,						lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: int  		=       logits_processor(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
									self.assertFalse(jnp.isinf(lowerCAmelCase__).any())



			def        _SCREAMING_SNAKE_CASE			( self  :     Optional[int]):
									SCREAMING_SNAKE_CASE_: int  		=       4
									SCREAMING_SNAKE_CASE_: List[Any]  		=       10
									SCREAMING_SNAKE_CASE_: int  		=       15
									SCREAMING_SNAKE_CASE_: Dict  		=       2
									SCREAMING_SNAKE_CASE_: int  		=       1
									SCREAMING_SNAKE_CASE_: List[Any]  		=       15

									# dummy input_ids and scores
									SCREAMING_SNAKE_CASE_: int  		=       ids_tensor((batch_size, sequence_length)   ,						lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: str  		=       input_ids.copy()

									SCREAMING_SNAKE_CASE_: List[Any]  		=       self._get_uniform_logits(lowerCAmelCase__   ,						lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: Optional[Any]  		=       scores.copy()

									# instantiate all dist processors
									SCREAMING_SNAKE_CASE_: Optional[int]  		=       FlaxTemperatureLogitsWarper(temperature=0.5)
									SCREAMING_SNAKE_CASE_: Tuple  		=       FlaxTopKLogitsWarper(3)
									SCREAMING_SNAKE_CASE_: Optional[int]  		=       FlaxTopPLogitsWarper(0.8)

									# instantiate all logits processors
									SCREAMING_SNAKE_CASE_: Optional[int]  		=       FlaxMinLengthLogitsProcessor(min_length=10   ,						eos_token_id=lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: List[Any]  		=       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: Dict  		=       FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase__   ,						eos_token_id=lowerCAmelCase__)

									SCREAMING_SNAKE_CASE_: str  		=       10

									# no processor list
									SCREAMING_SNAKE_CASE_: Dict  		=       temp_dist_warp(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: Tuple  		=       top_k_warp(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: Any  		=       top_p_warp(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: int  		=       min_dist_proc(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: Optional[int]  		=       bos_dist_proc(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: Tuple  		=       eos_dist_proc(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)

									# with processor list
									SCREAMING_SNAKE_CASE_: str  		=       FlaxLogitsProcessorList(
									    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc])
									SCREAMING_SNAKE_CASE_: Tuple  		=       processor(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)

									# scores should be equal
									self.assertTrue(jnp.allclose(lowerCAmelCase__   ,						lowerCAmelCase__   ,						atol=1E-3))

									# input_ids should never be changed
									self.assertListEqual(input_ids.tolist()   ,						input_ids_comp.tolist())




			def        _SCREAMING_SNAKE_CASE			( self  :     Any):
									SCREAMING_SNAKE_CASE_: Optional[int]  		=       4
									SCREAMING_SNAKE_CASE_: int  		=       10
									SCREAMING_SNAKE_CASE_: List[str]  		=       15
									SCREAMING_SNAKE_CASE_: List[Any]  		=       2
									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       1
									SCREAMING_SNAKE_CASE_: str  		=       15

									# dummy input_ids and scores
									SCREAMING_SNAKE_CASE_: Tuple  		=       ids_tensor((batch_size, sequence_length)   ,						lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: List[Any]  		=       input_ids.copy()

									SCREAMING_SNAKE_CASE_: List[Any]  		=       self._get_uniform_logits(lowerCAmelCase__   ,						lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: List[Any]  		=       scores.copy()

									# instantiate all dist processors
									SCREAMING_SNAKE_CASE_: Dict  		=       FlaxTemperatureLogitsWarper(temperature=0.5)
									SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       FlaxTopKLogitsWarper(3)
									SCREAMING_SNAKE_CASE_: Dict  		=       FlaxTopPLogitsWarper(0.8)

									# instantiate all logits processors
									SCREAMING_SNAKE_CASE_: int  		=       FlaxMinLengthLogitsProcessor(min_length=10   ,						eos_token_id=lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: List[str]  		=       FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: List[Any]  		=       FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase__   ,						eos_token_id=lowerCAmelCase__)

									SCREAMING_SNAKE_CASE_: str  		=       10

									# no processor list
									def run_no_processor_list(lowerCAmelCase__  :     Union[str, Any]   ,						lowerCAmelCase__  :     List[Any]   ,						lowerCAmelCase__  :     Dict):
															SCREAMING_SNAKE_CASE_: Any  		=       temp_dist_warp(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
															SCREAMING_SNAKE_CASE_: Optional[int]  		=       top_k_warp(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
															SCREAMING_SNAKE_CASE_: Optional[int]  		=       top_p_warp(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
															SCREAMING_SNAKE_CASE_: List[Any]  		=       min_dist_proc(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
															SCREAMING_SNAKE_CASE_: Union[str, Any]  		=       bos_dist_proc(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
															SCREAMING_SNAKE_CASE_: Optional[int]  		=       eos_dist_proc(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
															return scores

									# with processor list
									def run_processor_list(lowerCAmelCase__  :     int   ,						lowerCAmelCase__  :     List[Any]   ,						lowerCAmelCase__  :     Union[str, Any]):
															SCREAMING_SNAKE_CASE_: List[str]  		=       FlaxLogitsProcessorList(
															    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc])
															SCREAMING_SNAKE_CASE_: Dict  		=       processor(lowerCAmelCase__   ,						lowerCAmelCase__   ,						cur_len=lowerCAmelCase__)
															return scores

									SCREAMING_SNAKE_CASE_: str  		=       jax.jit(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: Optional[int]  		=       jax.jit(lowerCAmelCase__)

									SCREAMING_SNAKE_CASE_: Dict  		=       jitted_run_no_processor_list(lowerCAmelCase__   ,						lowerCAmelCase__   ,						lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_: List[Any]  		=       jitted_run_processor_list(lowerCAmelCase__   ,						lowerCAmelCase__   ,						lowerCAmelCase__)

									# scores should be equal
									self.assertTrue(jnp.allclose(lowerCAmelCase__   ,						lowerCAmelCase__   ,						atol=1E-3))

									# input_ids should never be changed
									self.assertListEqual(input_ids.tolist()   ,						input_ids_comp.tolist())






