


'''simple docstring'''



import unittest

from transformers import LiltConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import (
					    LiltForQuestionAnswering,
					    LiltForSequenceClassification,
					    LiltForTokenClassification,
					    LiltModel,
					)
					from transformers.models.lilt.modeling_lilt import LILT_PRETRAINED_MODEL_ARCHIVE_LIST






class 	UpperCamelCase__:




				def __init__(		self      :			Any				,   lowerCAmelCase      :			Union[str, Any]				,   lowerCAmelCase      :			Tuple=13				,   lowerCAmelCase      :			str=7				,   lowerCAmelCase      :			List[str]=True				,   lowerCAmelCase      :			List[Any]=True				,   lowerCAmelCase      :			Optional[Any]=True				,   lowerCAmelCase      :			List[Any]=True				,   lowerCAmelCase      :			List[str]=99				,   lowerCAmelCase      :			Optional[Any]=24				,   lowerCAmelCase      :			Optional[Any]=2				,   lowerCAmelCase      :			int=6				,   lowerCAmelCase      :			Any=37				,   lowerCAmelCase      :			Any="gelu"				,   lowerCAmelCase      :			List[Any]=0.1				,   lowerCAmelCase      :			Optional[int]=0.1				,   lowerCAmelCase      :			Dict=512				,   lowerCAmelCase      :			List[Any]=16				,   lowerCAmelCase      :			List[str]=2				,   lowerCAmelCase      :			List[Any]=0.02				,   lowerCAmelCase      :			List[str]=3				,   lowerCAmelCase      :			Any=None				,   lowerCAmelCase      :			Dict=1000				,   )->					Any:

										"""simple docstring"""




										UpperCAmelCase					    =     parent
										UpperCAmelCase					    =     batch_size
										UpperCAmelCase					    =     seq_length
										UpperCAmelCase					    =     is_training
										UpperCAmelCase					    =     use_input_mask
										UpperCAmelCase					    =     use_token_type_ids
										UpperCAmelCase					    =     use_labels
										UpperCAmelCase					    =     vocab_size
										UpperCAmelCase					    =     hidden_size
										UpperCAmelCase					    =     num_hidden_layers
										UpperCAmelCase					    =     num_attention_heads
										UpperCAmelCase					    =     intermediate_size
										UpperCAmelCase					    =     hidden_act
										UpperCAmelCase					    =     hidden_dropout_prob
										UpperCAmelCase					    =     attention_probs_dropout_prob
										UpperCAmelCase					    =     max_position_embeddings
										UpperCAmelCase					    =     type_vocab_size
										UpperCAmelCase					    =     type_sequence_label_size
										UpperCAmelCase					    =     initializer_range
										UpperCAmelCase					    =     num_labels
										UpperCAmelCase					    =     scope
										UpperCAmelCase					    =     range_bbox




				def 					a__(		self      :			int				)->					Union[str, Any]:

										"""simple docstring"""




										UpperCAmelCase					    =     ids_tensor([self.batch_size, self.seq_length]				,   self.vocab_size				)

										UpperCAmelCase					    =     ids_tensor([self.batch_size, self.seq_length, 4]				,   self.range_bbox				)
										# Ensure that bbox is legal
										for i in range(bbox.shape[0]				):
																for j in range(bbox.shape[1]				):
																						if bbox[i, j, 3] < bbox[i, j, 1]:
																												UpperCAmelCase					    =     bbox[i, j, 3]
																												UpperCAmelCase					    =     bbox[i, j, 1]
																												UpperCAmelCase					    =     t
																						if bbox[i, j, 2] < bbox[i, j, 0]:
																												UpperCAmelCase					    =     bbox[i, j, 2]
																												UpperCAmelCase					    =     bbox[i, j, 0]
																												UpperCAmelCase					    =     t

										UpperCAmelCase					    =     None
										if self.use_input_mask:
																UpperCAmelCase					    =     ids_tensor([self.batch_size, self.seq_length]				,   vocab_size=2				)

										UpperCAmelCase					    =     None
										if self.use_token_type_ids:
																UpperCAmelCase					    =     ids_tensor([self.batch_size, self.seq_length]				,   self.type_vocab_size				)

										UpperCAmelCase					    =     None
										UpperCAmelCase					    =     None
										if self.use_labels:
																UpperCAmelCase					    =     ids_tensor([self.batch_size]				,   self.type_sequence_label_size				)
																UpperCAmelCase					    =     ids_tensor([self.batch_size, self.seq_length]				,   self.num_labels				)

										UpperCAmelCase					    =     self.get_config()

										return config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels




				def 					a__(		self      :			Tuple				)->					str:

										"""simple docstring"""




										return LiltConfig(
										    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				,   initializer_range=self.initializer_range				,   )




				def 					a__(		self      :			Any				,   lowerCAmelCase      :			Optional[int]				,   lowerCAmelCase      :			List[str]				,   lowerCAmelCase      :			Optional[Any]				,   lowerCAmelCase      :			Optional[int]				,   lowerCAmelCase      :			List[str]				,   lowerCAmelCase      :			int				,   lowerCAmelCase      :			Union[str, Any]				,   )->					Any:

										"""simple docstring"""




										UpperCAmelCase					    =     LiltModel(config=lowerCAmelCase				)
										model.to(lowerCAmelCase				)
										model.eval()
										UpperCAmelCase					    =     model(lowerCAmelCase				,   bbox=lowerCAmelCase				,   attention_mask=lowerCAmelCase				,   token_type_ids=lowerCAmelCase				)
										UpperCAmelCase					    =     model(lowerCAmelCase				,   bbox=lowerCAmelCase				,   token_type_ids=lowerCAmelCase				)
										UpperCAmelCase					    =     model(lowerCAmelCase				,   bbox=lowerCAmelCase				)

										self.parent.assertEqual(result.last_hidden_state.shape				,   (self.batch_size, self.seq_length, self.hidden_size)				)
										self.parent.assertEqual(result.pooler_output.shape				,   (self.batch_size, self.hidden_size)				)




				def 					a__(		self      :			List[Any]				,   lowerCAmelCase      :			Union[str, Any]				,   lowerCAmelCase      :			Tuple				,   lowerCAmelCase      :			Tuple				,   lowerCAmelCase      :			List[Any]				,   lowerCAmelCase      :			Optional[Any]				,   lowerCAmelCase      :			Any				,   lowerCAmelCase      :			Optional[Any]				,   )->					Any:

										"""simple docstring"""




										UpperCAmelCase					    =     self.num_labels
										UpperCAmelCase					    =     LiltForTokenClassification(config=lowerCAmelCase				)
										model.to(lowerCAmelCase				)
										model.eval()
										UpperCAmelCase					    =     model(
										    lowerCAmelCase				,   bbox=lowerCAmelCase				,   attention_mask=lowerCAmelCase				,   token_type_ids=lowerCAmelCase				,   labels=lowerCAmelCase				)
										self.parent.assertEqual(result.logits.shape				,   (self.batch_size, self.seq_length, self.num_labels)				)




				def 					a__(		self      :			Union[str, Any]				,   lowerCAmelCase      :			Any				,   lowerCAmelCase      :			Optional[int]				,   lowerCAmelCase      :			List[Any]				,   lowerCAmelCase      :			List[Any]				,   lowerCAmelCase      :			Optional[Any]				,   lowerCAmelCase      :			int				,   lowerCAmelCase      :			Dict				,   )->					Tuple:

										"""simple docstring"""




										UpperCAmelCase					    =     LiltForQuestionAnswering(config=lowerCAmelCase				)
										model.to(lowerCAmelCase				)
										model.eval()
										UpperCAmelCase					    =     model(
										    lowerCAmelCase				,   bbox=lowerCAmelCase				,   attention_mask=lowerCAmelCase				,   token_type_ids=lowerCAmelCase				,   start_positions=lowerCAmelCase				,   end_positions=lowerCAmelCase				,   )
										self.parent.assertEqual(result.start_logits.shape				,   (self.batch_size, self.seq_length)				)
										self.parent.assertEqual(result.end_logits.shape				,   (self.batch_size, self.seq_length)				)





				def 					a__(		self      :			Union[str, Any]				)->					Tuple:

										"""simple docstring"""




										UpperCAmelCase					    =     self.prepare_config_and_inputs()
										(
										    (
										    UpperCAmelCase    
										)    ,    (
										    UpperCAmelCase    
										)    ,    (
										    UpperCAmelCase    
										)    ,    (
										    UpperCAmelCase    
										)    ,    (
										    UpperCAmelCase    
										)    ,    (
										    UpperCAmelCase    
										)    ,    (
										    UpperCAmelCase    
										)    ,        
										)					    =     config_and_inputs
										UpperCAmelCase					    =     {
										    '''input_ids''': input_ids,
										    '''bbox''': bbox,
										    '''token_type_ids''': token_type_ids,
										    '''attention_mask''': input_mask,
										}
										return config, inputs_dict







