

"""simple docstring"""
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, slow, torch_device

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin

class        _UpperCamelCase	(						lowerCamelCase__   ,				lowerCamelCase__   ,				unittest.TestCase	):


  """simple docstring"""





  snake_case_        =    UNetaDModel
  snake_case_        =    'sample'



  @property
  def 							_UpperCAmelCase      ( self					: Any      )	->			Optional[int]:


         '''simple docstring'''





         __magic_name__			:		List[Any]        =     4
         __magic_name__			:		Optional[int]        =     14
         __magic_name__			:		Tuple        =     16

         __magic_name__			:		Dict        =     floats_tensor((batch_size, num_features, seq_len)      ).to(snake_case      )
         __magic_name__			:		List[Any]        =     torch.tensor([10] * batch_size      ).to(snake_case      )

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



  @property
  def 							_UpperCAmelCase      ( self					: Optional[int]      )	->			List[Any]:


         '''simple docstring'''





         return (4, 14, 16)



  @property
  def 							_UpperCAmelCase      ( self					: Dict      )	->			Optional[int]:


         '''simple docstring'''





         return (4, 14, 16)



  def 							_UpperCAmelCase      ( self					: List[Any]      )	->			Tuple:


         '''simple docstring'''





         pass



  def 							_UpperCAmelCase      ( self					: Optional[int]      )	->			Optional[Any]:


         '''simple docstring'''





         pass



  @unittest.skipIf(torch_device == '''mps'''       ,       '''mish op not supported in MPS'''      )
  def 							_UpperCAmelCase      ( self					: Any      )	->			int:


         '''simple docstring'''





         super().test_determinism()



  @unittest.skipIf(torch_device == '''mps'''       ,       '''mish op not supported in MPS'''      )
  def 							_UpperCAmelCase      ( self					: Any      )	->			Dict:


         '''simple docstring'''





         super().test_outputs_equivalence()



  @unittest.skipIf(torch_device == '''mps'''       ,       '''mish op not supported in MPS'''      )
  def 							_UpperCAmelCase      ( self					: int      )	->			List[str]:


         '''simple docstring'''





         super().test_from_save_pretrained()



  @unittest.skipIf(torch_device == '''mps'''       ,       '''mish op not supported in MPS'''      )
  def 							_UpperCAmelCase      ( self					: List[str]      )	->			Optional[int]:


         '''simple docstring'''





         super().test_from_save_pretrained_variant()



  @unittest.skipIf(torch_device == '''mps'''       ,       '''mish op not supported in MPS'''      )
  def 							_UpperCAmelCase      ( self					: Dict      )	->			List[Any]:


         '''simple docstring'''





         super().test_model_from_pretrained()



  @unittest.skipIf(torch_device == '''mps'''       ,       '''mish op not supported in MPS'''      )
  def 							_UpperCAmelCase      ( self					: str      )	->			Any:


         '''simple docstring'''





         super().test_output()



  def 							_UpperCAmelCase      ( self					: List[str]      )	->			str:


         '''simple docstring'''





         __magic_name__			:		Optional[Any]        =     {
             """block_out_channels""": (32, 64, 128, 256),
             """in_channels""": 14,
             """out_channels""": 14,
             """time_embedding_type""": """positional""",
             """use_timestep_embedding""": True,
             """flip_sin_to_cos""": False,
             """freq_shift""": 1.0,
             """out_block_type""": """OutConv1DBlock""",
             """mid_block_type""": """MidResTemporalBlock1D""",
             """down_block_types""": ("""DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D"""),
             """up_block_types""": ("""UpResnetBlock1D""", """UpResnetBlock1D""", """UpResnetBlock1D"""),
             """act_fn""": """mish""",
         }
         __magic_name__			:		Union[str, Any]        =     self.dummy_input
         return init_dict, inputs_dict



  @unittest.skipIf(torch_device == '''mps'''       ,       '''mish op not supported in MPS'''      )
  def 							_UpperCAmelCase      ( self					: Dict      )	->			Dict:


         '''simple docstring'''





         __magic_name__			:		int        =     UNetaDModel.from_pretrained(
             '''bglick13/hopper-medium-v2-value-function-hor32'''       ,       output_loading_info=snake_case       ,       subfolder='''unet'''      )
         self.assertIsNotNone(snake_case      )
         self.assertEqual(len(loading_info['''missing_keys''']      )       ,       0      )

         model.to(snake_case      )
         __magic_name__			:		Any        =     model(**self.dummy_input      )

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



  @unittest.skipIf(torch_device == '''mps'''       ,       '''mish op not supported in MPS'''      )
  def 							_UpperCAmelCase      ( self					: Union[str, Any]      )	->			Optional[Any]:


         '''simple docstring'''





         __magic_name__			:		Any        =     UNetaDModel.from_pretrained('''bglick13/hopper-medium-v2-value-function-hor32'''       ,       subfolder='''unet'''      )
         torch.manual_seed(0      )
         if torch.cuda.is_available():
                torch.cuda.manual_seed_all(0      )

         __magic_name__			:		Union[str, Any]        =     model.config.in_channels
         __magic_name__			:		str        =     16
         __magic_name__			:		Any        =     torch.randn((1, seq_len, num_features)      ).permute(
             0       ,       2       ,       1      )  # match original, we can update values and remove
         __magic_name__			:		Dict        =     torch.full((num_features,)       ,       0      )

         with torch.no_grad():
                __magic_name__			:		Optional[int]        =     model(snake_case       ,       snake_case      ).sample.permute(0       ,       2       ,       1      )

         __magic_name__			:		str        =     output[0, -3:, -3:].flatten()
         # fmt: off
         __magic_name__			:		int        =     torch.tensor([-2.13_7172, 1.142_6016, 0.368_8687, -0.76_6922, 0.730_3146, 0.1103_8864, -0.476_0633, 0.1327_0172, 0.0259_1348]      )
         # fmt: on
         self.assertTrue(torch.allclose(snake_case       ,       snake_case       ,       rtol=1e-3      )      )



  def 							_UpperCAmelCase      ( self					: Dict      )	->			Tuple:


         '''simple docstring'''





         pass




  @slow
  def 							_UpperCAmelCase      ( self					: Optional[int]      )	->			List[Any]:


         '''simple docstring'''





         __magic_name__			:		Optional[Any]        =     """harmonai/maestro-150k"""
         __magic_name__			:		Dict        =     UNetaDModel.from_pretrained(snake_case       ,       subfolder='''unet'''      )
         model.to(snake_case      )

         __magic_name__			:		Optional[Any]        =     6_5536
         __magic_name__			:		List[str]        =     torch.sin(torch.arange(snake_case      )[None, None, :].repeat(1       ,       2       ,       1      )      ).to(snake_case      )
         __magic_name__			:		Optional[int]        =     torch.tensor([1]      ).to(snake_case      )

         with torch.no_grad():
                __magic_name__			:		Tuple        =     model(snake_case       ,       snake_case      ).sample

         __magic_name__			:		Optional[Any]        =     output.abs().sum()
         __magic_name__			:		Optional[int]        =     output.abs().max()

         assert (output_sum - 224.0896).abs() < 0.5
         assert (output_max - 0.0607).abs() < 4e-4

