



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 			_lowerCAmelCase   ( unittest.TestCase			):






			"""simple docstring"""





			def __init__(				self	:       Dict			,     lowercase	:       int			,     lowercase	:       int=13			,     lowercase	:       Any=7			,     lowercase	:       Optional[Any]=True			,     lowercase	:       List[str]=True			,     lowercase	:       int=True			,     lowercase	:       Optional[Any]=True			,     lowercase	:       Union[str, Any]=99			,     lowercase	:       Union[str, Any]=32			,     lowercase	:       List[Any]=5			,     lowercase	:       List[str]=4			,     lowercase	:       Union[str, Any]=37			,     lowercase	:       int="gelu"			,     lowercase	:       List[str]=0.1			,     lowercase	:       Any=0.1			,     lowercase	:       List[str]=512			,     lowercase	:       List[str]=16			,     lowercase	:       Dict=2			,     lowercase	:       Optional[Any]=0.02			,     lowercase	:       Dict=4			,     )						->					Optional[Any]:
									"""simple docstring"""







									__lowercase					=				parent
									__lowercase					=				batch_size
									__lowercase					=				seq_length
									__lowercase					=				is_training
									__lowercase					=				use_attention_mask
									__lowercase					=				use_token_type_ids
									__lowercase					=				use_labels
									__lowercase					=				vocab_size
									__lowercase					=				hidden_size
									__lowercase					=				num_hidden_layers
									__lowercase					=				num_attention_heads
									__lowercase					=				intermediate_size
									__lowercase					=				hidden_act
									__lowercase					=				hidden_dropout_prob
									__lowercase					=				attention_probs_dropout_prob
									__lowercase					=				max_position_embeddings
									__lowercase					=				type_vocab_size
									__lowercase					=				type_sequence_label_size
									__lowercase					=				initializer_range
									__lowercase					=				num_choices





			def 		snake_case__			(				self	:       Optional[int]   )						->					List[str]:
									"""simple docstring"""







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

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

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

									__lowercase					=				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=lowercase			,     initializer_range=self.initializer_range			,     )

									return config, input_ids, token_type_ids, attention_mask







			def 		snake_case__			(				self	:       Optional[Any]   )						->					Any:
									"""simple docstring"""







									__lowercase					=				self.prepare_config_and_inputs()
									__lowercase	,			__lowercase	,			__lowercase	,			__lowercase					=				config_and_inputs
									__lowercase					=				{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
									return config, inputs_dict





@require_flax
class 			_lowerCAmelCase   ( _UpperCAmelCase    ,					unittest.TestCase			):






			"""simple docstring"""





			lowercase__      :  Tuple      =			(
			    (
			        FlaxAlbertModel,
			        FlaxAlbertForPreTraining,
			        FlaxAlbertForMaskedLM,
			        FlaxAlbertForMultipleChoice,
			        FlaxAlbertForQuestionAnswering,
			        FlaxAlbertForSequenceClassification,
			        FlaxAlbertForTokenClassification,
			        FlaxAlbertForQuestionAnswering,
			    )
			    if is_flax_available()
			    else ()
			)





			def 		snake_case__			(				self	:       Tuple   )						->					Any:
									"""simple docstring"""







									__lowercase					=				FlaxAlbertModelTester(self   )







			@slow
			def 		snake_case__			(				self	:       List[str]   )						->					Optional[int]:
									"""simple docstring"""







									for model_class_name in self.all_model_classes:
															__lowercase					=				model_class_name.from_pretrained("""albert-base-v2"""   )
															__lowercase					=				model(np.ones((1, 1)   )   )
															self.assertIsNotNone(lowercase   )





@require_flax
class 			_lowerCAmelCase   ( unittest.TestCase			):






			"""simple docstring"""





			@slow
			def 		snake_case__			(				self	:       str   )						->					Dict:
									"""simple docstring"""







									__lowercase					=				FlaxAlbertModel.from_pretrained("""albert-base-v2"""   )
									__lowercase					=				np.array([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]   )
									__lowercase					=				np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]   )
									__lowercase					=				model(lowercase			,     attention_mask=lowercase   )[0]
									__lowercase					=				(1, 11, 768)
									self.assertEqual(output.shape			,     lowercase   )
									__lowercase					=				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]			,     lowercase			,     atol=1E-4   )   )










import importlib
import os
from dataclasses import dataclass
from enum import Enum
from typing import Any, Dict, Optional, Union

import torch

from ..utils import BaseOutput


UpperCamelCase__ =						"scheduler_config.json"
class 			_lowerCAmelCase   ( _UpperCAmelCase			):






			"""simple docstring"""





			lowercase__      :  List[Any]      =			1
			lowercase__      :  Tuple      =			2
			lowercase__      :  Union[str, Any]      =			3
			lowercase__      :  Union[str, Any]      =			4
			lowercase__      :  str      =			5
			lowercase__      :  Any      =			6
			lowercase__      :  Any      =			7
			lowercase__      :  List[str]      =			8
			lowercase__      :  Union[str, Any]      =			9
			lowercase__      :  int      =			10
			lowercase__      :  List[str]      =			11
			lowercase__      :  List[Any]      =			12
			lowercase__      :  str      =			13
			lowercase__      :  Optional[int]      =			14





@dataclass
class 			_lowerCAmelCase   ( _UpperCAmelCase			):






			"""simple docstring"""





			lowercase__      :  torch.FloatTensor





class 			_lowerCAmelCase   :






			"""simple docstring"""





			lowercase__      :  Optional[int]      =			SCHEDULER_CONFIG_NAME
			lowercase__      :  int      =			[]
			lowercase__      :  Dict      =			True





			@classmethod
			def 		snake_case__			(				cls	:       str			,     lowercase	:       Dict[str, Any] = None			,     lowercase	:       Optional[str] = None			,     lowercase	:       Any=False			,     **lowercase	:       List[str]			,     )						->					Union[str, Any]:
									"""simple docstring"""







									__lowercase	,			__lowercase	,			__lowercase					=				cls.load_config(
									    pretrained_model_name_or_path=lowercase			,     subfolder=lowercase			,     return_unused_kwargs=lowercase			,     return_commit_hash=lowercase			,     **lowercase			,     )
									return cls.from_config(lowercase			,     return_unused_kwargs=lowercase			,     **lowercase   )





			def 		snake_case__			(				self	:       Dict			,     lowercase	:       Union[str, os.PathLike]			,     lowercase	:       bool = False			,     **lowercase	:       List[str]   )						->					Optional[Any]:
									"""simple docstring"""







									self.save_config(save_directory=lowercase			,     push_to_hub=lowercase			,     **lowercase   )





			@property
			def 		snake_case__			(				self	:       Optional[Any]   )						->					Dict:
									"""simple docstring"""







									return self._get_compatibles()







			@classmethod
			def 		snake_case__			(				cls	:       List[str]   )						->					Union[str, Any]:
									"""simple docstring"""







									__lowercase					=				list(set([cls.__name__] + cls._compatibles   )   )
									__lowercase					=				importlib.import_module(__name__.split("""."""   )[0]   )
									__lowercase					=				[
									    getattr(lowercase			,     lowercase   ) for c in compatible_classes_str if hasattr(lowercase			,     lowercase   )
									]
									return compatible_classes





