


import unittest

import numpy as np

from transformers import AlbertConfig, 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.numpy as jnp

   from transformers.models.albert.modeling_flax_albert import (
       FlaxAlbertForMaskedLM,
       FlaxAlbertForMultipleChoice,
       FlaxAlbertForPreTraining,
       FlaxAlbertForQuestionAnswering,
       FlaxAlbertForSequenceClassification,
       FlaxAlbertForTokenClassification,
       FlaxAlbertModel,
   )






class     UpperCAmelCase__       (					unittest.TestCase		):



    """simple docstring"""






    def __init__(			self					,   A_					,   A_=13					,   A_=7					,   A_=True					,   A_=True					,   A_=True					,   A_=True					,   A_=99					,   A_=32					,   A_=5					,   A_=4					,   A_=37					,   A_="gelu"					,   A_=0.1					,   A_=0.1					,   A_=512					,   A_=16					,   A_=2					,   A_=0.02					,   A_=4					,   )  ->       Union[str, Any]:
         __UpperCamelCase    							=parent
         __UpperCamelCase    							=batch_size
         __UpperCamelCase    							=seq_length
         __UpperCamelCase    							=is_training
         __UpperCamelCase    							=use_attention_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_choices






    def _a					(			self    )  ->       Dict:
         __UpperCamelCase    							=ids_tensor([self.batch_size, self.seq_length]					,   self.vocab_size    )

         __UpperCamelCase    							=None
         if self.use_attention_mask:
              __UpperCamelCase    							=random_attention_mask([self.batch_size, self.seq_length]    )

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

         __UpperCamelCase    							=AlbertConfig(
             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=A_					,   initializer_range=self.initializer_range					,   )

         return config, input_ids, token_type_ids, attention_mask






    def _a					(			self    )  ->       Optional[Any]:
         __UpperCamelCase    							=self.prepare_config_and_inputs()
         __UpperCamelCase							,      __UpperCamelCase							,      __UpperCamelCase							,      __UpperCamelCase    							=config_and_inputs
         __UpperCamelCase    							={'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': attention_mask}
         return config, inputs_dict






@require_flax
class     UpperCAmelCase__       (					A_   ,  unittest.TestCase		):



    """simple docstring"""






    UpperCAmelCase__						:		List[str]      =  (
        (
            FlaxAlbertModel,
            FlaxAlbertForPreTraining,
            FlaxAlbertForMaskedLM,
            FlaxAlbertForMultipleChoice,
            FlaxAlbertForQuestionAnswering,
            FlaxAlbertForSequenceClassification,
            FlaxAlbertForTokenClassification,
            FlaxAlbertForQuestionAnswering,
        )
        if is_flax_available()
        else ()
    )






    def _a					(			self    )  ->       Tuple:
         __UpperCamelCase    							=FlaxAlbertModelTester(self    )






    @slow
    def _a					(			self    )  ->       str:
         for model_class_name in self.all_model_classes:
              __UpperCamelCase    							=model_class_name.from_pretrained('albert-base-v2'    )
              __UpperCamelCase    							=model(np.ones((1, 1)    )    )
              self.assertIsNotNone(A_    )






@require_flax
class     UpperCAmelCase__       (					unittest.TestCase		):



    """simple docstring"""






    @slow
    def _a					(			self    )  ->       str:
         __UpperCamelCase    							=FlaxAlbertModel.from_pretrained('albert-base-v2'    )
         __UpperCamelCase    							=np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]    )
         __UpperCamelCase    							=np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]    )
         __UpperCamelCase    							=model(A_					,   attention_mask=A_    )[0]
         __UpperCamelCase    							=(1, 11, 768)
         self.assertEqual(output.shape					,   A_    )
         __UpperCamelCase    							=np.array(
             [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]]    )

         self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]					,   A_					,   atol=1E-4    )    )









'''simple docstring'''



import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()





class    UpperCAmelCase_  (_UpperCAmelCase					,		_UpperCAmelCase					,		unittest.TestCase  ):





    """simple docstring"""






    lowerCamelCase  :					Union[str, Any]							     =			VQModel
    lowerCamelCase  :					Union[str, Any]							     =			'sample'







    @property
    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=(32, 32)	)				->     Any:
         __lowerCamelCase   :	Tuple								=		4
         __lowerCamelCase   :	Union[str, Any]								=		3

         __lowerCamelCase   :	List[Any]								=		floats_tensor((batch_size, num_channels) + sizes	).to(SCREAMING_SNAKE_CASE_	)

         return {"sample": image}







    @property
    def       lowercase_				(							self	)				->     Optional[int]:
         return (3, 32, 32)







    @property
    def       lowercase_				(							self	)				->     List[Any]:
         return (3, 32, 32)







    def       lowercase_				(							self	)				->     Optional[int]:
         __lowerCamelCase   :	Optional[Any]								=		{
             'block_out_channels': [32, 64],
             'in_channels': 3,
             'out_channels': 3,
             'down_block_types': ['DownEncoderBlock2D', 'DownEncoderBlock2D'],
             'up_block_types': ['UpDecoderBlock2D', 'UpDecoderBlock2D'],
             'latent_channels': 3,
         }
         __lowerCamelCase   :	Union[str, Any]								=		self.dummy_input
         return init_dict, inputs_dict







    def       lowercase_				(							self	)				->     str:
         pass







    def       lowercase_				(							self	)				->     Optional[int]:
         pass







    def       lowercase_				(							self	)				->     Any:
         __lowerCamelCase				,       __lowerCamelCase   :	Union[str, Any]								=		VQModel.from_pretrained('fusing/vqgan-dummy'			,     output_loading_info=SCREAMING_SNAKE_CASE_	)
         self.assertIsNotNone(SCREAMING_SNAKE_CASE_	)
         self.assertEqual(len(loading_info['missing_keys']	)			,     0	)

         model.to(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	List[Any]								=		model(**self.dummy_input	)

         assert image is not None, "Make sure output is not None"







    def       lowercase_				(							self	)				->     int:
         __lowerCamelCase   :	List[Any]								=		VQModel.from_pretrained('fusing/vqgan-dummy'	)
         model.to(SCREAMING_SNAKE_CASE_	).eval()

         torch.manual_seed(0	)
         if torch.cuda.is_available():
              torch.cuda.manual_seed_all(0	)

         __lowerCamelCase   :	List[str]								=		torch.randn(1			,     model.config.in_channels			,     model.config.sample_size			,     model.config.sample_size	)
         __lowerCamelCase   :	Optional[int]								=		image.to(SCREAMING_SNAKE_CASE_	)
         with torch.no_grad():
              __lowerCamelCase   :	int								=		model(SCREAMING_SNAKE_CASE_	).sample

         __lowerCamelCase   :	List[str]								=		output[0, -1, -3:, -3:].flatten().cpu()
         # fmt: off
         __lowerCamelCase   :	Union[str, Any]								=		torch.tensor([-0.0_1_5_3, -0.4_0_4_4, -0.1_8_8_0, -0.5_1_6_1, -0.2_4_1_8, -0.4_0_7_2, -0.1_6_1_2, -0.0_6_3_3, -0.0_1_4_3]	)
         # fmt: on
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     atol=1E-3	)	)






