import gc
import unittest

import numpy as np
import torch
import torch.nn.functional as F
from transformers import (
    ClapTextConfig,
    ClapTextModelWithProjection,
    RobertaTokenizer,
    SpeechTaHifiGan,
    SpeechTaHifiGanConfig,
)

from diffusers import (
    AudioLDMPipeline,
    AutoencoderKL,
    DDIMScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from ..pipeline_params import TEXT_TO_AUDIO_BATCH_PARAMS, TEXT_TO_AUDIO_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()




class 					_UpperCAmelCase			(     lowerCAmelCase,			unittest.TestCase  ):



 '''simple docstring'''




 __A			       =					AudioLDMPipeline
 __A			       =					TEXT_TO_AUDIO_PARAMS
 __A			       =					TEXT_TO_AUDIO_BATCH_PARAMS
 __A			       =					frozenset(
     [
         '''num_inference_steps''',
         '''num_waveforms_per_prompt''',
         '''generator''',
         '''latents''',
         '''output_type''',
         '''return_dict''',
         '''callback''',
         '''callback_steps''',
     ]  )







 def 		__UpperCAmelCase				(		self				:   Optional[int])						->	Optional[int]:


   """simple docstring"""







   torch.manual_seed(0)
   _UpperCamelCase           =    UNetaDConditionModel(
       block_out_channels=(32, 64)							, layers_per_block=2							, sample_size=32							, in_channels=4							, out_channels=4							, down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")							, up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")							, cross_attention_dim=(32, 64)							, class_embed_type="simple_projection"							, projection_class_embeddings_input_dim=32							, class_embeddings_concat=lowercase_							, )
   _UpperCamelCase           =    DDIMScheduler(
       beta_start=0.0_00_85							, beta_end=0.0_12							, beta_schedule="scaled_linear"							, clip_sample=lowercase_							, set_alpha_to_one=lowercase_							, )
   torch.manual_seed(0)
   _UpperCamelCase           =    AutoencoderKL(
       block_out_channels=[32, 64]							, in_channels=1							, out_channels=1							, down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]							, up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]							, latent_channels=4							, )
   torch.manual_seed(0)
   _UpperCamelCase           =    ClapTextConfig(
       bos_token_id=0							, eos_token_id=2							, hidden_size=32							, intermediate_size=37							, layer_norm_eps=1e-0_5							, num_attention_heads=4							, num_hidden_layers=5							, pad_token_id=1							, vocab_size=1000							, projection_dim=32							, )
   _UpperCamelCase           =    ClapTextModelWithProjection(lowercase_)
   _UpperCamelCase           =    RobertaTokenizer.from_pretrained("hf-internal-testing/tiny-random-roberta"							, model_max_length=77)

   _UpperCamelCase           =    SpeechTaHifiGanConfig(
       model_in_dim=8							, sampling_rate=16000							, upsample_initial_channel=16							, upsample_rates=[2, 2]							, upsample_kernel_sizes=[4, 4]							, resblock_kernel_sizes=[3, 7]							, resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5]]							, normalize_before=lowercase_							, )

   _UpperCamelCase           =    SpeechTaHifiGan(lowercase_)

   _UpperCamelCase           =    {
       "unet": unet,
       "scheduler": scheduler,
       "vae": vae,
       "text_encoder": text_encoder,
       "tokenizer": tokenizer,
       "vocoder": vocoder,
   }
   return components







 def 		__UpperCAmelCase				(		self				:   Optional[int]							, lowercase_				:   int							, lowercase_				:   int=0)						->	List[Any]:


   """simple docstring"""







   if str(lowercase_).startswith("mps"):
     _UpperCamelCase           =    torch.manual_seed(lowercase_)
   else:
     _UpperCamelCase           =    torch.Generator(device=lowercase_).manual_seed(lowercase_)
   _UpperCamelCase           =    {
       "prompt": "A hammer hitting a wooden surface",
       "generator": generator,
       "num_inference_steps": 2,
       "guidance_scale": 6.0,
   }
   return inputs







 def 		__UpperCAmelCase				(		self				:   List[str])						->	Union[str, Any]:


   """simple docstring"""







   _UpperCamelCase           =    "cpu"  # ensure determinism for the device-dependent torch.Generator

   _UpperCamelCase           =    self.get_dummy_components()
   _UpperCamelCase           =    AudioLDMPipeline(**lowercase_)
   _UpperCamelCase           =    audioldm_pipe.to(lowercase_)
   audioldm_pipe.set_progress_bar_config(disable=lowercase_)

   _UpperCamelCase           =    self.get_dummy_inputs(lowercase_)
   _UpperCamelCase           =    audioldm_pipe(**lowercase_)
   _UpperCamelCase           =    output.audios[0]

   assert audio.ndim == 1
   assert len(lowercase_) == 256

   _UpperCamelCase           =    audio[:10]
   _UpperCamelCase           =    np.array(
       [-0.00_50, 0.00_50, -0.00_60, 0.00_33, -0.00_26, 0.00_33, -0.00_27, 0.00_33, -0.00_28, 0.00_33])

   assert np.abs(audio_slice - expected_slice).max() < 1e-2







 def 		__UpperCAmelCase				(		self				:   List[str])						->	List[Any]:


   """simple docstring"""







   _UpperCamelCase           =    self.get_dummy_components()
   _UpperCamelCase           =    AudioLDMPipeline(**lowercase_)
   _UpperCamelCase           =    audioldm_pipe.to(lowercase_)
   _UpperCamelCase           =    audioldm_pipe.to(lowercase_)
   audioldm_pipe.set_progress_bar_config(disable=lowercase_)

   _UpperCamelCase           =    self.get_dummy_inputs(lowercase_)
   _UpperCamelCase           =    3 * [inputs["prompt"]]

   # forward
   _UpperCamelCase           =    audioldm_pipe(**lowercase_)
   _UpperCamelCase           =    output.audios[0]

   _UpperCamelCase           =    self.get_dummy_inputs(lowercase_)
   _UpperCamelCase           =    3 * [inputs.pop("prompt")]

   _UpperCamelCase           =    audioldm_pipe.tokenizer(
       lowercase_							, padding="max_length"							, max_length=audioldm_pipe.tokenizer.model_max_length							, truncation=lowercase_							, return_tensors="pt"							, )
   _UpperCamelCase           =    text_inputs["input_ids"].to(lowercase_)

   _UpperCamelCase           =    audioldm_pipe.text_encoder(
       lowercase_							, )
   _UpperCamelCase           =    prompt_embeds.text_embeds
   # additional L_2 normalization over each hidden-state
   _UpperCamelCase           =    F.normalize(lowercase_							, dim=-1)

   _UpperCamelCase           =    prompt_embeds

   # forward
   _UpperCamelCase           =    audioldm_pipe(**lowercase_)
   _UpperCamelCase           =    output.audios[0]

   assert np.abs(audio_a - audio_a).max() < 1e-2







 def 		__UpperCAmelCase				(		self				:   List[str])						->	Optional[Any]:


   """simple docstring"""







   _UpperCamelCase           =    self.get_dummy_components()
   _UpperCamelCase           =    AudioLDMPipeline(**lowercase_)
   _UpperCamelCase           =    audioldm_pipe.to(lowercase_)
   _UpperCamelCase           =    audioldm_pipe.to(lowercase_)
   audioldm_pipe.set_progress_bar_config(disable=lowercase_)

   _UpperCamelCase           =    self.get_dummy_inputs(lowercase_)
   _UpperCamelCase           =    3 * ["this is a negative prompt"]
   _UpperCamelCase           =    negative_prompt
   _UpperCamelCase           =    3 * [inputs["prompt"]]

   # forward
   _UpperCamelCase           =    audioldm_pipe(**lowercase_)
   _UpperCamelCase           =    output.audios[0]

   _UpperCamelCase           =    self.get_dummy_inputs(lowercase_)
   _UpperCamelCase           =    3 * [inputs.pop("prompt")]

   _UpperCamelCase           =    []
   for p in [prompt, negative_prompt]:
     _UpperCamelCase           =    audioldm_pipe.tokenizer(
         lowercase_							, padding="max_length"							, max_length=audioldm_pipe.tokenizer.model_max_length							, truncation=lowercase_							, return_tensors="pt"							, )
     _UpperCamelCase           =    text_inputs["input_ids"].to(lowercase_)

     _UpperCamelCase           =    audioldm_pipe.text_encoder(
         lowercase_							, )
     _UpperCamelCase           =    text_embeds.text_embeds
     # additional L_2 normalization over each hidden-state
     _UpperCamelCase           =    F.normalize(lowercase_							, dim=-1)

     embeds.append(lowercase_)

   _UpperCamelCase      ,						_UpperCamelCase           =    embeds

   # forward
   _UpperCamelCase           =    audioldm_pipe(**lowercase_)
   _UpperCamelCase           =    output.audios[0]

   assert np.abs(audio_a - audio_a).max() < 1e-2







 def 		__UpperCAmelCase				(		self				:   Tuple)						->	List[Any]:


   """simple docstring"""







   _UpperCamelCase           =    "cpu"  # ensure determinism for the device-dependent torch.Generator
   _UpperCamelCase           =    self.get_dummy_components()
   _UpperCamelCase           =    PNDMScheduler(skip_prk_steps=lowercase_)
   _UpperCamelCase           =    AudioLDMPipeline(**lowercase_)
   _UpperCamelCase           =    audioldm_pipe.to(lowercase_)
   audioldm_pipe.set_progress_bar_config(disable=lowercase_)

   _UpperCamelCase           =    self.get_dummy_inputs(lowercase_)
   _UpperCamelCase           =    "egg cracking"
   _UpperCamelCase           =    audioldm_pipe(**lowercase_							, negative_prompt=lowercase_)
   _UpperCamelCase           =    output.audios[0]

   assert audio.ndim == 1
   assert len(lowercase_) == 256

   _UpperCamelCase           =    audio[:10]
   _UpperCamelCase           =    np.array(
       [-0.00_51, 0.00_50, -0.00_60, 0.00_34, -0.00_26, 0.00_33, -0.00_27, 0.00_33, -0.00_28, 0.00_32])

   assert np.abs(audio_slice - expected_slice).max() < 1e-2







 def 		__UpperCAmelCase				(		self				:   Union[str, Any])						->	Any:


   """simple docstring"""







   _UpperCamelCase           =    "cpu"  # ensure determinism for the device-dependent torch.Generator
   _UpperCamelCase           =    self.get_dummy_components()
   _UpperCamelCase           =    PNDMScheduler(skip_prk_steps=lowercase_)
   _UpperCamelCase           =    AudioLDMPipeline(**lowercase_)
   _UpperCamelCase           =    audioldm_pipe.to(lowercase_)
   audioldm_pipe.set_progress_bar_config(disable=lowercase_)

   _UpperCamelCase           =    "A hammer hitting a wooden surface"

   # test num_waveforms_per_prompt=1 (default)
   _UpperCamelCase           =    audioldm_pipe(lowercase_							, num_inference_steps=2).audios

   assert audios.shape == (1, 256)

   # test num_waveforms_per_prompt=1 (default) for batch of prompts
   _UpperCamelCase           =    2
   _UpperCamelCase           =    audioldm_pipe([prompt] * batch_size							, num_inference_steps=2).audios

   assert audios.shape == (batch_size, 256)

   # test num_waveforms_per_prompt for single prompt
   _UpperCamelCase           =    2
   _UpperCamelCase           =    audioldm_pipe(lowercase_							, num_inference_steps=2							, num_waveforms_per_prompt=lowercase_).audios

   assert audios.shape == (num_waveforms_per_prompt, 256)

   # test num_waveforms_per_prompt for batch of prompts
   _UpperCamelCase           =    2
   _UpperCamelCase           =    audioldm_pipe(
       [prompt] * batch_size							, num_inference_steps=2							, num_waveforms_per_prompt=lowercase_).audios

   assert audios.shape == (batch_size * num_waveforms_per_prompt, 256)







 def 		__UpperCAmelCase				(		self				:   List[str])						->	str:


   """simple docstring"""







   _UpperCamelCase           =    "cpu"  # ensure determinism for the device-dependent torch.Generator
   _UpperCamelCase           =    self.get_dummy_components()
   _UpperCamelCase           =    AudioLDMPipeline(**lowercase_)
   _UpperCamelCase           =    audioldm_pipe.to(lowercase_)
   audioldm_pipe.set_progress_bar_config(disable=lowercase_)
   _UpperCamelCase           =    audioldm_pipe.vocoder.config.sampling_rate

   _UpperCamelCase           =    self.get_dummy_inputs(lowercase_)
   _UpperCamelCase           =    audioldm_pipe(audio_length_in_s=0.0_16							, **lowercase_)
   _UpperCamelCase           =    output.audios[0]

   assert audio.ndim == 1
   assert len(lowercase_) / vocoder_sampling_rate == 0.0_16

   _UpperCamelCase           =    audioldm_pipe(audio_length_in_s=0.0_32							, **lowercase_)
   _UpperCamelCase           =    output.audios[0]

   assert audio.ndim == 1
   assert len(lowercase_) / vocoder_sampling_rate == 0.0_32







 def 		__UpperCAmelCase				(		self				:   Dict)						->	Dict:


   """simple docstring"""







   _UpperCamelCase           =    self.get_dummy_components()
   _UpperCamelCase           =    AudioLDMPipeline(**lowercase_)
   _UpperCamelCase           =    audioldm_pipe.to(lowercase_)
   audioldm_pipe.set_progress_bar_config(disable=lowercase_)

   _UpperCamelCase           =    ["hey"]

   _UpperCamelCase           =    audioldm_pipe(lowercase_							, num_inference_steps=1)
   _UpperCamelCase           =    output.audios.shape
   assert audio_shape == (1, 256)

   _UpperCamelCase           =    audioldm_pipe.vocoder.config
   config.model_in_dim *= 2
   _UpperCamelCase           =    SpeechTaHifiGan(lowercase_).to(lowercase_)
   _UpperCamelCase           =    audioldm_pipe(lowercase_							, num_inference_steps=1)
   _UpperCamelCase           =    output.audios.shape
   # waveform shape is unchanged, we just have 2x the number of mel channels in the spectrogram
   assert audio_shape == (1, 256)







 def 		__UpperCAmelCase				(		self				:   List[str])						->	Optional[int]:


   """simple docstring"""







   self._test_attention_slicing_forward_pass(test_mean_pixel_difference=lowercase_)







 def 		__UpperCAmelCase				(		self				:   List[Any])						->	Union[str, Any]:


   """simple docstring"""







   self._test_inference_batch_single_identical(test_mean_pixel_difference=lowercase_)







 @unittest.skipIf(
     torch_device != "cuda" or not is_xformers_available()							, reason="XFormers attention is only available with CUDA and `xformers` installed"							, )
 def 		__UpperCAmelCase				(		self				:   Optional[Any])						->	Any:


   """simple docstring"""







   self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=lowercase_)