@require_torch
class 	UpperCamelCase__(       lowerCAmelCase						,    lowerCAmelCase						,    lowerCAmelCase						,    unittest.TestCase	):
				__magic_name__    :				Dict  		=	(
				    (
				        LiltModel,
				        LiltForSequenceClassification,
				        LiltForTokenClassification,
				        LiltForQuestionAnswering,
				    )
				    if is_torch_available()
				    else ()
				)
				__magic_name__    :				Optional[Any]  		=	(
				    {
				        "feature-extraction": LiltModel,
				        "question-answering": LiltForQuestionAnswering,
				        "text-classification": LiltForSequenceClassification,
				        "token-classification": LiltForTokenClassification,
				        "zero-shot": LiltForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				__magic_name__    :				List[Any]  		=	False
				__magic_name__    :				Optional[int]  		=	False




				def 					a__(		self      :			Dict				,   lowerCAmelCase      :			List[str]				,   lowerCAmelCase      :			Dict				,   lowerCAmelCase      :			Union[str, Any]				,   lowerCAmelCase      :			Union[str, Any]				,   lowerCAmelCase      :			int				)->					Tuple:

										"""simple docstring"""




										return True




				def 					a__(		self      :			List[Any]				)->					Optional[int]:

										"""simple docstring"""




										UpperCAmelCase					    =     LiltModelTester(self				)
										UpperCAmelCase					    =     ConfigTester(self				,   config_class=lowerCAmelCase				,   hidden_size=37				)




				def 					a__(		self      :			Optional[Any]				)->					Union[str, Any]:

										"""simple docstring"""




										self.config_tester.run_common_tests()




				def 					a__(		self      :			Any				)->					Any:

										"""simple docstring"""




										UpperCAmelCase					    =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*lowerCAmelCase				)




				def 					a__(		self      :			Union[str, Any]				)->					Any:

										"""simple docstring"""




										UpperCAmelCase					    =     self.model_tester.prepare_config_and_inputs()
										for type in ["absolute", "relative_key", "relative_key_query"]:
																UpperCAmelCase					    =     type
																self.model_tester.create_and_check_model(*lowerCAmelCase				)




				def 					a__(		self      :			str				)->					List[str]:

										"""simple docstring"""




										UpperCAmelCase					    =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase				)




				def 					a__(		self      :			List[Any]				)->					Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase					    =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_question_answering(*lowerCAmelCase				)





				@slow
				def 					a__(		self      :			Optional[Any]				)->					List[Any]:

										"""simple docstring"""




										for model_name in LILT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																UpperCAmelCase					    =     LiltModel.from_pretrained(lowerCAmelCase				)
																self.assertIsNotNone(lowerCAmelCase				)







