



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





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

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

           __a       =		rescale_embeddings
           __a       =		attention_type
           __a       =		use_bias
           __a       =		block_size
           __a       =		num_random_blocks





       def 						UpperCamelCase_					(   self		:		List[Any]	):

           '''simple docstring'''
           __a       =		ids_tensor([self.batch_size, self.seq_length]      ,				self.vocab_size	)

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

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

           __a       =		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=_SCREAMING_SNAKE_CASE      ,				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		:		Union[str, Any]	):

           '''simple docstring'''
           __a       =		self.prepare_config_and_inputs()
           __a     ,					__a     ,					__a     ,					__a       =		config_and_inputs
           __a       =		{
               """input_ids""": input_ids,
               """token_type_ids""": token_type_ids,
               """attention_mask""": attention_mask,
           }
           return config, inputs_dict






@require_flax
class    SCREAMING_SNAKE_CASE			(							lowerCamelCase__			,      unittest.TestCase					):
       __lowerCamelCase	:	Optional[int]    =(
           (
               FlaxBigBirdForCausalLM,
               FlaxBigBirdModel,
               FlaxBigBirdForPreTraining,
               FlaxBigBirdForMaskedLM,
               FlaxBigBirdForMultipleChoice,
               FlaxBigBirdForQuestionAnswering,
               FlaxBigBirdForSequenceClassification,
               FlaxBigBirdForTokenClassification,
           )
           if is_flax_available()
           else ()
       )

       __lowerCamelCase	:	Tuple    =False
       __lowerCamelCase	:	Tuple    =False





       def 						UpperCamelCase_					(   self		:		Optional[Any]	):

           '''simple docstring'''
           __a       =		FlaxBigBirdModelTester(self	)





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

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

           '''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		:		Optional[int]	):

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

           '''simple docstring'''
           super().test_hidden_states_output()





       @slow
       def 						UpperCamelCase_					(   self		:		Optional[int]	):

           '''simple docstring'''






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





       def 						UpperCamelCase_					(   self		:		Union[str, Any]	):

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

           '''simple docstring'''
           __a     ,					__a       =		self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
               with self.subTest(model_class.__name__	):
                   __a       =		self._prepare_for_class(_SCREAMING_SNAKE_CASE      ,				_SCREAMING_SNAKE_CASE	)
                   __a       =		model_class(_SCREAMING_SNAKE_CASE	)

                   @jax.jit
                   def model_jitted(__lowercase		:		Any      ,				__lowercase		:		List[str]=None      ,				**__lowercase		:		List[Any]	):
                       return model(input_ids=_SCREAMING_SNAKE_CASE      ,				attention_mask=_SCREAMING_SNAKE_CASE      ,				**_SCREAMING_SNAKE_CASE	)

                   with self.subTest("""JIT Enabled"""	):
                       __a       =		model_jitted(**_SCREAMING_SNAKE_CASE	).to_tuple()

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

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







       def 						UpperCamelCase_					(   self		:		Tuple      ,				__lowercase		:		Optional[Any]      ,				__lowercase		:		Dict      ,				__lowercase		:		List[Any]      ,				__lowercase		:		int=1E-5      ,				__lowercase		:		List[str]="outputs"      ,				__lowercase		:		int=None	):

           '''simple docstring'''






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





from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .modeling_utils import ModelMixin






@dataclass
class    SCREAMING_SNAKE_CASE			(							lowerCamelCase__					):
							__lowerCamelCase	:	torch.FloatTensor






class    SCREAMING_SNAKE_CASE			(							lowerCamelCase__			,      lowerCamelCase__					):





							@register_to_config
							def __init__(   self		:		Tuple      ,				__lowercase		:		int = 16      ,				__lowercase		:		int = 88      ,				__lowercase		:		Optional[int] = None      ,				__lowercase		:		Optional[int] = None      ,				__lowercase		:		int = 1      ,				__lowercase		:		float = 0.0      ,				__lowercase		:		int = 32      ,				__lowercase		:		Optional[int] = None      ,				__lowercase		:		bool = False      ,				__lowercase		:		Optional[int] = None      ,				__lowercase		:		str = "geglu"      ,				__lowercase		:		bool = True      ,				__lowercase		:		bool = True      ,				):

											'''simple docstring'''
											super().__init__()
											__a       =		num_attention_heads
											__a       =		attention_head_dim
											__a       =		num_attention_heads * attention_head_dim

											__a       =		in_channels

											__a       =		torch.nn.GroupNorm(num_groups=__lowercase      ,				num_channels=__lowercase      ,				eps=1E-6      ,				affine=__lowercase	)
											__a       =		nn.Linear(__lowercase      ,				__lowercase	)

											# 3. Define transformers blocks
											__a       =		nn.ModuleList(
											    [
											        BasicTransformerBlock(
											            __lowercase      ,				__lowercase      ,				__lowercase      ,				dropout=__lowercase      ,				cross_attention_dim=__lowercase      ,				activation_fn=__lowercase      ,				attention_bias=__lowercase      ,				double_self_attention=__lowercase      ,				norm_elementwise_affine=__lowercase      ,				)
											        for d in range(__lowercase	)
											    ]	)

											__a       =		nn.Linear(__lowercase      ,				__lowercase	)







							def 						UpperCamelCase_					(   self		:		Any      ,				__lowercase		:		Optional[Any]      ,				__lowercase		:		str=None      ,				__lowercase		:		List[Any]=None      ,				__lowercase		:		List[Any]=None      ,				__lowercase		:		int=1      ,				__lowercase		:		Union[str, Any]=None      ,				__lowercase		:		bool = True      ,				):

											'''simple docstring'''
											__a     ,					__a     ,					__a     ,					__a       =		hidden_states.shape
											__a       =		batch_frames // num_frames

											__a       =		hidden_states

											__a       =		hidden_states[None, :].reshape(__lowercase      ,				__lowercase      ,				__lowercase      ,				__lowercase      ,				__lowercase	)
											__a       =		hidden_states.permute(0      ,				2      ,				1      ,				3      ,				4	)

											__a       =		self.norm(__lowercase	)
											__a       =		hidden_states.permute(0      ,				3      ,				4      ,				2      ,				1	).reshape(batch_size * height * width      ,				__lowercase      ,				__lowercase	)

											__a       =		self.proj_in(__lowercase	)

											# 2. Blocks
											for block in self.transformer_blocks:
															__a       =		block(
															    __lowercase      ,				encoder_hidden_states=__lowercase      ,				timestep=__lowercase      ,				cross_attention_kwargs=__lowercase      ,				class_labels=__lowercase      ,				)

											# 3. Output
											__a       =		self.proj_out(__lowercase	)
											__a       =		(
											    hidden_states[None, None, :]
											    .reshape(__lowercase      ,				__lowercase      ,				__lowercase      ,				__lowercase      ,				__lowercase	)
											    .permute(0      ,				3      ,				4      ,				1      ,				2	)
											    .contiguous()
											)
											__a       =		hidden_states.reshape(__lowercase      ,				__lowercase      ,				__lowercase      ,				__lowercase	)

											__a       =		hidden_states + residual

											if not return_dict:
															return (output,)

											return TransformerTemporalModelOutput(sample=__lowercase	)
