


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





     def __init__(    self					:				Any    ,	SCREAMING_SNAKE_CASE_					:				Tuple    ,	SCREAMING_SNAKE_CASE_					:				Optional[int]=13    ,	SCREAMING_SNAKE_CASE_					:				Union[str, Any]=7    ,	SCREAMING_SNAKE_CASE_					:				Optional[int]=True    ,	SCREAMING_SNAKE_CASE_					:				Optional[int]=True    ,	SCREAMING_SNAKE_CASE_					:				Optional[int]=True    ,	SCREAMING_SNAKE_CASE_					:				Optional[Any]=True    ,	SCREAMING_SNAKE_CASE_					:				Tuple=99    ,	SCREAMING_SNAKE_CASE_					:				List[Any]=32    ,	SCREAMING_SNAKE_CASE_					:				List[str]=5    ,	SCREAMING_SNAKE_CASE_					:				Any=4    ,	SCREAMING_SNAKE_CASE_					:				Union[str, Any]=37    ,	SCREAMING_SNAKE_CASE_					:				str="gelu"    ,	SCREAMING_SNAKE_CASE_					:				List[Any]=0.1    ,	SCREAMING_SNAKE_CASE_					:				Dict=0.1    ,	SCREAMING_SNAKE_CASE_					:				Dict=512    ,	SCREAMING_SNAKE_CASE_					:				Any=16    ,	SCREAMING_SNAKE_CASE_					:				int=2    ,	SCREAMING_SNAKE_CASE_					:				str=0.02    ,	SCREAMING_SNAKE_CASE_					:				Tuple=4    ,	):
        __lowerCamelCase:				str     =	parent
        __lowerCamelCase:				Union[str, Any]     =	batch_size
        __lowerCamelCase:				List[Any]     =	seq_length
        __lowerCamelCase:				List[Any]     =	is_training
        __lowerCamelCase:				Any     =	use_attention_mask
        __lowerCamelCase:				str     =	use_token_type_ids
        __lowerCamelCase:				Optional[Any]     =	use_labels
        __lowerCamelCase:				List[str]     =	vocab_size
        __lowerCamelCase:				Union[str, Any]     =	hidden_size
        __lowerCamelCase:				Tuple     =	num_hidden_layers
        __lowerCamelCase:				Optional[int]     =	num_attention_heads
        __lowerCamelCase:				Dict     =	intermediate_size
        __lowerCamelCase:				str     =	hidden_act
        __lowerCamelCase:				Optional[Any]     =	hidden_dropout_prob
        __lowerCamelCase:				Optional[int]     =	attention_probs_dropout_prob
        __lowerCamelCase:				Dict     =	max_position_embeddings
        __lowerCamelCase:				str     =	type_vocab_size
        __lowerCamelCase:				str     =	type_sequence_label_size
        __lowerCamelCase:				List[Any]     =	initializer_range
        __lowerCamelCase:				Optional[int]     =	num_choices





     def  SCREAMING_SNAKE_CASE__       (    self					:				int							):
        __lowerCamelCase:				Optional[int]     =	ids_tensor([self.batch_size, self.seq_length]    ,	self.vocab_size							)

        __lowerCamelCase:				Union[str, Any]     =	None
        if self.use_attention_mask:
           __lowerCamelCase:				List[Any]     =	random_attention_mask([self.batch_size, self.seq_length]							)

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

        __lowerCamelCase:				List[str]     =	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=SCREAMING_SNAKE_CASE_    ,	initializer_range=self.initializer_range    ,	)

        return config, input_ids, token_type_ids, attention_mask







     def  SCREAMING_SNAKE_CASE__       (    self					:				Optional[int]							):
        __lowerCamelCase:				Dict     =	self.prepare_config_and_inputs()
        __lowerCamelCase   ,     __lowerCamelCase   ,     __lowerCamelCase   ,     __lowerCamelCase:				Any     =	config_and_inputs
        __lowerCamelCase:				Optional[int]     =	{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
        return config, inputs_dict



@require_flax
class        a					(			_UpperCAmelCase     ,unittest.TestCase   ):
     UpperCAmelCase__ :  Union[str, Any]						      =	(
         (
             FlaxAlbertModel,
             FlaxAlbertForPreTraining,
             FlaxAlbertForMaskedLM,
             FlaxAlbertForMultipleChoice,
             FlaxAlbertForQuestionAnswering,
             FlaxAlbertForSequenceClassification,
             FlaxAlbertForTokenClassification,
             FlaxAlbertForQuestionAnswering,
         )
         if is_flax_available()
         else ()
     )





     def  SCREAMING_SNAKE_CASE__       (    self					:				int							):
        __lowerCamelCase:				Optional[Any]     =	FlaxAlbertModelTester(self							)







     @slow
     def  SCREAMING_SNAKE_CASE__       (    self					:				List[Any]							):
        for model_class_name in self.all_model_classes:
           __lowerCamelCase:				Union[str, Any]     =	model_class_name.from_pretrained("""albert-base-v2"""							)
           __lowerCamelCase:				List[str]     =	model(np.ones((1, 1)							)							)
           self.assertIsNotNone(SCREAMING_SNAKE_CASE_							)



@require_flax
class        a					(			unittest.TestCase   ):





     @slow
     def  SCREAMING_SNAKE_CASE__       (    self					:				Tuple							):
        __lowerCamelCase:				List[Any]     =	FlaxAlbertModel.from_pretrained("""albert-base-v2"""							)
        __lowerCamelCase:				int     =	np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]							)
        __lowerCamelCase:				str     =	np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]							)
        __lowerCamelCase:				Any     =	model(SCREAMING_SNAKE_CASE_    ,	attention_mask=SCREAMING_SNAKE_CASE_							)[0]
        __lowerCamelCase:				Dict     =	(1, 11, 768)
        self.assertEqual(output.shape    ,	SCREAMING_SNAKE_CASE_							)
        __lowerCamelCase:				List[str]     =	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]    ,	SCREAMING_SNAKE_CASE_    ,	atol=1E-4							)							)






