

import unittest

from transformers import TrOCRConfig
from transformers.testing_utils import is_torch_available, require_torch, 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.models.trocr.modeling_trocr import TrOCRDecoder, TrOCRForCausalLM




@require_torch
class     a			:






						"""simple docstring"""




						def __init__(							self					:      List[Any]			,			snake_case_					:      List[Any]			,			snake_case_					:      Dict=9_9			,			snake_case_					:      Dict=1_3			,			snake_case_					:      Union[str, Any]=1_6			,			snake_case_					:      List[Any]=7			,			snake_case_					:      Union[str, Any]=True			,			snake_case_					:      Tuple=True			,			snake_case_					:      Any=True			,			snake_case_					:      Optional[int]=False			,			snake_case_					:      int=True			,			snake_case_					:      List[Any]=2			,			snake_case_					:      List[Any]=3_2			,			snake_case_					:      Optional[Any]=4			,			snake_case_					:      List[str]=4			,			snake_case_					:      int=3_0			,			snake_case_					:      Tuple=0			,			snake_case_					:      Optional[Any]=1			,			snake_case_					:      int=2			,			snake_case_					:      Tuple=None			,			):



								'''simple docstring'''



								snake_case__						:       Any       						=  parent
								snake_case__						:       str       						=  batch_size
								snake_case__						:       Any       						=  decoder_seq_length
								# For common tests
								snake_case__						:       Optional[int]       						=  self.decoder_seq_length
								snake_case__						:       List[Any]       						=  is_training
								snake_case__						:       Tuple       						=  use_attention_mask
								snake_case__						:       List[Any]       						=  use_labels

								snake_case__						:       Optional[int]       						=  vocab_size
								snake_case__						:       Optional[int]       						=  d_model
								snake_case__						:       Tuple       						=  d_model
								snake_case__						:       Any       						=  decoder_layers
								snake_case__						:       int       						=  decoder_layers
								snake_case__						:       Tuple       						=  decoder_ffn_dim
								snake_case__						:       Any       						=  decoder_attention_heads
								snake_case__						:       Dict       						=  decoder_attention_heads
								snake_case__						:       int       						=  eos_token_id
								snake_case__						:       Tuple       						=  bos_token_id
								snake_case__						:       int       						=  pad_token_id
								snake_case__						:       List[Any]       						=  decoder_start_token_id
								snake_case__						:       Dict       						=  use_cache
								snake_case__						:       Dict       						=  max_position_embeddings

								snake_case__						:       Tuple       						=  None
								snake_case__						:       int       						=  decoder_seq_length
								snake_case__						:       Tuple       						=  2
								snake_case__						:       Tuple       						=  1







						def       __magic_name__							(							self					:      Tuple  ):



								'''simple docstring'''



								snake_case__						:       Optional[int]       						=  ids_tensor([self.batch_size, self.decoder_seq_length]			,			self.vocab_size  )

								snake_case__						:       Optional[int]       						=  None
								if self.use_attention_mask:
										snake_case__						:       Optional[int]       						=  ids_tensor([self.batch_size, self.decoder_seq_length]			,			vocab_size=2  )

								snake_case__						:       List[Any]       						=  None
								if self.use_labels:
										snake_case__						:       Any       						=  ids_tensor([self.batch_size, self.decoder_seq_length]			,			self.vocab_size  )

								snake_case__						:       int       						=  TrOCRConfig(
								    vocab_size=self.vocab_size			,			d_model=self.d_model			,			decoder_layers=self.decoder_layers			,			decoder_ffn_dim=self.decoder_ffn_dim			,			decoder_attention_heads=self.decoder_attention_heads			,			eos_token_id=self.eos_token_id			,			bos_token_id=self.bos_token_id			,			use_cache=self.use_cache			,			pad_token_id=self.pad_token_id			,			decoder_start_token_id=self.decoder_start_token_id			,			max_position_embeddings=self.max_position_embeddings			,			)

								return (config, input_ids, attention_mask, lm_labels)







						def       __magic_name__							(							self					:      Dict			,			snake_case_					:      List[Any]			,			snake_case_					:      int			,			snake_case_					:      Union[str, Any]			,			snake_case_					:      str			,			):



								'''simple docstring'''



								snake_case__						:       Optional[int]       						=  True
								snake_case__						:       Union[str, Any]       						=  TrOCRDecoder(config=snake_case_  ).to(snake_case_  ).eval()
								snake_case__						:       str       						=  input_ids[:2]

								input_ids[input_ids == 0] += 1
								# first forward pass
								snake_case__						:       Optional[Any]       						=  model(snake_case_			,			use_cache=snake_case_  )
								snake_case__						:       Any       						=  model(snake_case_  )
								snake_case__						:       List[str]       						=  model(snake_case_			,			use_cache=snake_case_  )

								self.parent.assertTrue(len(snake_case_  ) == len(snake_case_  )  )
								self.parent.assertTrue(len(snake_case_  ) == len(snake_case_  ) + 1  )

								snake_case__						:       Optional[int]       						=  outputs['''past_key_values''']

								# create hypothetical next token and extent to next_input_ids
								snake_case__						:       Union[str, Any]       						=  ids_tensor((2, 1)			,			config.vocab_size - 1  ) + 1

								# append to next input_ids and
								snake_case__						:       str       						=  torch.cat([input_ids, next_tokens]			,			dim=-1  )

								snake_case__						:       str       						=  model(snake_case_  )['''last_hidden_state''']
								snake_case__						:       str       						=  model(snake_case_			,			past_key_values=snake_case_  )['''last_hidden_state''']

								# select random slice
								snake_case__						:       Dict       						=  ids_tensor((1,)			,			output_from_past.shape[-1]  ).item()
								snake_case__						:       int       						=  output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
								snake_case__						:       Optional[Any]       						=  output_from_past[:, 0, random_slice_idx].detach()

								# test that outputs are equal for slice
								assert torch.allclose(snake_case_			,			snake_case_			,			atol=1e-3  )







						def       __magic_name__							(							self					:      int  ):



								'''simple docstring'''



								snake_case__						:       str       						=  self.prepare_config_and_inputs()
								snake_case__						:       Tuple       						=  config_and_inputs

								snake_case__						:       Any       						=  {'''input_ids''': input_ids, '''attention_mask''': attention_mask}
								return config, inputs_dict





