





'''simple docstring'''






import unittest

from diffusers.models.unet_ad_blocks import *  # noqa F403
from diffusers.utils import torch_device

from .test_unet_blocks_common import UNetBlockTesterMixin







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				DownBlockaD  # noqa F405
  lowerCAmelCase_					=				"""down"""


  def   UpperCamelCase__      (				self    :					str  )     ->						List[str]:
    _UpperCamelCase       			=[-0.0232, -0.9869, 0.8054, -0.0637, -0.1688, -1.4264, 0.4470, -1.3394, 0.0904]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				ResnetDownsampleBlockaD  # noqa F405
  lowerCAmelCase_					=				"""down"""


  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						Optional[int]:
    _UpperCamelCase       			=[0.0710, 0.2410, -0.7320, -1.0757, -1.1343, 0.3540, -0.0133, -0.2576, 0.0948]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				AttnDownBlockaD  # noqa F405
  lowerCAmelCase_					=				"""down"""


  def   UpperCamelCase__      (				self    :					int  )     ->						str:
    _UpperCamelCase       			=[0.0636, 0.8964, -0.6234, -1.0131, 0.0844, 0.4935, 0.3437, 0.0911, -0.2957]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				CrossAttnDownBlockaD  # noqa F405
  lowerCAmelCase_					=				"""down"""


  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						int:
    _UpperCamelCase			,     _UpperCamelCase       			=super().prepare_init_args_and_inputs_for_common()
    _UpperCamelCase       			=32
    return init_dict, inputs_dict




  def   UpperCamelCase__      (				self    :					Any  )     ->						List[Any]:
    _UpperCamelCase       			=[0.2238, -0.7396, -0.2255, -0.3829, 0.1925, 1.1665, 0.0603, -0.7295, 0.1983]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				SimpleCrossAttnDownBlockaD  # noqa F405
  lowerCAmelCase_					=				"""down"""


  @property
  def   UpperCamelCase__      (				self    :					str  )     ->						int:
    return super().get_dummy_input(include_encoder_hidden_states=UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					str  )     ->						str:
    _UpperCamelCase			,     _UpperCamelCase       			=super().prepare_init_args_and_inputs_for_common()
    _UpperCamelCase       			=32
    return init_dict, inputs_dict




  @unittest.skipIf(torch_device == '''mps'''							,	'''MPS result is not consistent'''  )
  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						Union[str, Any]:
    _UpperCamelCase       			=[0.7921, -0.0992, -0.1962, -0.7695, -0.4242, 0.7804, 0.4737, 0.2765, 0.3338]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				SkipDownBlockaD  # noqa F405
  lowerCAmelCase_					=				"""down"""


  @property
  def   UpperCamelCase__      (				self    :					Tuple  )     ->						Any:
    return super().get_dummy_input(include_skip_sample=UpperCamelCase__  )




  def   UpperCamelCase__      (				self    :					Dict  )     ->						Optional[Any]:
    _UpperCamelCase       			=[-0.0845, -0.2087, -0.2465, 0.0971, 0.1900, -0.0484, 0.2664, 0.4179, 0.5069]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				AttnSkipDownBlockaD  # noqa F405
  lowerCAmelCase_					=				"""down"""


  @property
  def   UpperCamelCase__      (				self    :					int  )     ->						Union[str, Any]:
    return super().get_dummy_input(include_skip_sample=UpperCamelCase__  )




  def   UpperCamelCase__      (				self    :					Any  )     ->						str:
    _UpperCamelCase       			=[0.5539, 0.1609, 0.4924, 0.0537, -0.1995, 0.4050, 0.0979, -0.2721, -0.0642]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				DownEncoderBlockaD  # noqa F405
  lowerCAmelCase_					=				"""down"""


  @property
  def   UpperCamelCase__      (				self    :					Any  )     ->						Dict:
    return super().get_dummy_input(include_temb=UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					int  )     ->						str:
    _UpperCamelCase       			={
        '''in_channels''': 32,
        '''out_channels''': 32,
    }
    _UpperCamelCase       			=self.dummy_input
    return init_dict, inputs_dict




  def   UpperCamelCase__      (				self    :					Any  )     ->						List[str]:
    _UpperCamelCase       			=[1.1102, 0.5302, 0.4872, -0.0023, -0.8042, 0.0483, -0.3489, -0.5632, 0.7626]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				AttnDownEncoderBlockaD  # noqa F405
  lowerCAmelCase_					=				"""down"""


  @property
  def   UpperCamelCase__      (				self    :					Dict  )     ->						Optional[int]:
    return super().get_dummy_input(include_temb=UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						Optional[int]:
    _UpperCamelCase       			={
        '''in_channels''': 32,
        '''out_channels''': 32,
    }
    _UpperCamelCase       			=self.dummy_input
    return init_dict, inputs_dict




  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						int:
    _UpperCamelCase       			=[0.8966, -0.1486, 0.8568, 0.8141, -0.9046, -0.1342, -0.0972, -0.7417, 0.1538]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				UNetMidBlockaD  # noqa F405
  lowerCAmelCase_					=				"""mid"""


  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						Any:
    _UpperCamelCase       			={
        '''in_channels''': 32,
        '''temb_channels''': 128,
    }
    _UpperCamelCase       			=self.dummy_input
    return init_dict, inputs_dict




  def   UpperCamelCase__      (				self    :					List[str]  )     ->						Union[str, Any]:
    _UpperCamelCase       			=[-0.1062, 1.7248, 0.3494, 1.4569, -0.0910, -1.2421, -0.9984, 0.6736, 1.0028]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				UNetMidBlockaDCrossAttn  # noqa F405
  lowerCAmelCase_					=				"""mid"""


  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						Union[str, Any]:
    _UpperCamelCase			,     _UpperCamelCase       			=super().prepare_init_args_and_inputs_for_common()
    _UpperCamelCase       			=32
    return init_dict, inputs_dict




  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						Union[str, Any]:
    _UpperCamelCase       			=[0.0187, 2.4220, 0.4484, 1.1203, -0.6121, -1.5122, -0.8270, 0.7851, 1.8335]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				UNetMidBlockaDSimpleCrossAttn  # noqa F405
  lowerCAmelCase_					=				"""mid"""


  @property
  def   UpperCamelCase__      (				self    :					int  )     ->						int:
    return super().get_dummy_input(include_encoder_hidden_states=UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						Optional[int]:
    _UpperCamelCase			,     _UpperCamelCase       			=super().prepare_init_args_and_inputs_for_common()
    _UpperCamelCase       			=32
    return init_dict, inputs_dict




  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						Tuple:
    _UpperCamelCase       			=[0.7143, 1.9974, 0.5448, 1.3977, 0.1282, -1.1237, -1.4238, 0.5530, 0.8880]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				UpBlockaD  # noqa F405
  lowerCAmelCase_					=				"""up"""


  @property
  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						Any:
    return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase__  )




  def   UpperCamelCase__      (				self    :					Tuple  )     ->						Tuple:
    _UpperCamelCase       			=[-0.2041, -0.4165, -0.3022, 0.0041, -0.6628, -0.7053, 0.1928, -0.0325, 0.0523]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				ResnetUpsampleBlockaD  # noqa F405
  lowerCAmelCase_					=				"""up"""


  @property
  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						Dict:
    return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase__  )




  def   UpperCamelCase__      (				self    :					str  )     ->						Dict:
    _UpperCamelCase       			=[0.2287, 0.3549, -0.1346, 0.4797, -0.1715, -0.9649, 0.7305, -0.5864, -0.6244]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				CrossAttnUpBlockaD  # noqa F405
  lowerCAmelCase_					=				"""up"""


  @property
  def   UpperCamelCase__      (				self    :					Any  )     ->						int:
    return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					Any  )     ->						Any:
    _UpperCamelCase			,     _UpperCamelCase       			=super().prepare_init_args_and_inputs_for_common()
    _UpperCamelCase       			=32
    return init_dict, inputs_dict




  def   UpperCamelCase__      (				self    :					Tuple  )     ->						List[Any]:
    _UpperCamelCase       			=[-0.1403, -0.3515, -0.0420, -0.1425, 0.3167, 0.5094, -0.2181, 0.5931, 0.5582]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				SimpleCrossAttnUpBlockaD  # noqa F405
  lowerCAmelCase_					=				"""up"""


  @property
  def   UpperCamelCase__      (				self    :					List[str]  )     ->						Tuple:
    return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase__							,	include_encoder_hidden_states=UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						Optional[Any]:
    _UpperCamelCase			,     _UpperCamelCase       			=super().prepare_init_args_and_inputs_for_common()
    _UpperCamelCase       			=32
    return init_dict, inputs_dict




  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						Any:
    _UpperCamelCase       			=[0.2645, 0.1480, 0.0909, 0.8044, -0.9758, -0.9083, 0.0994, -1.1453, -0.7402]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				AttnUpBlockaD  # noqa F405
  lowerCAmelCase_					=				"""up"""


  @property
  def   UpperCamelCase__      (				self    :					Tuple  )     ->						Any:
    return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase__  )




  @unittest.skipIf(torch_device == '''mps'''							,	'''MPS result is not consistent'''  )
  def   UpperCamelCase__      (				self    :					str  )     ->						Tuple:
    _UpperCamelCase       			=[0.0979, 0.1326, 0.0021, 0.0659, 0.2249, 0.0059, 0.1132, 0.5952, 0.1033]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				SkipUpBlockaD  # noqa F405
  lowerCAmelCase_					=				"""up"""


  @property
  def   UpperCamelCase__      (				self    :					List[str]  )     ->						Tuple:
    return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase__  )




  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						Optional[int]:
    _UpperCamelCase       			=[-0.0893, -0.1234, -0.1506, -0.0332, 0.0123, -0.0211, 0.0566, 0.0143, 0.0362]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				AttnSkipUpBlockaD  # noqa F405
  lowerCAmelCase_					=				"""up"""


  @property
  def   UpperCamelCase__      (				self    :					List[str]  )     ->						Optional[Any]:
    return super().get_dummy_input(include_res_hidden_states_tuple=UpperCamelCase__  )




  def   UpperCamelCase__      (				self    :					Union[str, Any]  )     ->						int:
    _UpperCamelCase       			=[0.0361, 0.0617, 0.2787, -0.0350, 0.0342, 0.3421, -0.0843, 0.0913, 0.3015]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				UpDecoderBlockaD  # noqa F405
  lowerCAmelCase_					=				"""up"""


  @property
  def   UpperCamelCase__      (				self    :					Dict  )     ->						List[Any]:
    return super().get_dummy_input(include_temb=UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						Any:
    _UpperCamelCase       			={'''in_channels''': 32, '''out_channels''': 32}

    _UpperCamelCase       			=self.dummy_input
    return init_dict, inputs_dict




  def   UpperCamelCase__      (				self    :					List[str]  )     ->						str:
    _UpperCamelCase       			=[0.4404, 0.1998, -0.9886, -0.3320, -0.3128, -0.7034, -0.6955, -0.2338, -0.3137]
    super().test_output(UpperCamelCase__  )







class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				AttnUpDecoderBlockaD  # noqa F405
  lowerCAmelCase_					=				"""up"""


  @property
  def   UpperCamelCase__      (				self    :					List[Any]  )     ->						str:
    return super().get_dummy_input(include_temb=UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						Tuple:
    _UpperCamelCase       			={'''in_channels''': 32, '''out_channels''': 32}

    _UpperCamelCase       			=self.dummy_input
    return init_dict, inputs_dict




  def   UpperCamelCase__      (				self    :					List[str]  )     ->						List[str]:
    _UpperCamelCase       			=[0.6738, 0.4491, 0.1055, 1.0710, 0.7316, 0.3339, 0.3352, 0.1023, 0.3568]
    super().test_output(UpperCamelCase__  )













'''simple docstring'''






from __future__ import absolute_import, division, print_function, unicode_literals

from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers import RobertaConfig
from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.roberta.modeling_roberta import (
    ROBERTA_INPUTS_DOCSTRING,
    ROBERTA_START_DOCSTRING,
    RobertaEmbeddings,
)

from .modeling_highway_bert import BertPreTrainedModel, DeeBertModel, HighwayException, entropy







@add_start_docstrings(
    """The RoBERTa Model transformer with early exiting (DeeRoBERTa). """   ,     lowercase_   ,     )
class 		UpperCAmelCase   (				lowercase_):


  """simple docstring"""







  lowerCAmelCase_					=				RobertaConfig
  lowerCAmelCase_					=				"""roberta"""


  def __init__(				self    :					Dict							,	UpperCamelCase__    :					List[str]  )     ->						List[Any]:
    super().__init__(UpperCamelCase__  )

    _UpperCamelCase       			=RobertaEmbeddings(UpperCamelCase__  )
    self.init_weights()







@add_start_docstrings(
    """RoBERTa Model (with early exiting - DeeRoBERTa) with a classifier on top,
    also takes care of multi-layer training. """   ,     lowercase_   ,     )
class 		UpperCAmelCase   (				lowercase_):


  """simple docstring"""







  lowerCAmelCase_					=				RobertaConfig
  lowerCAmelCase_					=				"""roberta"""


  def __init__(				self    :					Tuple							,	UpperCamelCase__    :					Any  )     ->						Dict:
    super().__init__(UpperCamelCase__  )
    _UpperCamelCase       			=config.num_labels
    _UpperCamelCase       			=config.num_hidden_layers

    _UpperCamelCase       			=DeeRobertaModel(UpperCamelCase__  )
    _UpperCamelCase       			=nn.Dropout(config.hidden_dropout_prob  )
    _UpperCamelCase       			=nn.Linear(config.hidden_size							,	self.config.num_labels  )




  @add_start_docstrings_to_model_forward(UpperCamelCase__  )
  def   UpperCamelCase__      (				self    :					List[Any]							,	UpperCamelCase__    :					Optional[int]=None							,	UpperCamelCase__    :					Union[str, Any]=None							,	UpperCamelCase__    :					List[str]=None							,	UpperCamelCase__    :					Optional[Any]=None							,	UpperCamelCase__    :					int=None							,	UpperCamelCase__    :					Union[str, Any]=None							,	UpperCamelCase__    :					Tuple=None							,	UpperCamelCase__    :					Union[str, Any]=-1							,	UpperCamelCase__    :					Dict=False							,	)     ->						Union[str, Any]:
    _UpperCamelCase       			=self.num_layers
    try:
      _UpperCamelCase       			=self.roberta(
          UpperCamelCase__							,	attention_mask=UpperCamelCase__							,	token_type_ids=UpperCamelCase__							,	position_ids=UpperCamelCase__							,	head_mask=UpperCamelCase__							,	inputs_embeds=UpperCamelCase__							,	)

      _UpperCamelCase       			=outputs[1]

      _UpperCamelCase       			=self.dropout(UpperCamelCase__  )
      _UpperCamelCase       			=self.classifier(UpperCamelCase__  )
      _UpperCamelCase       			=(logits,) + outputs[2:]  # add hidden states and attention if they are here
    except HighwayException as e:
      _UpperCamelCase       			=e.message
      _UpperCamelCase       			=e.exit_layer
      _UpperCamelCase       			=outputs[0]

    if not self.training:
      _UpperCamelCase       			=entropy(UpperCamelCase__  )
      _UpperCamelCase       			=[]
      _UpperCamelCase       			=[]
    if labels is not None:
      if self.num_labels == 1:
        #  We are doing regression
        _UpperCamelCase       			=MSELoss()
        _UpperCamelCase       			=loss_fct(logits.view(-1  )							,	labels.view(-1  )  )
      else:
        _UpperCamelCase       			=CrossEntropyLoss()
        _UpperCamelCase       			=loss_fct(logits.view(-1							,	self.num_labels  )							,	labels.view(-1  )  )

      # work with highway exits
      _UpperCamelCase       			=[]
      for highway_exit in outputs[-1]:
        _UpperCamelCase       			=highway_exit[0]
        if not self.training:
          highway_logits_all.append(UpperCamelCase__  )
          highway_entropy.append(highway_exit[2]  )
        if self.num_labels == 1:
          #  We are doing regression
          _UpperCamelCase       			=MSELoss()
          _UpperCamelCase       			=loss_fct(highway_logits.view(-1  )							,	labels.view(-1  )  )
        else:
          _UpperCamelCase       			=CrossEntropyLoss()
          _UpperCamelCase       			=loss_fct(highway_logits.view(-1							,	self.num_labels  )							,	labels.view(-1  )  )
        highway_losses.append(UpperCamelCase__  )

      if train_highway:
        _UpperCamelCase       			=(sum(highway_losses[:-1]  ),) + outputs
        # exclude the final highway, of course
      else:
        _UpperCamelCase       			=(loss,) + outputs
    if not self.training:
      _UpperCamelCase       			=outputs + ((original_entropy, highway_entropy), exit_layer)
      if output_layer >= 0:
        _UpperCamelCase       			=(
            (outputs[0],) + (highway_logits_all[output_layer],) + outputs[2:]
        )  # use the highway of the last layer

    return outputs  # (loss), logits, (hidden_states), (attentions), entropy