@slow
class 					_UpperCAmelCase			(     unittest.TestCase  ):



 '''simple docstring'''




 def 		__UpperCAmelCase				(		self				:   str)						->	Optional[int]:


   """simple docstring"""







   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()







 def 		__UpperCAmelCase				(		self				:   List[Any]							, lowercase_				:   List[str]							, lowercase_				:   List[str]="cpu"							, lowercase_				:   int=torch.floataa							, lowercase_				:   Tuple=0)						->	Optional[Any]:


   """simple docstring"""







   _UpperCamelCase           =    torch.Generator(device=lowercase_).manual_seed(lowercase_)
   _UpperCamelCase           =    np.random.RandomState(lowercase_).standard_normal((1, 8, 128, 16))
   _UpperCamelCase           =    torch.from_numpy(lowercase_).to(device=lowercase_							, dtype=lowercase_)
   _UpperCamelCase           =    {
       "prompt": "A hammer hitting a wooden surface",
       "latents": latents,
       "generator": generator,
       "num_inference_steps": 3,
       "guidance_scale": 2.5,
   }
   return inputs







 def 		__UpperCAmelCase				(		self				:   List[Any])						->	str:


   """simple docstring"""







   _UpperCamelCase           =    AudioLDMPipeline.from_pretrained("cvssp/audioldm")
   _UpperCamelCase           =    audioldm_pipe.to(lowercase_)
   audioldm_pipe.set_progress_bar_config(disable=lowercase_)

   _UpperCamelCase           =    self.get_inputs(lowercase_)
   _UpperCamelCase           =    25
   _UpperCamelCase           =    audioldm_pipe(**lowercase_).audios[0]

   assert audio.ndim == 1
   assert len(lowercase_) == 81920

   _UpperCamelCase           =    audio[77230:77240]
   _UpperCamelCase           =    np.array(
       [-0.48_84, -0.46_07, 0.00_23, 0.50_07, 0.58_96, 0.51_51, 0.38_13, -0.02_08, -0.36_87, -0.43_15])
   _UpperCamelCase           =    np.abs(expected_slice - audio_slice).max()
   assert max_diff < 1e-2







 def 		__UpperCAmelCase				(		self				:   str)						->	Optional[int]:


   """simple docstring"""







   _UpperCamelCase           =    AudioLDMPipeline.from_pretrained("cvssp/audioldm")
   _UpperCamelCase           =    LMSDiscreteScheduler.from_config(audioldm_pipe.scheduler.config)
   _UpperCamelCase           =    audioldm_pipe.to(lowercase_)
   audioldm_pipe.set_progress_bar_config(disable=lowercase_)

   _UpperCamelCase           =    self.get_inputs(lowercase_)
   _UpperCamelCase           =    audioldm_pipe(**lowercase_).audios[0]

   assert audio.ndim == 1
   assert len(lowercase_) == 81920

   _UpperCamelCase           =    audio[27780:27790]
   _UpperCamelCase           =    np.array([-0.21_31, -0.08_73, -0.01_24, -0.01_89, 0.05_69, 0.13_73, 0.18_83, 0.28_86, 0.32_97, 0.22_12])
   _UpperCamelCase           =    np.abs(expected_slice - audio_slice).max()
   assert max_diff < 3e-2