class        _UpperCamelCase	(						lowerCamelCase__   ,				lowerCamelCase__   ,				unittest.TestCase	):


  """simple docstring"""





  snake_case_        =    UNetaDModel
  snake_case_        =    'sample'



  @property
  def 							_UpperCAmelCase      ( self					: Optional[Any]      )	->			int:


         '''simple docstring'''





         __magic_name__			:		Optional[Any]        =     4
         __magic_name__			:		List[Any]        =     14
         __magic_name__			:		str        =     16

         __magic_name__			:		List[str]        =     floats_tensor((batch_size, num_features, seq_len)      ).to(snake_case      )
         __magic_name__			:		List[str]        =     torch.tensor([10] * batch_size      ).to(snake_case      )

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



  @property
  def 							_UpperCAmelCase      ( self					: Dict      )	->			Union[str, Any]:


         '''simple docstring'''





         return (4, 14, 16)



  @property
  def 							_UpperCAmelCase      ( self					: int      )	->			Union[str, Any]:


         '''simple docstring'''





         return (4, 14, 1)



  @unittest.skipIf(torch_device == '''mps'''       ,       '''mish op not supported in MPS'''      )
  def 							_UpperCAmelCase      ( self					: Any      )	->			Dict:


         '''simple docstring'''





         super().test_determinism()



  @unittest.skipIf(torch_device == '''mps'''       ,       '''mish op not supported in MPS'''      )
  def 							_UpperCAmelCase      ( self					: Tuple      )	->			str:


         '''simple docstring'''





         super().test_outputs_equivalence()



  @unittest.skipIf(torch_device == '''mps'''       ,       '''mish op not supported in MPS'''      )
  def 							_UpperCAmelCase      ( self					: Optional[Any]      )	->			List[str]:


         '''simple docstring'''





         super().test_from_save_pretrained()



  @unittest.skipIf(torch_device == '''mps'''       ,       '''mish op not supported in MPS'''      )
  def 							_UpperCAmelCase      ( self					: List[str]      )	->			Tuple:


         '''simple docstring'''





         super().test_from_save_pretrained_variant()



  @unittest.skipIf(torch_device == '''mps'''       ,       '''mish op not supported in MPS'''      )
  def 							_UpperCAmelCase      ( self					: Any      )	->			int:


         '''simple docstring'''





         super().test_model_from_pretrained()



  @unittest.skipIf(torch_device == '''mps'''       ,       '''mish op not supported in MPS'''      )
  def 							_UpperCAmelCase      ( self					: List[Any]      )	->			Dict:


         '''simple docstring'''





         __magic_name__			:		Any        =     self.prepare_init_args_and_inputs_for_common()
         __magic_name__			:		List[Any]        =     self.model_class(**snake_case      )
         model.to(snake_case      )
         model.eval()

         with torch.no_grad():
                __magic_name__			:		Dict        =     model(**snake_case      )

                if isinstance(snake_case       ,       snake_case      ):
                       __magic_name__			:		Optional[int]        =     output.sample

         self.assertIsNotNone(snake_case      )
         __magic_name__			:		Union[str, Any]        =     torch.Size((inputs_dict['''sample'''].shape[0], 1)      )
         self.assertEqual(output.shape       ,       snake_case       ,       '''Input and output shapes do not match'''      )



  def 							_UpperCAmelCase      ( self					: Any      )	->			List[str]:


         '''simple docstring'''





         pass



  def 							_UpperCAmelCase      ( self					: Optional[int]      )	->			Tuple:


         '''simple docstring'''





         pass



  def 							_UpperCAmelCase      ( self					: int      )	->			Optional[int]:


         '''simple docstring'''





         __magic_name__			:		Union[str, Any]        =     {
             """in_channels""": 14,
             """out_channels""": 14,
             """down_block_types""": ["""DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D"""],
             """up_block_types""": [],
             """out_block_type""": """ValueFunction""",
             """mid_block_type""": """ValueFunctionMidBlock1D""",
             """block_out_channels""": [32, 64, 128, 256],
             """layers_per_block""": 1,
             """downsample_each_block""": True,
             """use_timestep_embedding""": True,
             """freq_shift""": 1.0,
             """flip_sin_to_cos""": False,
             """time_embedding_type""": """positional""",
             """act_fn""": """mish""",
         }
         __magic_name__			:		str        =     self.dummy_input
         return init_dict, inputs_dict



  @unittest.skipIf(torch_device == '''mps'''       ,       '''mish op not supported in MPS'''      )
  def 							_UpperCAmelCase      ( self					: Optional[Any]      )	->			Union[str, Any]:


         '''simple docstring'''





         __magic_name__			:		List[Any]        =     UNetaDModel.from_pretrained(
             '''bglick13/hopper-medium-v2-value-function-hor32'''       ,       output_loading_info=snake_case       ,       subfolder='''value_function'''      )
         self.assertIsNotNone(snake_case      )
         self.assertEqual(len(vf_loading_info['''missing_keys''']      )       ,       0      )

         value_function.to(snake_case      )
         __magic_name__			:		List[Any]        =     value_function(**self.dummy_input      )

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



  @unittest.skipIf(torch_device == '''mps'''       ,       '''mish op not supported in MPS'''      )
  def 							_UpperCAmelCase      ( self					: int      )	->			Optional[Any]:


         '''simple docstring'''





         __magic_name__			:		Tuple        =     UNetaDModel.from_pretrained(
             '''bglick13/hopper-medium-v2-value-function-hor32'''       ,       output_loading_info=snake_case       ,       subfolder='''value_function'''      )
         torch.manual_seed(0      )
         if torch.cuda.is_available():
                torch.cuda.manual_seed_all(0      )

         __magic_name__			:		Optional[int]        =     value_function.config.in_channels
         __magic_name__			:		List[str]        =     14
         __magic_name__			:		Optional[int]        =     torch.randn((1, seq_len, num_features)      ).permute(
             0       ,       2       ,       1      )  # match original, we can update values and remove
         __magic_name__			:		Union[str, Any]        =     torch.full((num_features,)       ,       0      )

         with torch.no_grad():
                __magic_name__			:		Tuple        =     value_function(snake_case       ,       snake_case      ).sample

         # fmt: off
         __magic_name__			:		Optional[int]        =     torch.tensor([165.25] * seq_len      )
         # fmt: on
         self.assertTrue(torch.allclose(snake_case       ,       snake_case       ,       rtol=1e-3      )      )




  def 							_UpperCAmelCase      ( self					: Tuple      )	->			Union[str, Any]:


         '''simple docstring'''





         pass








"""simple docstring"""



from itertools import combinations, product




def   _lowerCAmelCase		(				UpperCAmelCase__							:       int,    UpperCAmelCase__							:       int,    UpperCAmelCase__							:       int,    UpperCAmelCase__							:       int   )		->bool:
						if xa == ya == 0 or xa == ya == 0:
												return False
						A__    :	List[str]			=	xa * xa + ya * ya
						A__    :	List[Any]			=	xa * xa + ya * ya
						A__    :	Optional[int]			=	(xa - xa) * (xa - xa) + (ya - ya) * (ya - ya)
						return (
						    a_square + b_square == c_square
						    or a_square + c_square == b_square
						    or b_square + c_square == a_square
						)




def   _lowerCAmelCase		(				UpperCAmelCase__							:       int = 5_0   )		->int:
						return sum(
						    1
						    for pta, pta in combinations(product(range(limit + 1   ),    repeat=2   ),    2   )
						    if is_right(*UpperCAmelCase__,    *UpperCAmelCase__   )   )


if __name__ == "__main__":
						print(F'{solution() = }')






