
"""simple docstring"""

import unittest

from transformers import FunnelConfig, FunnelTokenizer, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import (
          MODEL_FOR_PRETRAINING_MAPPING,
          FunnelBaseModel,
          FunnelForMaskedLM,
          FunnelForMultipleChoice,
          FunnelForPreTraining,
          FunnelForQuestionAnswering,
          FunnelForSequenceClassification,
          FunnelForTokenClassification,
          FunnelModel,
      )

class __lowerCAmelCase						:





     def __init__(      self      ,	__UpperCAmelCase      ,	__UpperCAmelCase=13      ,	__UpperCAmelCase=7      ,	__UpperCAmelCase=True      ,	__UpperCAmelCase=True      ,	__UpperCAmelCase=True      ,	__UpperCAmelCase=True      ,	__UpperCAmelCase=99      ,	__UpperCAmelCase=[1, 1, 2]      ,	__UpperCAmelCase=1      ,	__UpperCAmelCase=32      ,	__UpperCAmelCase=4      ,	__UpperCAmelCase=8      ,	__UpperCAmelCase=37      ,	__UpperCAmelCase="gelu_new"      ,	__UpperCAmelCase=0.1      ,	__UpperCAmelCase=0.1      ,	__UpperCAmelCase=0.0      ,	__UpperCAmelCase=512      ,	__UpperCAmelCase=3      ,	__UpperCAmelCase=0.0_2      ,	__UpperCAmelCase=3      ,	__UpperCAmelCase=4      ,	__UpperCAmelCase=None      ,	__UpperCAmelCase=False      ,	):

      '''simple docstring'''




      __UpperCamelCase							    =			parent
      __UpperCamelCase							    =			batch_size
      __UpperCamelCase							    =			seq_length
      __UpperCamelCase							    =			is_training
      __UpperCamelCase							    =			use_input_mask
      __UpperCamelCase							    =			use_token_type_ids
      __UpperCamelCase							    =			use_labels
      __UpperCamelCase							    =			vocab_size
      __UpperCamelCase							    =			block_sizes
      __UpperCamelCase							    =			num_decoder_layers
      __UpperCamelCase							    =			d_model
      __UpperCamelCase							    =			n_head
      __UpperCamelCase							    =			d_head
      __UpperCamelCase							    =			d_inner
      __UpperCamelCase							    =			hidden_act
      __UpperCamelCase							    =			hidden_dropout
      __UpperCamelCase							    =			attention_dropout
      __UpperCamelCase							    =			activation_dropout
      __UpperCamelCase							    =			max_position_embeddings
      __UpperCamelCase							    =			type_vocab_size
      __UpperCamelCase							    =			2
      __UpperCamelCase							    =			num_labels
      __UpperCamelCase							    =			num_choices
      __UpperCamelCase							    =			scope
      __UpperCamelCase							    =			initializer_std

      # Used in the tests to check the size of the first attention layer
      __UpperCamelCase							    =			n_head
      # Used in the tests to check the size of the first hidden state
      __UpperCamelCase							    =			self.d_model
      # Used in the tests to check the number of output hidden states/attentions
      __UpperCamelCase							    =			sum(self.block_sizes ) + (0 if base else self.num_decoder_layers)
      # FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
      # the last hidden state of the first block (which is the first hidden state of the decoder).
      if not base:
       __UpperCamelCase							    =			self.num_hidden_layers + 2





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      __UpperCamelCase							    =			ids_tensor([self.batch_size, self.seq_length]      ,	self.vocab_size )

      __UpperCamelCase							    =			None
      if self.use_input_mask:
       __UpperCamelCase							    =			ids_tensor([self.batch_size, self.seq_length]      ,	vocab_size=2 )

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

      __UpperCamelCase							    =			None
      __UpperCamelCase							    =			None
      __UpperCamelCase							    =			None
      if self.use_labels:
       __UpperCamelCase							    =			ids_tensor([self.batch_size]      ,	self.type_sequence_label_size )
       __UpperCamelCase							    =			ids_tensor([self.batch_size, self.seq_length]      ,	self.num_labels )
       __UpperCamelCase							    =			ids_tensor([self.batch_size]      ,	self.num_choices )
       __UpperCamelCase							    =			ids_tensor([self.batch_size, self.seq_length]      ,	1 )

      __UpperCamelCase							    =			self.get_config()

      return (
          config,
          input_ids,
          token_type_ids,
          input_mask,
          sequence_labels,
          token_labels,
          choice_labels,
          fake_token_labels,
      )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      return FunnelConfig(
          vocab_size=self.vocab_size      ,	block_sizes=self.block_sizes      ,	num_decoder_layers=self.num_decoder_layers      ,	d_model=self.d_model      ,	n_head=self.n_head      ,	d_head=self.d_head      ,	d_inner=self.d_inner      ,	hidden_act=self.hidden_act      ,	hidden_dropout=self.hidden_dropout      ,	attention_dropout=self.attention_dropout      ,	activation_dropout=self.activation_dropout      ,	max_position_embeddings=self.max_position_embeddings      ,	type_vocab_size=self.type_vocab_size      ,	initializer_std=self.initializer_std      ,	)





     def 		UpperCAmelCase   (      self      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	):

      '''simple docstring'''




      __UpperCamelCase							    =			FunnelModel(config=__UpperCAmelCase )
      model.to(__UpperCAmelCase )
      model.eval()
      __UpperCamelCase							    =			model(__UpperCAmelCase      ,	attention_mask=__UpperCAmelCase      ,	token_type_ids=__UpperCAmelCase )
      __UpperCamelCase							    =			model(__UpperCAmelCase      ,	token_type_ids=__UpperCAmelCase )
      __UpperCamelCase							    =			model(__UpperCAmelCase )
      self.parent.assertEqual(result.last_hidden_state.shape      ,	(self.batch_size, self.seq_length, self.d_model) )

      __UpperCamelCase							    =			False
      __UpperCamelCase							    =			model(__UpperCAmelCase )
      self.parent.assertEqual(result.last_hidden_state.shape      ,	(self.batch_size, self.seq_length, self.d_model) )

      __UpperCamelCase							    =			False
      __UpperCamelCase							    =			model(__UpperCAmelCase )
      self.parent.assertEqual(result.last_hidden_state.shape      ,	(self.batch_size, self.seq_length, self.d_model) )





     def 		UpperCAmelCase   (      self      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	):

      '''simple docstring'''




      __UpperCamelCase							    =			FunnelBaseModel(config=__UpperCAmelCase )
      model.to(__UpperCAmelCase )
      model.eval()
      __UpperCamelCase							    =			model(__UpperCAmelCase      ,	attention_mask=__UpperCAmelCase      ,	token_type_ids=__UpperCAmelCase )
      __UpperCamelCase							    =			model(__UpperCAmelCase      ,	token_type_ids=__UpperCAmelCase )
      __UpperCamelCase							    =			model(__UpperCAmelCase )
      self.parent.assertEqual(result.last_hidden_state.shape      ,	(self.batch_size, 2, self.d_model) )

      __UpperCamelCase							    =			False
      __UpperCamelCase							    =			model(__UpperCAmelCase )
      self.parent.assertEqual(result.last_hidden_state.shape      ,	(self.batch_size, 3, self.d_model) )

      __UpperCamelCase							    =			False
      __UpperCamelCase							    =			model(__UpperCAmelCase )
      self.parent.assertEqual(result.last_hidden_state.shape      ,	(self.batch_size, 2, self.d_model) )





     def 		UpperCAmelCase   (      self      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	):

      '''simple docstring'''




      __UpperCamelCase							    =			self.num_labels
      __UpperCamelCase							    =			FunnelForPreTraining(config=__UpperCAmelCase )
      model.to(__UpperCAmelCase )
      model.eval()
      __UpperCamelCase							    =			model(__UpperCAmelCase      ,	attention_mask=__UpperCAmelCase      ,	token_type_ids=__UpperCAmelCase      ,	labels=__UpperCAmelCase )
      self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.seq_length) )





     def 		UpperCAmelCase   (      self      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	):

      '''simple docstring'''




      __UpperCamelCase							    =			FunnelForMaskedLM(config=__UpperCAmelCase )
      model.to(__UpperCAmelCase )
      model.eval()
      __UpperCamelCase							    =			model(__UpperCAmelCase      ,	attention_mask=__UpperCAmelCase      ,	token_type_ids=__UpperCAmelCase      ,	labels=__UpperCAmelCase )
      self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.seq_length, self.vocab_size) )





     def 		UpperCAmelCase   (      self      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	):

      '''simple docstring'''




      __UpperCamelCase							    =			self.num_labels
      __UpperCamelCase							    =			FunnelForSequenceClassification(__UpperCAmelCase )
      model.to(__UpperCAmelCase )
      model.eval()
      __UpperCamelCase							    =			model(__UpperCAmelCase      ,	attention_mask=__UpperCAmelCase      ,	token_type_ids=__UpperCAmelCase      ,	labels=__UpperCAmelCase )
      self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.num_labels) )





     def 		UpperCAmelCase   (      self      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	):

      '''simple docstring'''




      __UpperCamelCase							    =			self.num_choices
      __UpperCamelCase							    =			FunnelForMultipleChoice(config=__UpperCAmelCase )
      model.to(__UpperCAmelCase )
      model.eval()
      __UpperCamelCase							    =			input_ids.unsqueeze(1 ).expand(-1      ,	self.num_choices      ,	-1 ).contiguous()
      __UpperCamelCase							    =			token_type_ids.unsqueeze(1 ).expand(-1      ,	self.num_choices      ,	-1 ).contiguous()
      __UpperCamelCase							    =			input_mask.unsqueeze(1 ).expand(-1      ,	self.num_choices      ,	-1 ).contiguous()
      __UpperCamelCase							    =			model(
          __UpperCAmelCase      ,	attention_mask=__UpperCAmelCase      ,	token_type_ids=__UpperCAmelCase      ,	labels=__UpperCAmelCase      ,	)
      self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.num_choices) )





     def 		UpperCAmelCase   (      self      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	):

      '''simple docstring'''




      __UpperCamelCase							    =			self.num_labels
      __UpperCamelCase							    =			FunnelForTokenClassification(config=__UpperCAmelCase )
      model.to(__UpperCAmelCase )
      model.eval()
      __UpperCamelCase							    =			model(__UpperCAmelCase      ,	attention_mask=__UpperCAmelCase      ,	token_type_ids=__UpperCAmelCase      ,	labels=__UpperCAmelCase )
      self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.seq_length, self.num_labels) )





     def 		UpperCAmelCase   (      self      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	):

      '''simple docstring'''




      __UpperCamelCase							    =			FunnelForQuestionAnswering(config=__UpperCAmelCase )
      model.to(__UpperCAmelCase )
      model.eval()
      __UpperCamelCase							    =			model(
          __UpperCAmelCase      ,	attention_mask=__UpperCAmelCase      ,	token_type_ids=__UpperCAmelCase      ,	start_positions=__UpperCAmelCase      ,	end_positions=__UpperCAmelCase      ,	)
      self.parent.assertEqual(result.start_logits.shape      ,	(self.batch_size, self.seq_length) )
      self.parent.assertEqual(result.end_logits.shape      ,	(self.batch_size, self.seq_length) )







     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      __UpperCamelCase							    =			self.prepare_config_and_inputs()
      (
          (
          __UpperCamelCase
      )    ,      (
          __UpperCamelCase
      )    ,      (
          __UpperCamelCase
      )    ,      (
          __UpperCamelCase
      )    ,      (
          __UpperCamelCase
      )    ,      (
          __UpperCamelCase
      )    ,      (
          __UpperCamelCase
      )    ,      (
          __UpperCamelCase
      )    ,      
      )							    =			config_and_inputs
      __UpperCamelCase							    =			{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
      return config, inputs_dict


@require_torch
class __lowerCAmelCase						( __SCREAMING_SNAKE_CASE     ,    __SCREAMING_SNAKE_CASE     ,    unittest.TestCase     ):
     lowercase		  =		False
     lowercase		  =		False
     lowercase		  =		(
         (
             FunnelModel,
             FunnelForMaskedLM,
             FunnelForPreTraining,
             FunnelForQuestionAnswering,
             FunnelForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     lowercase		  =		(
         {
             "feature-extraction": (FunnelBaseModel, FunnelModel),
             "fill-mask": FunnelForMaskedLM,
             "question-answering": FunnelForQuestionAnswering,
             "text-classification": FunnelForSequenceClassification,
             "token-classification": FunnelForTokenClassification,
             "zero-shot": FunnelForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )





     def 		UpperCAmelCase   (      self      ,	__UpperCAmelCase      ,	__UpperCAmelCase      ,	__UpperCAmelCase=False ):

      '''simple docstring'''




      __UpperCamelCase							    =			super()._prepare_for_class(__UpperCAmelCase      ,	__UpperCAmelCase      ,	return_labels=__UpperCAmelCase )

      if return_labels:
       if model_class in get_values(__UpperCAmelCase ):
        __UpperCamelCase							    =			torch.zeros(
            (self.model_tester.batch_size, self.model_tester.seq_length)      ,	dtype=torch.long      ,	device=__UpperCAmelCase )
      return inputs_dict





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      __UpperCamelCase							    =			FunnelModelTester(self )
      __UpperCamelCase							    =			ConfigTester(self      ,	config_class=__UpperCAmelCase )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      self.config_tester.run_common_tests()





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      __UpperCamelCase							    =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__UpperCAmelCase )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      __UpperCamelCase							    =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_pretraining(*__UpperCAmelCase )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      __UpperCamelCase							    =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*__UpperCAmelCase )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      __UpperCamelCase							    =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*__UpperCAmelCase )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      __UpperCamelCase							    =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_question_answering(*__UpperCAmelCase )







     def 		UpperCAmelCase   (      self      ,	__UpperCAmelCase ):

      '''simple docstring'''




      if hasattr(__UpperCAmelCase      ,	'weight' ) and module.weight is not None:
       module.weight.data.fill_(3 )
      if hasattr(__UpperCAmelCase      ,	'bias' ) and module.bias is not None:
       module.bias.data.fill_(3 )

      for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
       if hasattr(__UpperCAmelCase      ,	__UpperCAmelCase ) and getattr(__UpperCAmelCase      ,	__UpperCAmelCase ) is not None:
        __UpperCamelCase							    =			getattr(__UpperCAmelCase      ,	__UpperCAmelCase )
        weight.data.fill_(3 )


@require_torch
class __lowerCAmelCase						( __SCREAMING_SNAKE_CASE     ,    unittest.TestCase     ):
     lowercase		  =		False
     lowercase		  =		False
     lowercase		  =		(
         (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else ()
     )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      __UpperCamelCase							    =			FunnelModelTester(self      ,	base=__UpperCAmelCase )
      __UpperCamelCase							    =			ConfigTester(self      ,	config_class=__UpperCAmelCase )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      self.config_tester.run_common_tests()





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      __UpperCamelCase							    =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_base_model(*__UpperCAmelCase )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      __UpperCamelCase							    =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_sequence_classification(*__UpperCAmelCase )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      __UpperCamelCase							    =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_multiple_choice(*__UpperCAmelCase )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      __UpperCamelCase    ,      __UpperCamelCase							    =			self.model_tester.prepare_config_and_inputs_for_common()
      __UpperCamelCase							    =			True

      for model_class in self.all_model_classes:
       if model_class.__name__ == "FunnelBaseModel":
        continue
       __UpperCamelCase							    =			model_class(__UpperCAmelCase )
       model.to(__UpperCAmelCase )
       model.train()
       __UpperCamelCase							    =			self._prepare_for_class(__UpperCAmelCase      ,	__UpperCAmelCase      ,	return_labels=__UpperCAmelCase )
       __UpperCamelCase							    =			model(**__UpperCAmelCase ).loss
       loss.backward()







     def 		UpperCAmelCase   (      self      ,	__UpperCAmelCase ):

      '''simple docstring'''




      if hasattr(__UpperCAmelCase      ,	'weight' ) and module.weight is not None:
       module.weight.data.fill_(3 )
      if hasattr(__UpperCAmelCase      ,	'bias' ) and module.bias is not None:
       module.bias.data.fill_(3 )

      for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
       if hasattr(__UpperCAmelCase      ,	__UpperCAmelCase ) and getattr(__UpperCAmelCase      ,	__UpperCAmelCase ) is not None:
        __UpperCamelCase							    =			getattr(__UpperCAmelCase      ,	__UpperCAmelCase )
        weight.data.fill_(3 )


@require_torch
@require_sentencepiece
@require_tokenizers
class __lowerCAmelCase						( unittest.TestCase     ):





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      __UpperCamelCase							    =			13
      __UpperCamelCase							    =			7
      __UpperCamelCase							    =			torch.arange(0      ,	batch_size * sequence_length ).long().reshape(__UpperCAmelCase      ,	__UpperCAmelCase )
      __UpperCamelCase							    =			[0, 1, 2, 3, 4, 5, 6, 4, 1, 3, 5, 0, 1]
      __UpperCamelCase							    =			torch.tensor([[2] + [0] * a + [1] * (sequence_length - a - 1) for a in lengths] )

      __UpperCamelCase							    =			FunnelModel.from_pretrained('sgugger/funnel-random-tiny' )
      __UpperCamelCase							    =			model(__UpperCAmelCase      ,	token_type_ids=__UpperCAmelCase )[0].abs()

      __UpperCamelCase							    =			torch.tensor(2_3_4_4.8_3_5_2 )
      __UpperCamelCase							    =			torch.tensor(0.8_0_5_2 )
      self.assertTrue(torch.allclose(output.sum()      ,	__UpperCAmelCase      ,	atol=1E-4 ) )
      self.assertTrue(torch.allclose(output.mean()      ,	__UpperCAmelCase      ,	atol=1E-4 ) )

      __UpperCamelCase							    =			torch.tensor([[1] * 7, [1] * 4 + [0] * 3] * 6 + [[0, 1, 1, 0, 0, 1, 1]] )
      __UpperCamelCase							    =			model(__UpperCAmelCase      ,	attention_mask=__UpperCAmelCase      ,	token_type_ids=__UpperCAmelCase )[0].abs()

      __UpperCamelCase							    =			torch.tensor(2_3_4_3.8_4_2_5 )
      __UpperCamelCase							    =			torch.tensor(0.8_0_4_9 )
      self.assertTrue(torch.allclose(output.sum()      ,	__UpperCAmelCase      ,	atol=1E-4 ) )
      self.assertTrue(torch.allclose(output.mean()      ,	__UpperCAmelCase      ,	atol=1E-4 ) )







     @slow
     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      __UpperCamelCase							    =			FunnelTokenizer.from_pretrained('huggingface/funnel-small' )
      __UpperCamelCase							    =			FunnelModel.from_pretrained('huggingface/funnel-small' )
      __UpperCamelCase							    =			tokenizer('Hello! I am the Funnel Transformer model.'      ,	return_tensors='pt' )
      __UpperCamelCase							    =			model(**__UpperCAmelCase )[0]

      __UpperCamelCase							    =			torch.tensor(2_3_5.7_2_4_6 )
      __UpperCamelCase							    =			torch.tensor(0.0_2_5_6 )
      self.assertTrue(torch.allclose(output.sum()      ,	__UpperCAmelCase      ,	atol=1E-4 ) )
      self.assertTrue(torch.allclose(output.mean()      ,	__UpperCAmelCase      ,	atol=1E-4 ) )





"""simple docstring"""

import random
import unittest

import torch

from diffusers import IFSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin

@skip_mps
class __lowerCAmelCase						( __SCREAMING_SNAKE_CASE     ,    __SCREAMING_SNAKE_CASE     ,    unittest.TestCase     ):
     lowercase		  =		IFSuperResolutionPipeline
     lowercase		  =		TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"width", "height"}
     lowercase		  =		TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
     lowercase		  =		PipelineTesterMixin.required_optional_params - {"latents"}





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      return self._get_superresolution_dummy_components()





     def 		UpperCAmelCase   (      self      ,	__UpperCAmelCase      ,	__UpperCAmelCase=0 ):

      '''simple docstring'''




      if str(__UpperCAmelCase ).startswith('mps' ):
       __UpperCamelCase							    =			torch.manual_seed(__UpperCAmelCase )
      else:
       __UpperCamelCase							    =			torch.Generator(device=__UpperCAmelCase ).manual_seed(__UpperCAmelCase )

      __UpperCamelCase							    =			floats_tensor((1, 3, 32, 32)      ,	rng=random.Random(__UpperCAmelCase ) ).to(__UpperCAmelCase )

      __UpperCamelCase							    =			{
          'prompt': 'A painting of a squirrel eating a burger',
          'image': image,
          'generator': generator,
          'num_inference_steps': 2,
          'output_type': 'numpy',
      }

      return inputs





     @unittest.skipIf(
         torch_device != 'cuda' or not is_xformers_available()      ,	reason='XFormers attention is only available with CUDA and `xformers` installed'      ,	)
     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3 )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      self._test_save_load_optional_components()





     @unittest.skipIf(torch_device != 'cuda'      ,	reason='float16 requires CUDA' )
     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      super().test_save_load_floataa(expected_max_diff=1E-1 )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      self._test_attention_slicing_forward_pass(expected_max_diff=1E-2 )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      self._test_save_load_local()







     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      self._test_inference_batch_single_identical(
          expected_max_diff=1E-2      ,	)