import unittest

from transformers import BigBirdConfig, 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

   from transformers.models.big_bird.modeling_flax_big_bird import (
       FlaxBigBirdForCausalLM,
       FlaxBigBirdForMaskedLM,
       FlaxBigBirdForMultipleChoice,
       FlaxBigBirdForPreTraining,
       FlaxBigBirdForQuestionAnswering,
       FlaxBigBirdForSequenceClassification,
       FlaxBigBirdForTokenClassification,
       FlaxBigBirdModel,
   )




class 					_UpperCAmelCase			(     unittest.TestCase  ):



 '''simple docstring'''




 def __init__(		self				:   List[Any]							, lowercase_				:   Union[str, Any]							, lowercase_				:   Union[str, Any]=2							, lowercase_				:   Dict=56							, lowercase_				:   str=True							, lowercase_				:   Optional[Any]=True							, lowercase_				:   List[Any]=True							, lowercase_				:   Union[str, Any]=True							, lowercase_				:   Any=99							, lowercase_				:   Optional[int]=32							, lowercase_				:   Tuple=2							, lowercase_				:   int=2							, lowercase_				:   List[str]=7							, lowercase_				:   Any="gelu_new"							, lowercase_				:   List[str]=0.1							, lowercase_				:   str=0.1							, lowercase_				:   List[Any]=512							, lowercase_				:   List[str]=16							, lowercase_				:   Optional[int]=2							, lowercase_				:   Union[str, Any]=0.02							, lowercase_				:   Union[str, Any]=4							, lowercase_				:   Union[str, Any]="block_sparse"							, lowercase_				:   Tuple=True							, lowercase_				:   Dict=False							, lowercase_				:   Dict=2							, lowercase_				:   Dict=3							, )						->	List[str]:


   """simple docstring"""







   _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

   _UpperCamelCase           =    rescale_embeddings
   _UpperCamelCase           =    attention_type
   _UpperCamelCase           =    use_bias
   _UpperCamelCase           =    block_size
   _UpperCamelCase           =    num_random_blocks







 def 		__UpperCAmelCase				(		self				:   List[Any])						->	Dict:


   """simple docstring"""







   _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           =    BigBirdConfig(
       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							, attention_type=self.attention_type							, block_size=self.block_size							, num_random_blocks=self.num_random_blocks							, use_bias=self.use_bias							, rescale_embeddings=self.rescale_embeddings							, )

   return config, input_ids, token_type_ids, attention_mask







 def 		__UpperCAmelCase				(		self				:   Optional[Any])						->	Optional[Any]:


   """simple docstring"""







   _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			(     lowerCAmelCase,			unittest.TestCase  ):



 '''simple docstring'''




 __A			       =					(
     (
         FlaxBigBirdForCausalLM,
         FlaxBigBirdModel,
         FlaxBigBirdForPreTraining,
         FlaxBigBirdForMaskedLM,
         FlaxBigBirdForMultipleChoice,
         FlaxBigBirdForQuestionAnswering,
         FlaxBigBirdForSequenceClassification,
         FlaxBigBirdForTokenClassification,
     )
     if is_flax_available()
     else ()
 )

 __A			       =					False
 __A			       =					False







 def 		__UpperCAmelCase				(		self				:   Optional[int])						->	Optional[Any]:


   """simple docstring"""







   _UpperCamelCase           =    FlaxBigBirdModelTester(self)







 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def 		__UpperCAmelCase				(		self				:   Union[str, Any])						->	List[str]:


   """simple docstring"""







   super().test_from_pretrained_save_pretrained()







 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def 		__UpperCAmelCase				(		self				:   Any)						->	int:


   """simple docstring"""







   super().test_from_pretrained_with_no_automatic_init()







 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def 		__UpperCAmelCase				(		self				:   List[Any])						->	Union[str, Any]:


   """simple docstring"""







   super().test_no_automatic_init()







 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def 		__UpperCAmelCase				(		self				:   str)						->	Union[str, Any]:


   """simple docstring"""







   super().test_hidden_states_output()







 @slow
 def 		__UpperCAmelCase				(		self				:   Tuple)						->	Optional[int]:


   """simple docstring"""







   for model_class_name in self.all_model_classes:
     _UpperCamelCase           =    model_class_name.from_pretrained("google/bigbird-roberta-base")
     self.assertIsNotNone(lowercase_)







 def 		__UpperCAmelCase				(		self				:   Tuple)						->	str:


   """simple docstring"""







   if self.test_attn_probs:
     super().test_attention_outputs()







 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def 		__UpperCAmelCase				(		self				:   Tuple)						->	Dict:


   """simple docstring"""







   _UpperCamelCase      ,						_UpperCamelCase           =    self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     with self.subTest(model_class.__name__):
       _UpperCamelCase           =    self._prepare_for_class(lowercase_							, lowercase_)
       _UpperCamelCase           =    model_class(lowercase_)

       @jax.jit
       def model_jitted(lowercase_				:   Dict							, lowercase_				:   List[Any]=None							, **lowercase_				:   Tuple):
         return model(input_ids=lowercase_							, attention_mask=lowercase_							, **lowercase_)

       with self.subTest("JIT Enabled"):
         _UpperCamelCase           =    model_jitted(**lowercase_).to_tuple()

       with self.subTest("JIT Disabled"):
         with jax.disable_jit():
           _UpperCamelCase           =    model_jitted(**lowercase_).to_tuple()

       self.assertEqual(len(lowercase_)							, len(lowercase_))
       for jitted_output, output in zip(lowercase_							, lowercase_):
         self.assertEqual(jitted_output.shape							, output.shape)







 def 		__UpperCAmelCase				(		self				:   Any							, lowercase_				:   List[str]							, lowercase_				:   List[str]							, lowercase_				:   List[str]							, lowercase_				:   str=1e-5							, lowercase_				:   int="outputs"							, lowercase_				:   List[str]=None)						->	Tuple:


   """simple docstring"""







   if name.startswith("outputs.attentions"):
     return
   else:
     super().check_pt_flax_outputs(lowercase_							, lowercase_							, lowercase_							, lowercase_							, lowercase_							, lowercase_)




