"""simple docstring"""




import gc
import math
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, logging, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


_a      :		Optional[int]           =   logging.get_logger(__name__)

enable_full_determinism()
class 						__A							( SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	unittest.TestCase		):
     _UpperCamelCase							:			Optional[int]							=		UNetaDModel
     _UpperCamelCase							:			int							=		"sample"






     @property
     def     __A  (		self  ):
       _lowerCAmelCase						:       Optional[int]  			=     4
       _lowerCAmelCase						:       Optional[int]  			=     3
       _lowerCAmelCase						:       List[str]  			=     (32, 32)

       _lowerCAmelCase						:       str  			=     floats_tensor((batch_size, num_channels) + sizes  ).to(a__  )
       _lowerCAmelCase						:       Dict  			=     torch.tensor([10]  ).to(a__  )

       return {"sample": noise, "timestep": time_step}






     @property
     def     __A  (		self  ):
       return (3, 32, 32)






     @property
     def     __A  (		self  ):
       return (3, 32, 32)






     def     __A  (		self  ):
       _lowerCAmelCase						:       str  			=     {
           """block_out_channels""": (32, 64),
           """down_block_types""": ("""DownBlock2D""", """AttnDownBlock2D"""),
           """up_block_types""": ("""AttnUpBlock2D""", """UpBlock2D"""),
           """attention_head_dim""": 3,
           """out_channels""": 3,
           """in_channels""": 3,
           """layers_per_block""": 2,
           """sample_size""": 32,
       }
       _lowerCAmelCase						:       int  			=     self.dummy_input
       return init_dict, inputs_dict







class 						__A							( SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	unittest.TestCase		):
     _UpperCamelCase							:			Tuple							=		UNetaDModel
     _UpperCamelCase							:			Dict							=		"sample"






     @property
     def     __A  (		self  ):
       _lowerCAmelCase						:       Optional[int]  			=     4
       _lowerCAmelCase						:       int  			=     4
       _lowerCAmelCase						:       Optional[int]  			=     (32, 32)

       _lowerCAmelCase						:       List[Any]  			=     floats_tensor((batch_size, num_channels) + sizes  ).to(a__  )
       _lowerCAmelCase						:       int  			=     torch.tensor([10]  ).to(a__  )

       return {"sample": noise, "timestep": time_step}






     @property
     def     __A  (		self  ):
       return (4, 32, 32)






     @property
     def     __A  (		self  ):
       return (4, 32, 32)






     def     __A  (		self  ):
       _lowerCAmelCase						:       Dict  			=     {
           """sample_size""": 32,
           """in_channels""": 4,
           """out_channels""": 4,
           """layers_per_block""": 2,
           """block_out_channels""": (32, 64),
           """attention_head_dim""": 32,
           """down_block_types""": ("""DownBlock2D""", """DownBlock2D"""),
           """up_block_types""": ("""UpBlock2D""", """UpBlock2D"""),
       }
       _lowerCAmelCase						:       int  			=     self.dummy_input
       return init_dict, inputs_dict






     def     __A  (		self  ):
       _lowerCAmelCase			,	_lowerCAmelCase						:       List[str]  			=     UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""				,     output_loading_info=a__  )

       self.assertIsNotNone(a__  )
       self.assertEqual(len(loading_info["""missing_keys"""]  )				,     0  )

       model.to(a__  )
       _lowerCAmelCase						:       Any  			=     model(**self.dummy_input  ).sample

       assert image is not None, "Make sure output is not None"






     @unittest.skipIf(torch_device != """cuda"""				,     """This test is supposed to run on GPU"""  )
     def     __A  (		self  ):
       _lowerCAmelCase			,	_lowerCAmelCase						:       Any  			=     UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""				,     output_loading_info=a__  )
       model.to(a__  )
       _lowerCAmelCase						:       Union[str, Any]  			=     model(**self.dummy_input  ).sample

       assert image is not None, "Make sure output is not None"






     @unittest.skipIf(torch_device != """cuda"""				,     """This test is supposed to run on GPU"""  )
     def     __A  (		self  ):
       # by defautl model loading will use accelerate as `low_cpu_mem_usage=True`
       _lowerCAmelCase			,	_lowerCAmelCase						:       List[Any]  			=     UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""				,     output_loading_info=a__  )
       model_accelerate.to(a__  )
       model_accelerate.eval()

       _lowerCAmelCase						:       int  			=     torch.randn(
           1				,     model_accelerate.config.in_channels				,     model_accelerate.config.sample_size				,     model_accelerate.config.sample_size				,     generator=torch.manual_seed(0  )				,     )
       _lowerCAmelCase						:       Optional[int]  			=     noise.to(a__  )
       _lowerCAmelCase						:       Dict  			=     torch.tensor([10] * noise.shape[0]  ).to(a__  )

       _lowerCAmelCase						:       Dict  			=     model_accelerate(a__				,     a__  )["""sample"""]

       # two models don't need to stay in the device at the same time
       del model_accelerate
       torch.cuda.empty_cache()
       gc.collect()

       _lowerCAmelCase			,	_lowerCAmelCase						:       Union[str, Any]  			=     UNetaDModel.from_pretrained(
           """fusing/unet-ldm-dummy-update"""				,     output_loading_info=a__				,     low_cpu_mem_usage=a__  )
       model_normal_load.to(a__  )
       model_normal_load.eval()
       _lowerCAmelCase						:       Dict  			=     model_normal_load(a__				,     a__  )["""sample"""]

       assert torch_all_close(a__				,     a__				,     rtol=1e-3  )






     def     __A  (		self  ):
       _lowerCAmelCase						:       Optional[Any]  			=     UNetaDModel.from_pretrained("""fusing/unet-ldm-dummy-update"""  )
       model.eval()
       model.to(a__  )

       _lowerCAmelCase						:       List[Any]  			=     torch.randn(
           1				,     model.config.in_channels				,     model.config.sample_size				,     model.config.sample_size				,     generator=torch.manual_seed(0  )				,     )
       _lowerCAmelCase						:       List[str]  			=     noise.to(a__  )
       _lowerCAmelCase						:       Any  			=     torch.tensor([10] * noise.shape[0]  ).to(a__  )

       with torch.no_grad():
         _lowerCAmelCase						:       Union[str, Any]  			=     model(a__				,     a__  ).sample

       _lowerCAmelCase						:       int  			=     output[0, -1, -3:, -3:].flatten().cpu()
       # fmt: off
       _lowerCAmelCase						:       Optional[int]  			=     torch.tensor([-1_3.3_2_5_8, -2_0.1_1_0_0, -1_5.9_8_7_3, -1_7.6_6_1_7, -2_3.0_5_9_6, -1_7.9_4_1_9, -1_3.3_6_7_5, -1_6.1_8_8_9, -1_2.3_8_0_0]  )
       # fmt: on

       self.assertTrue(torch_all_close(a__				,     a__				,     rtol=1e-3  )  )







