






'''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    ):






  """simple docstring"""





  def __init__(     self		,		_UpperCAmelCase		,		_UpperCAmelCase=13		,		_UpperCAmelCase=7		,		_UpperCAmelCase=True		,		_UpperCAmelCase=True		,		_UpperCAmelCase=True		,		_UpperCAmelCase=True		,		_UpperCAmelCase=99		,		_UpperCAmelCase=32		,		_UpperCAmelCase=5		,		_UpperCAmelCase=4		,		_UpperCAmelCase=37		,		_UpperCAmelCase="gelu"		,		_UpperCAmelCase=0.1		,		_UpperCAmelCase=0.1		,		_UpperCAmelCase=512		,		_UpperCAmelCase=16		,		_UpperCAmelCase=2		,		_UpperCAmelCase=0.0_2		,		_UpperCAmelCase=4		,		)			-> List[str]:
     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    _UpperCAmelCase     (     self )			-> List[Any]:
     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_									=     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=_snake_case		,		initializer_range=self.initializer_range		,		)

     return config, input_ids, token_type_ids, attention_mask




  def    _UpperCAmelCase     (     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




  def    _UpperCAmelCase     (     self )			-> str:
     UpperCamelCase_									=     self.prepare_config_and_inputs()
     UpperCamelCase_		,      UpperCamelCase_		,      UpperCamelCase_		,      UpperCamelCase_									=     config_and_inputs

     UpperCamelCase_									=     True
     UpperCamelCase_									=     floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
     UpperCamelCase_									=     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       (  UpperCAmelCase_						,       unittest.TestCase    ):






  """simple docstring"""





  A_						=   True

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




  def    _UpperCAmelCase     (     self )			-> List[str]:
     UpperCamelCase_									=     FlaxBertModelTester(self )




  @slow
  def    _UpperCAmelCase     (     self )			-> Tuple:
     UpperCamelCase_									=     FlaxBertModel.from_pretrained('bert-base-cased' )
     UpperCamelCase_									=     model(np.ones((1, 1) ) )
     self.assertIsNotNone(_snake_case )













from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
 if not (is_transformers_available() and is_torch_available()):
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 from ...utils.dummy_torch_and_transformers_objects import (
     ImageTextPipelineOutput,
     UniDiffuserPipeline,
 )
else:
 from .modeling_text_decoder import UniDiffuserTextDecoder
 from .modeling_uvit import UniDiffuserModel, UTransformeraDModel
 from .pipeline_unidiffuser import ImageTextPipelineOutput, UniDiffuserPipeline





