




"""simple docstring"""





import unittest

import numpy as np

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
					from transformers.models.bert.modeling_flax_bert import (
					    FlaxBertForMaskedLM,
					    FlaxBertForMultipleChoice,
					    FlaxBertForNextSentencePrediction,
					    FlaxBertForPreTraining,
					    FlaxBertForQuestionAnswering,
					    FlaxBertForSequenceClassification,
					    FlaxBertForTokenClassification,
					    FlaxBertModel,
					)



class   a__					(						unittest.TestCase ):


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


						"""simple docstring"""




						__UpperCAmelCase						:			str												=       parent
						__UpperCAmelCase						:			Dict												=       batch_size
						__UpperCAmelCase						:			Dict												=       seq_length
						__UpperCAmelCase						:			List[str]												=       is_training
						__UpperCAmelCase						:			List[Any]												=       use_attention_mask
						__UpperCAmelCase						:			Union[str, Any]												=       use_token_type_ids
						__UpperCAmelCase						:			Dict												=       use_labels
						__UpperCAmelCase						:			Optional[Any]												=       vocab_size
						__UpperCAmelCase						:			List[Any]												=       hidden_size
						__UpperCAmelCase						:			Optional[int]												=       num_hidden_layers
						__UpperCAmelCase						:			List[Any]												=       num_attention_heads
						__UpperCAmelCase						:			int												=       intermediate_size
						__UpperCAmelCase						:			Union[str, Any]												=       hidden_act
						__UpperCAmelCase						:			Dict												=       hidden_dropout_prob
						__UpperCAmelCase						:			Optional[Any]												=       attention_probs_dropout_prob
						__UpperCAmelCase						:			Any												=       max_position_embeddings
						__UpperCAmelCase						:			Any												=       type_vocab_size
						__UpperCAmelCase						:			Optional[int]												=       type_sequence_label_size
						__UpperCAmelCase						:			Any												=       initializer_range
						__UpperCAmelCase						:			str												=       num_choices


			def    a_  (   self   :     int):


						"""simple docstring"""




						__UpperCAmelCase						:			List[Any]												=       ids_tensor([self.batch_size, self.seq_length]   ,	self.vocab_size)

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

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

						__UpperCAmelCase						:			Union[str, Any]												=       BertConfig(
						    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=UpperCamelCase_   ,	initializer_range=self.initializer_range   ,	)

						return config, input_ids, token_type_ids, attention_mask


			def    a_  (   self   :     str):


						"""simple docstring"""




						__UpperCAmelCase						:			Union[str, Any]												=       self.prepare_config_and_inputs()
						__UpperCAmelCase						,  __UpperCAmelCase						,  __UpperCAmelCase						,  __UpperCAmelCase						:			List[Any]												=       config_and_inputs
						__UpperCAmelCase						:			int												=       {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
						return config, inputs_dict


			def    a_  (   self   :     Dict):


						"""simple docstring"""




						__UpperCAmelCase						:			int												=       self.prepare_config_and_inputs()
						__UpperCAmelCase						,  __UpperCAmelCase						,  __UpperCAmelCase						,  __UpperCAmelCase						:			Tuple												=       config_and_inputs

						__UpperCAmelCase						:			str												=       True
						__UpperCAmelCase						:			int												=       floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
						__UpperCAmelCase						:			List[str]												=       ids_tensor([self.batch_size, self.seq_length]   ,	vocab_size=2)

						return (
						    config,
						    input_ids,
						    attention_mask,
						    encoder_hidden_states,
						    encoder_attention_mask,
						)






@require_flax
class   a__					(						__magic_name__				,     unittest.TestCase ):
			lowercase_		=						True

			lowercase_		=						(
			    (
			        FlaxBertModel,
			        FlaxBertForPreTraining,
			        FlaxBertForMaskedLM,
			        FlaxBertForMultipleChoice,
			        FlaxBertForQuestionAnswering,
			        FlaxBertForNextSentencePrediction,
			        FlaxBertForSequenceClassification,
			        FlaxBertForTokenClassification,
			        FlaxBertForQuestionAnswering,
			    )
			    if is_flax_available()
			    else ()
			)


			def    a_  (   self   :     str):


						"""simple docstring"""




						__UpperCAmelCase						:			List[Any]												=       FlaxBertModelTester(self)


			@slow
			def    a_  (   self   :     Dict):


						"""simple docstring"""




						__UpperCAmelCase						:			Any												=       FlaxBertModel.from_pretrained("bert-base-cased")
						__UpperCAmelCase						:			Union[str, Any]												=       model(np.ones((1, 1)))
						self.assertIsNotNone(UpperCamelCase_)







"""simple docstring"""







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