class 						__A							( SCREAMING_SNAKE_CASE_    ,	SCREAMING_SNAKE_CASE_    ,	unittest.TestCase		):
     _UpperCamelCase							:			List[str]							=		UNetaDModel
     _UpperCamelCase							:			Any							=		"sample"






     @property
     def     __A  (		self				,     a__=(32, 32)  ):
       _lowerCAmelCase						:       Optional[int]  			=     4
       _lowerCAmelCase						:       Dict  			=     3

       _lowerCAmelCase						:       List[Any]  			=     floats_tensor((batch_size, num_channels) + sizes  ).to(a__  )
       _lowerCAmelCase						:       Any  			=     torch.tensor(batch_size * [10]  ).to(dtype=torch.intaa				,     device=a__  )

       return {"sample": noise, "timestep": time_step}






     @property
     def     __A  (		self  ):
       return (3, 32, 32)






     @property
     def     __A  (		self  ):
       return (3, 32, 32)






     def     __A  (		self  ):
       _lowerCAmelCase						:       List[Any]  			=     {
           """block_out_channels""": [32, 64, 64, 64],
           """in_channels""": 3,
           """layers_per_block""": 1,
           """out_channels""": 3,
           """time_embedding_type""": """fourier""",
           """norm_eps""": 1e-6,
           """mid_block_scale_factor""": math.sqrt(2.0  ),
           """norm_num_groups""": None,
           """down_block_types""": [
               """SkipDownBlock2D""",
               """AttnSkipDownBlock2D""",
               """SkipDownBlock2D""",
               """SkipDownBlock2D""",
           ],
           """up_block_types""": [
               """SkipUpBlock2D""",
               """SkipUpBlock2D""",
               """AttnSkipUpBlock2D""",
               """SkipUpBlock2D""",
           ],
       }
       _lowerCAmelCase						:       Optional[int]  			=     self.dummy_input
       return init_dict, inputs_dict






     @slow
     def     __A  (		self  ):
       _lowerCAmelCase			,	_lowerCAmelCase						:       Union[str, Any]  			=     UNetaDModel.from_pretrained("""google/ncsnpp-celebahq-256"""				,     output_loading_info=a__  )
       self.assertIsNotNone(a__  )
       self.assertEqual(len(loading_info["""missing_keys"""]  )				,     0  )

       model.to(a__  )
       _lowerCAmelCase						:       str  			=     self.dummy_input
       _lowerCAmelCase						:       Optional[Any]  			=     floats_tensor((4, 3) + (256, 256)  ).to(a__  )
       _lowerCAmelCase						:       Any  			=     noise
       _lowerCAmelCase						:       str  			=     model(**a__  )

       assert image is not None, "Make sure output is not None"






     @slow
     def     __A  (		self  ):
       _lowerCAmelCase						:       Optional[Any]  			=     UNetaDModel.from_pretrained("""google/ncsnpp-celebahq-256"""  )
       model.to(a__  )

       _lowerCAmelCase						:       str  			=     4
       _lowerCAmelCase						:       List[str]  			=     3
       _lowerCAmelCase						:       Any  			=     (256, 256)

       _lowerCAmelCase						:       List[Any]  			=     torch.ones((batch_size, num_channels) + sizes  ).to(a__  )
       _lowerCAmelCase						:       str  			=     torch.tensor(batch_size * [1e-4]  ).to(a__  )

       with torch.no_grad():
         _lowerCAmelCase						:       Optional[Any]  			=     model(a__				,     a__  ).sample

       _lowerCAmelCase						:       Optional[Any]  			=     output[0, -3:, -3:, -1].flatten().cpu()
       # fmt: off
       _lowerCAmelCase						:       int  			=     torch.tensor([-4_8_4_2.8_6_9_1, -6_4_9_9.6_6_3_1, -3_8_0_0.1_9_5_3, -7_9_7_8.2_6_8_6, -1_0_9_8_0.7_1_2_9, -2_0_0_2_8.8_5_3_5, 8_1_4_8.2_8_2_2, 2_3_4_2.2_9_0_5, 5_6_7.7_6_0_8]  )
       # fmt: on

       self.assertTrue(torch_all_close(a__				,     a__				,     rtol=1e-2  )  )






     def     __A  (		self  ):
       _lowerCAmelCase						:       int  			=     UNetaDModel.from_pretrained("""fusing/ncsnpp-ffhq-ve-dummy-update"""  )
       model.to(a__  )

       _lowerCAmelCase						:       Union[str, Any]  			=     4
       _lowerCAmelCase						:       Tuple  			=     3
       _lowerCAmelCase						:       int  			=     (32, 32)

       _lowerCAmelCase						:       List[str]  			=     torch.ones((batch_size, num_channels) + sizes  ).to(a__  )
       _lowerCAmelCase						:       List[Any]  			=     torch.tensor(batch_size * [1e-4]  ).to(a__  )

       with torch.no_grad():
         _lowerCAmelCase						:       int  			=     model(a__				,     a__  ).sample

       _lowerCAmelCase						:       Tuple  			=     output[0, -3:, -3:, -1].flatten().cpu()
       # fmt: off
       _lowerCAmelCase						:       Optional[Any]  			=     torch.tensor([-0.0_3_2_5, -0.0_9_0_0, -0.0_8_6_9, -0.0_3_3_2, -0.0_7_2_5, -0.0_2_7_0, -0.0_1_0_1, 0.0_2_2_7, 0.0_2_5_6]  )
       # fmt: on

       self.assertTrue(torch_all_close(a__				,     a__				,     rtol=1e-2  )  )






     def     __A  (		self  ):
       # not required for this model
       pass





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






   def __init__(      self				:     Tuple	,						lowercase_				:     List[Any]	,						lowercase_				:     Union[str, Any]=2	,						lowercase_				:     Union[str, Any]=56	,						lowercase_				:     Tuple=True	,						lowercase_				:     Optional[Any]=True	,						lowercase_				:     Optional[Any]=True	,						lowercase_				:     int=True	,						lowercase_				:     Any=99	,						lowercase_				:     int=32	,						lowercase_				:     str=2	,						lowercase_				:     Union[str, Any]=2	,						lowercase_				:     Dict=7	,						lowercase_				:     Dict="gelu_new"	,						lowercase_				:     Tuple=0.1	,						lowercase_				:     List[Any]=0.1	,						lowercase_				:     Tuple=5_12	,						lowercase_				:     Optional[Any]=16	,						lowercase_				:     List[Any]=2	,						lowercase_				:     Dict=0.02	,						lowercase_				:     int=4	,						lowercase_				:     Tuple="block_sparse"	,						lowercase_				:     Dict=True	,						lowercase_				:     Optional[int]=False	,						lowercase_				:     Dict=2	,						lowercase_				:     int=3	,						)					->						Union[str, Any]:
       lowercase__ :   Dict    					=				parent
       lowercase__ :   Dict    					=				batch_size
       lowercase__ :   Tuple    					=				seq_length
       lowercase__ :   Dict    					=				is_training
       lowercase__ :   Dict    					=				use_attention_mask
       lowercase__ :   Tuple    					=				use_token_type_ids
       lowercase__ :   Optional[int]    					=				use_labels
       lowercase__ :   List[Any]    					=				vocab_size
       lowercase__ :   Any    					=				hidden_size
       lowercase__ :   List[Any]    					=				num_hidden_layers
       lowercase__ :   Union[str, Any]    					=				num_attention_heads
       lowercase__ :   str    					=				intermediate_size
       lowercase__ :   int    					=				hidden_act
       lowercase__ :   str    					=				hidden_dropout_prob
       lowercase__ :   List[str]    					=				attention_probs_dropout_prob
       lowercase__ :   Optional[Any]    					=				max_position_embeddings
       lowercase__ :   Union[str, Any]    					=				type_vocab_size
       lowercase__ :   Dict    					=				type_sequence_label_size
       lowercase__ :   Any    					=				initializer_range
       lowercase__ :   List[str]    					=				num_choices

       lowercase__ :   str    					=				rescale_embeddings
       lowercase__ :   Optional[Any]    					=				attention_type
       lowercase__ :   Optional[int]    					=				use_bias
       lowercase__ :   Optional[int]    					=				block_size
       lowercase__ :   str    					=				num_random_blocks






   def       __UpperCamelCase     (      self				:     str  )					->						Optional[Any]:
       lowercase__ :   List[Any]    					=				ids_tensor([self.batch_size, self.seq_length]	,						self.vocab_size  )

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

       lowercase__ :   Optional[int]    					=				None
       if self.use_token_type_ids:
           lowercase__ :   Any    					=				ids_tensor([self.batch_size, self.seq_length]	,						self.type_vocab_size  )

       lowercase__ :   int    					=				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				:     Union[str, Any]  )					->						int:
       lowercase__ :   int    					=				self.prepare_config_and_inputs()
       lowercase__					,   lowercase__					,   lowercase__					,   lowercase__ :   Dict    					=				config_and_inputs
       lowercase__ :   Union[str, Any]    					=				{
           "input_ids": input_ids,
           "token_type_ids": token_type_ids,
           "attention_mask": attention_mask,
       }
       return config, inputs_dict



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

   __A	:		List[str]         =				False
   __A	:		Any         =				False






   def       __UpperCamelCase     (      self				:     List[str]  )					->						List[Any]:
       lowercase__ :   Union[str, Any]    					=				FlaxBigBirdModelTester(self  )






   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def       __UpperCamelCase     (      self				:     Optional[int]  )					->						Dict:
       super().test_from_pretrained_save_pretrained()






   @slow
   # copied from `test_modeling_flax_common` because it takes much longer than other models
   def       __UpperCamelCase     (      self				:     List[str]  )					->						Any:
       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				:     Tuple  )					->						str:
       super().test_no_automatic_init()






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






   @slow
   def       __UpperCamelCase     (      self				:     Optional[int]  )					->						Tuple:
       for model_class_name in self.all_model_classes:
           lowercase__ :   Optional[Any]    					=				model_class_name.from_pretrained("google/bigbird-roberta-base"  )
           self.assertIsNotNone(lowercase_  )






   def       __UpperCamelCase     (      self				:     int  )					->						Optional[int]:
       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				:     str  )					->						Any:
       lowercase__					,   lowercase__ :   Dict    					=				self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
           with self.subTest(model_class.__name__  ):
               lowercase__ :   Union[str, Any]    					=				self._prepare_for_class(lowercase_	,						lowercase_  )
               lowercase__ :   Optional[Any]    					=				model_class(lowercase_  )

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

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

               with self.subTest("JIT Disabled"  ):
                   with jax.disable_jit():
                       lowercase__ :   Any    					=				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				:     List[Any]	,						lowercase_				:     str	,						lowercase_				:     Union[str, Any]	,						lowercase_				:     Optional[int]	,						lowercase_				:     List[Any]=1E-5	,						lowercase_				:     Any="outputs"	,						lowercase_				:     List[str]=None  )					->						List[Any]:
       # `bigbird_block_sparse_attention` in `FlaxBigBird` returns `attention_probs = None`, while in PyTorch version,
       # an effort was done to return `attention_probs` (yet to be verified).
       if name.startswith("outputs.attentions"  ):
           return
       else:
           super().check_pt_flax_outputs(lowercase_	,						lowercase_	,						lowercase_	,						lowercase_	,						lowercase_	,						lowercase_  )