@require_torch
@slow
class 	UpperCamelCase__(       unittest.TestCase	):




				def 					a__(		self      :			int				)->					str:

										"""simple docstring"""




										UpperCAmelCase					    =     LiltModel.from_pretrained('''SCUT-DLVCLab/lilt-roberta-en-base'''				).to(lowerCAmelCase				)

										UpperCAmelCase					    =     torch.tensor([[1, 2]]				,   device=lowerCAmelCase				)
										UpperCAmelCase					    =     torch.tensor([[[1, 2, 3, 4], [5, 6, 7, 8]]]				,   device=lowerCAmelCase				)

										# forward pass
										with torch.no_grad():
																UpperCAmelCase					    =     model(input_ids=lowerCAmelCase				,   bbox=lowerCAmelCase				)

										UpperCAmelCase					    =     torch.Size([1, 2, 768]				)
										UpperCAmelCase					    =     torch.tensor(
										    [[-0.0653, 0.0950, -0.0061], [-0.0545, 0.0926, -0.0324]]				,   device=lowerCAmelCase				,   )

										self.assertTrue(outputs.last_hidden_state.shape				,   lowerCAmelCase				)
										self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :, :3]				,   lowerCAmelCase				,   atol=1E-3				)				)









'''simple docstring'''



import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
	import os

	import jax.numpy as jnp
	from jax import jit

	from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
	from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

	_lowercase      :      List[str]            =					"""0.12"""  # assumed parallelism: 8