from ..utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_scipy_available,
    is_torch_available,
    is_torchsde_available,
)


try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ..utils.dummy_pt_objects import *  # noqa F403
else:
   from .scheduling_consistency_models import CMStochasticIterativeScheduler
   from .scheduling_ddim import DDIMScheduler
   from .scheduling_ddim_inverse import DDIMInverseScheduler
   from .scheduling_ddim_parallel import DDIMParallelScheduler
   from .scheduling_ddpm import DDPMScheduler
   from .scheduling_ddpm_parallel import DDPMParallelScheduler
   from .scheduling_deis_multistep import DEISMultistepScheduler
   from .scheduling_dpmsolver_multistep import DPMSolverMultistepScheduler
   from .scheduling_dpmsolver_multistep_inverse import DPMSolverMultistepInverseScheduler
   from .scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
   from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler
   from .scheduling_euler_discrete import EulerDiscreteScheduler
   from .scheduling_heun_discrete import HeunDiscreteScheduler
   from .scheduling_ipndm import IPNDMScheduler
   from .scheduling_k_dpm_2_ancestral_discrete import KDPMaAncestralDiscreteScheduler
   from .scheduling_k_dpm_2_discrete import KDPMaDiscreteScheduler
   from .scheduling_karras_ve import KarrasVeScheduler
   from .scheduling_pndm import PNDMScheduler
   from .scheduling_repaint import RePaintScheduler
   from .scheduling_sde_ve import ScoreSdeVeScheduler
   from .scheduling_sde_vp import ScoreSdeVpScheduler
   from .scheduling_unclip import UnCLIPScheduler
   from .scheduling_unipc_multistep import UniPCMultistepScheduler
   from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
   from .scheduling_vq_diffusion import VQDiffusionScheduler

try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ..utils.dummy_flax_objects import *  # noqa F403
else:
   from .scheduling_ddim_flax import FlaxDDIMScheduler
   from .scheduling_ddpm_flax import FlaxDDPMScheduler
   from .scheduling_dpmsolver_multistep_flax import FlaxDPMSolverMultistepScheduler
   from .scheduling_karras_ve_flax import FlaxKarrasVeScheduler
   from .scheduling_lms_discrete_flax import FlaxLMSDiscreteScheduler
   from .scheduling_pndm_flax import FlaxPNDMScheduler
   from .scheduling_sde_ve_flax import FlaxScoreSdeVeScheduler
   from .scheduling_utils_flax import (
       FlaxKarrasDiffusionSchedulers,
       FlaxSchedulerMixin,
       FlaxSchedulerOutput,
       broadcast_to_shape_from_left,
   )


try:
   if not (is_torch_available() and is_scipy_available()):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ..utils.dummy_torch_and_scipy_objects import *  # noqa F403
else:
   from .scheduling_lms_discrete import LMSDiscreteScheduler

try:
   if not (is_torch_available() and is_torchsde_available()):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ..utils.dummy_torch_and_torchsde_objects import *  # noqa F403
else:
   from .scheduling_dpmsolver_sde import DPMSolverSDEScheduler


