

import unittest
from typing import Tuple

import torch

from diffusers.utils import floats_tensor, randn_tensor, torch_all_close, torch_device
from diffusers.utils.testing_utils import require_torch





@require_torch
class A__    :



 """simple docstring"""







 @property
 def    __magic_name__       (		self :	str       ):






      '''simple docstring'''




      return self.get_dummy_input()

 @property
 def    __magic_name__       (		self :	int       ):






      '''simple docstring'''




      if self.block_type == "down":
           return (4, 3_2, 1_6, 1_6)
      elif self.block_type == "mid":
           return (4, 3_2, 3_2, 3_2)
      elif self.block_type == "up":
           return (4, 3_2, 6_4, 6_4)

      raise ValueError(F'''\'{self.block_type}\' is not a supported block_type. Set it to \'up\', \'mid\', or \'down\'.'''       )

 def    __magic_name__       (		self :	Union[str, Any]     ,						A_ :	str=True     ,						A_ :	List[str]=False     ,						A_ :	Any=False     ,						A_ :	Optional[int]=False     ,						):






      '''simple docstring'''




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

      _lowerCAmelCase   :					Optional[Any]        =    torch.manual_seed(0       )
      _lowerCAmelCase   :					Any        =    torch.device(A_       )
      _lowerCAmelCase   :					int        =    (batch_size, num_channels) + sizes
      _lowerCAmelCase   :					str        =    randn_tensor(A_     ,						generator=A_     ,						device=A_       )
      _lowerCAmelCase   :					List[str]        =    {"hidden_states": hidden_states}

      if include_temb:
           _lowerCAmelCase   :					Optional[Any]        =    1_2_8
           _lowerCAmelCase   :					Union[str, Any]        =    randn_tensor((batch_size, temb_channels)     ,						generator=A_     ,						device=A_       )

      if include_res_hidden_states_tuple:
           _lowerCAmelCase   :					List[str]        =    torch.manual_seed(1       )
           _lowerCAmelCase   :					str        =    (randn_tensor(A_     ,						generator=A_     ,						device=A_       ),)

      if include_encoder_hidden_states:
           _lowerCAmelCase   :					str        =    floats_tensor((batch_size, 3_2, 3_2)       ).to(A_       )

      if include_skip_sample:
           _lowerCAmelCase   :					Any        =    randn_tensor(((batch_size, 3) + sizes)     ,						generator=A_     ,						device=A_       )

      return dummy_input

 def    __magic_name__       (		self :	Any       ):






      '''simple docstring'''




      _lowerCAmelCase   :					List[str]        =    {
          "in_channels": 3_2,
          "out_channels": 3_2,
          "temb_channels": 1_2_8,
      }
      if self.block_type == "up":
           _lowerCAmelCase   :					Optional[Any]        =    3_2

      if self.block_type == "mid":
           init_dict.pop("out_channels"       )

      _lowerCAmelCase   :					Optional[Any]        =    self.dummy_input
      return init_dict, inputs_dict

 def    __magic_name__       (		self :	Optional[int]     ,						A_ :	Union[str, Any]       ):






      '''simple docstring'''




      _lowerCAmelCase      ,			_lowerCAmelCase   :					List[Any]        =    self.prepare_init_args_and_inputs_for_common()
      _lowerCAmelCase   :					Any        =    self.block_class(**A_       )
      unet_block.to(A_       )
      unet_block.eval()

      with torch.no_grad():
           _lowerCAmelCase   :					Optional[Any]        =    unet_block(**A_       )

      if isinstance(A_     ,						A_       ):
           _lowerCAmelCase   :					List[str]        =    output[0]

      self.assertEqual(output.shape     ,						self.output_shape       )

      _lowerCAmelCase   :					int        =    output[0, -1, -3:, -3:]
      _lowerCAmelCase   :					Union[str, Any]        =    torch.tensor(A_       ).to(A_       )
      assert torch_all_close(output_slice.flatten()     ,						A_     ,						atol=5E-3       )


 @unittest.skipIf(torch_device == "mps"     ,						"Training is not supported in mps"       )
 def    __magic_name__       (		self :	Any       ):






      '''simple docstring'''




      _lowerCAmelCase      ,			_lowerCAmelCase   :					Optional[int]        =    self.prepare_init_args_and_inputs_for_common()
      _lowerCAmelCase   :					Union[str, Any]        =    self.block_class(**A_       )
      model.to(A_       )
      model.train()
      _lowerCAmelCase   :					Optional[int]        =    model(**A_       )

      if isinstance(A_     ,						A_       ):
           _lowerCAmelCase   :					Optional[int]        =    output[0]

      _lowerCAmelCase   :					Optional[int]        =    torch.device(A_       )
      _lowerCAmelCase   :					Dict        =    randn_tensor(output.shape     ,						device=A_       )
      _lowerCAmelCase   :					Dict        =    torch.nn.functional.mse_loss(A_     ,						A_       )
      loss.backward()









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tensorflow_text_available, is_torch_available


__UpperCAmelCase    				=     {
    'configuration_ernie': ['ERNIE_PRETRAINED_CONFIG_ARCHIVE_MAP', 'ErnieConfig', 'ErnieOnnxConfig'],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase    				=     [
      'ERNIE_PRETRAINED_MODEL_ARCHIVE_LIST',
      'ErnieForCausalLM',
      'ErnieForMaskedLM',
      'ErnieForMultipleChoice',
      'ErnieForNextSentencePrediction',
      'ErnieForPreTraining',
      'ErnieForQuestionAnswering',
      'ErnieForSequenceClassification',
      'ErnieForTokenClassification',
      'ErnieModel',
      'ErniePreTrainedModel',
  ]

if TYPE_CHECKING:
  from .configuration_ernie import ERNIE_PRETRAINED_CONFIG_ARCHIVE_MAP, ErnieConfig, ErnieOnnxConfig

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_ernie import (
        ERNIE_PRETRAINED_MODEL_ARCHIVE_LIST,
        ErnieForCausalLM,
        ErnieForMaskedLM,
        ErnieForMultipleChoice,
        ErnieForNextSentencePrediction,
        ErnieForPreTraining,
        ErnieForQuestionAnswering,
        ErnieForSequenceClassification,
        ErnieForTokenClassification,
        ErnieModel,
        ErniePreTrainedModel,
    )

else:
  import sys

  __UpperCAmelCase    				=     _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)