if is_torch_available():
	import torch


def 				lowerCamelCase__				(	A       :      str				,       A       :      str				,       A       :      List[Any]=None					):




						'''simple docstring'''






						if rng is None:
												UpperCAmelCase					    =     random.Random()

						UpperCAmelCase					    =     1
						for dim in shape:
												total_dims *= dim

						UpperCAmelCase					    =     []
						for _ in range(A					):
												values.append(rng.randint(0				,       vocab_size - 1					)					)

						UpperCAmelCase					    =     np.array(A				,       dtype=jnp.intaa					).reshape(A					)

						return output


def 				lowerCamelCase__				(	A       :      int				,       A       :      Optional[int]=None					):




						'''simple docstring'''






						UpperCAmelCase					    =     ids_tensor(A				,       vocab_size=2				,       rng=A					)
						# make sure that at least one token is attended to for each batch
						UpperCAmelCase					    =     1
						return attn_mask






@require_flax
class 	UpperCamelCase__:
				__magic_name__    :				Optional[int]  		=	None
				__magic_name__    :				Optional[Any]  		=	()




				def 					a__(		self      :			str				)->					Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase					    =     self.model_tester.prepare_config_and_inputs_for_common()

										# cut to half length & take max batch_size 3
										UpperCAmelCase					    =     2
										UpperCAmelCase					    =     inputs['''input_ids'''].shape[-1] // 2
										UpperCAmelCase					    =     inputs['''input_ids'''][:max_batch_size, :sequence_length]

										UpperCAmelCase					    =     jnp.ones_like(lowerCAmelCase				)
										UpperCAmelCase					    =     attention_mask[:max_batch_size, :sequence_length]

										# generate max 5 tokens
										UpperCAmelCase					    =     input_ids.shape[-1] + 5
										if config.eos_token_id is not None and config.pad_token_id is None:
																# hack to allow generate for models such as GPT2 as is done in `generate()`
																UpperCAmelCase					    =     config.eos_token_id
										return config, input_ids, attention_mask, max_length




				@is_pt_flax_cross_test
				def 					a__(		self      :			Dict				)->					Optional[int]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()
										UpperCAmelCase					    =     False
										UpperCAmelCase					    =     max_length
										UpperCAmelCase					    =     0

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model_class.__name__[4:]  # Skip the "Flax" at the beginning
																UpperCAmelCase					    =     getattr(lowerCAmelCase				,   lowerCAmelCase				)
																UpperCAmelCase					    =     pt_model_class(lowerCAmelCase				).eval()
																UpperCAmelCase					    =     load_flax_weights_in_pytorch_model(lowerCAmelCase				,   flax_model.params				)

																UpperCAmelCase					    =     flax_model.generate(lowerCAmelCase				).sequences
																UpperCAmelCase					    =     pt_model.generate(torch.tensor(lowerCAmelCase				,   dtype=torch.long				)				)

																if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
																						UpperCAmelCase					    =     flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

																self.assertListEqual(pt_generation_outputs.numpy().tolist()				,   flax_generation_outputs.tolist()				)




				def 					a__(		self      :			Any				)->					Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()
										UpperCAmelCase					    =     False
										UpperCAmelCase					    =     max_length

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)




				def 					a__(		self      :			Optional[Any]				)->					int:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()
										UpperCAmelCase					    =     True
										UpperCAmelCase					    =     max_length

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)




				def 					a__(		self      :			str				)->					List[str]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()
										UpperCAmelCase					    =     False
										UpperCAmelCase					    =     max_length
										UpperCAmelCase					    =     2

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)




				def 					a__(		self      :			List[Any]				)->					List[str]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()
										UpperCAmelCase					    =     False
										UpperCAmelCase					    =     max_length
										UpperCAmelCase					    =     2
										UpperCAmelCase					    =     2

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[0]				,   input_ids.shape[0] * config.num_return_sequences				)




				def 					a__(		self      :			Tuple				)->					List[str]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()
										UpperCAmelCase					    =     True
										UpperCAmelCase					    =     max_length
										UpperCAmelCase					    =     0.8
										UpperCAmelCase					    =     10
										UpperCAmelCase					    =     0.3
										UpperCAmelCase					    =     1
										UpperCAmelCase					    =     8
										UpperCAmelCase					    =     9

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)




				def 					a__(		self      :			Optional[Any]				)->					Optional[int]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()
										UpperCAmelCase					    =     max_length
										UpperCAmelCase					    =     1
										UpperCAmelCase					    =     8
										UpperCAmelCase					    =     9

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)




				def 					a__(		self      :			Tuple				)->					Tuple:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()
										UpperCAmelCase					    =     max_length
										UpperCAmelCase					    =     2
										UpperCAmelCase					    =     1
										UpperCAmelCase					    =     8
										UpperCAmelCase					    =     9

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)




				def 					a__(		self      :			Union[str, Any]				)->					Any:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()

										# pad attention mask on the left
										UpperCAmelCase					    =     attention_mask.at[(0, 0)].set(0				)

										UpperCAmelCase					    =     False
										UpperCAmelCase					    =     max_length

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				,   attention_mask=lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				,   attention_mask=lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)




				def 					a__(		self      :			Optional[Any]				)->					int:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()

										# pad attention mask on the left
										UpperCAmelCase					    =     attention_mask.at[(0, 0)].set(0				)

										UpperCAmelCase					    =     True
										UpperCAmelCase					    =     max_length

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				,   attention_mask=lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				,   attention_mask=lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)





				def 					a__(		self      :			Tuple				)->					Union[str, Any]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase    ,    UpperCAmelCase					    =     self._get_input_ids_and_config()

										# pad attention mask on the left
										UpperCAmelCase					    =     attention_mask.at[(0, 0)].set(0				)

										UpperCAmelCase					    =     2
										UpperCAmelCase					    =     max_length

										for model_class in self.all_generative_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)

																UpperCAmelCase					    =     model.generate(lowerCAmelCase				,   attention_mask=lowerCAmelCase				).sequences
																self.assertEqual(generation_outputs.shape[-1]				,   lowerCAmelCase				)

																UpperCAmelCase					    =     jit(model.generate				)
																UpperCAmelCase					    =     jit_generate(lowerCAmelCase				,   attention_mask=lowerCAmelCase				).sequences

																self.assertListEqual(generation_outputs.tolist()				,   jit_generation_outputs.tolist()				)







@require_flax
class 	UpperCamelCase__(       unittest.TestCase	):




				def 					a__(		self      :			Union[str, Any]				)->					Optional[int]:

										"""simple docstring"""




										UpperCAmelCase					    =     AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-bert'''				)
										UpperCAmelCase					    =     FlaxAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''				)

										UpperCAmelCase					    =     '''Hello world'''
										UpperCAmelCase					    =     tokenizer(lowerCAmelCase				,   return_tensors='''np'''				).input_ids

										# typos are quickly detected (the correct argument is `do_sample`)
										with self.assertRaisesRegex(lowerCAmelCase				,   '''do_samples'''				):
																model.generate(lowerCAmelCase				,   do_samples=lowerCAmelCase				)

										# arbitrary arguments that will not be used anywhere are also not accepted
										with self.assertRaisesRegex(lowerCAmelCase				,   '''foo'''				):
																UpperCAmelCase					    =     {'''foo''': '''bar'''}
																model.generate(lowerCAmelCase				,   **lowerCAmelCase				)