@require_torch
class     a			(				SCREAMING_SNAKE_CASE			,			SCREAMING_SNAKE_CASE			,			SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    (TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
						__UpperCAmelCase			     =    (TrOCRForCausalLM,) if is_torch_available() else ()
						__UpperCAmelCase			     =    {"""text-generation""": TrOCRForCausalLM} if is_torch_available() else {}
						__UpperCAmelCase			     =    True
						__UpperCAmelCase			     =    False







						def       __magic_name__							(							self					:      Any  ):



								'''simple docstring'''



								snake_case__						:       Union[str, Any]       						=  TrOCRStandaloneDecoderModelTester(self			,			is_training=snake_case_  )
								snake_case__						:       Optional[int]       						=  ConfigTester(self			,			config_class=snake_case_  )







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



								'''simple docstring'''



								pass







						def       __magic_name__							(							self					:      Optional[int]  ):



								'''simple docstring'''



								pass







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



								'''simple docstring'''



								pass







						def       __magic_name__							(							self					:      Optional[int]  ):



								'''simple docstring'''



								self.config_tester.run_common_tests()







						def       __magic_name__							(							self					:      List[Any]  ):



								'''simple docstring'''



								snake_case__						:       List[str]       						=  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_decoder_model_past(*snake_case_  )







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



								'''simple docstring'''



								return







						@unittest.skip('''The model doesn\'t support left padding'''  )  # and it's not used enough to be worth fixing :)
						def       __magic_name__							(							self					:      int  ):



								'''simple docstring'''



								pass






'''simple docstring'''


import unittest

from diffusers import FlaxAutoencoderKL
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax

from .test_modeling_common_flax import FlaxModelTesterMixin


if is_flax_available():
					import jax




@require_flax
class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    FlaxAutoencoderKL







						@property
						def       __magic_name__							(							self					:      str  ):



								'''simple docstring'''



								snake_case__						:       Any       						=  4
								snake_case__						:       Optional[Any]       						=  3
								snake_case__						:       Optional[int]       						=  (3_2, 3_2)

								snake_case__						:       Optional[int]       						=  jax.random.PRNGKey(0  )
								snake_case__						:       Union[str, Any]       						=  jax.random.uniform(snake_case_			,			((batch_size, num_channels) + sizes)  )

								return {"sample": image, "prng_key": prng_key}







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



								'''simple docstring'''



								snake_case__						:       Any       						=  {
								    '''block_out_channels''': [3_2, 6_4],
								    '''in_channels''': 3,
								    '''out_channels''': 3,
								    '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
								    '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
								    '''latent_channels''': 4,
								}
								snake_case__						:       Dict       						=  self.dummy_input
								return init_dict, inputs_dict



