



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 	a__       (       _snake_case  ,			_snake_case  ,			unittest.TestCase       ):


					"""simple docstring"""



					A__  :				Any					 =				UNetaDModel
					A__  :				Any					 =				'''sample'''

					@property
					def 				__UpperCAmelCase		(		self     :str    ):
										lowercase						   =       4
										lowercase						   =       14
										lowercase						   =       16

										lowercase						   =       floats_tensor((batch_size, num_features, seq_len)    ).to(lowercase__    )
										lowercase						   =       torch.tensor([10] * batch_size    ).to(lowercase__    )

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

					@property
					def 				__UpperCAmelCase		(		self     :str    ):
										return (4, 14, 16)

					@property
					def 				__UpperCAmelCase		(		self     :Optional[int]    ):
										return (4, 14, 16)

					def 				__UpperCAmelCase		(		self     :Any    ):
										pass

					def 				__UpperCAmelCase		(		self     :Dict    ):
										pass

					@unittest.skipIf(torch_device == 'mps'      ,  'mish op not supported in MPS'    )
					def 				__UpperCAmelCase		(		self     :List[Any]    ):
										super().test_determinism()

					@unittest.skipIf(torch_device == 'mps'      ,  'mish op not supported in MPS'    )
					def 				__UpperCAmelCase		(		self     :Dict    ):
										super().test_outputs_equivalence()

					@unittest.skipIf(torch_device == 'mps'      ,  'mish op not supported in MPS'    )
					def 				__UpperCAmelCase		(		self     :Optional[int]    ):
										super().test_from_save_pretrained()

					@unittest.skipIf(torch_device == 'mps'      ,  'mish op not supported in MPS'    )
					def 				__UpperCAmelCase		(		self     :List[str]    ):
										super().test_from_save_pretrained_variant()

					@unittest.skipIf(torch_device == 'mps'      ,  'mish op not supported in MPS'    )
					def 				__UpperCAmelCase		(		self     :int    ):
										super().test_model_from_pretrained()

					@unittest.skipIf(torch_device == 'mps'      ,  'mish op not supported in MPS'    )
					def 				__UpperCAmelCase		(		self     :int    ):
										super().test_output()

					def 				__UpperCAmelCase		(		self     :List[str]    ):
										lowercase						   =       {
										    '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',
										}
										lowercase						   =       self.dummy_input
										return init_dict, inputs_dict

					@unittest.skipIf(torch_device == 'mps'      ,  'mish op not supported in MPS'    )
					def 				__UpperCAmelCase		(		self     :Optional[int]    ):
										lowercase   ,      lowercase						   =       UNetaDModel.from_pretrained(
										    'bglick13/hopper-medium-v2-value-function-hor32'      ,  output_loading_info=lowercase__      ,  subfolder='unet'    )
										self.assertIsNotNone(lowercase__    )
										self.assertEqual(len(loading_info['missing_keys']    )      ,  0    )

										model.to(lowercase__    )
										lowercase						   =       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     :List[str]    ):
										lowercase						   =       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    )

										lowercase						   =       model.config.in_channels
										lowercase						   =       16
										lowercase						   =       torch.randn((1, seq_len, num_features)    ).permute(
										    0      ,  2      ,  1    )  # match original, we can update values and remove
										lowercase						   =       torch.full((num_features,)      ,  0    )

										with torch.no_grad():
															lowercase						   =       model(lowercase__      ,  lowercase__    ).sample.permute(0      ,  2      ,  1    )

										lowercase						   =       output[0, -3:, -3:].flatten()
										# fmt: off
										lowercase						   =       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(lowercase__      ,  lowercase__      ,  rtol=1E-3    )    )

					def 				__UpperCAmelCase		(		self     :int    ):
										# Not implemented yet for this UNet
										pass



					@slow
					def 				__UpperCAmelCase		(		self     :str    ):
										lowercase						   =       'harmonai/maestro-150k'
										lowercase						   =       UNetaDModel.from_pretrained(lowercase__      ,  subfolder='unet'    )
										model.to(lowercase__    )

										lowercase						   =       6_5536
										lowercase						   =       torch.sin(torch.arange(lowercase__    )[None, None, :].repeat(1      ,  2      ,  1    )    ).to(lowercase__    )
										lowercase						   =       torch.tensor([1]    ).to(lowercase__    )

										with torch.no_grad():
															lowercase						   =       model(lowercase__      ,  lowercase__    ).sample

										lowercase						   =       output.abs().sum()
										lowercase						   =       output.abs().max()

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




class 	a__       (       _snake_case  ,			_snake_case  ,			unittest.TestCase       ):


					"""simple docstring"""



					A__  :				int					 =				UNetaDModel
					A__  :				Any					 =				'''sample'''

					@property
					def 				__UpperCAmelCase		(		self     :Optional[int]    ):
										lowercase						   =       4
										lowercase						   =       14
										lowercase						   =       16

										lowercase						   =       floats_tensor((batch_size, num_features, seq_len)    ).to(lowercase__    )
										lowercase						   =       torch.tensor([10] * batch_size    ).to(lowercase__    )

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

					@property
					def 				__UpperCAmelCase		(		self     :Dict    ):
										return (4, 14, 16)

					@property
					def 				__UpperCAmelCase		(		self     :List[Any]    ):
										return (4, 14, 1)

					@unittest.skipIf(torch_device == 'mps'      ,  'mish op not supported in MPS'    )
					def 				__UpperCAmelCase		(		self     :List[str]    ):
										super().test_determinism()

					@unittest.skipIf(torch_device == 'mps'      ,  'mish op not supported in MPS'    )
					def 				__UpperCAmelCase		(		self     :Union[str, Any]    ):
										super().test_outputs_equivalence()

					@unittest.skipIf(torch_device == 'mps'      ,  'mish op not supported in MPS'    )
					def 				__UpperCAmelCase		(		self     :Optional[int]    ):
										super().test_from_save_pretrained()

					@unittest.skipIf(torch_device == 'mps'      ,  'mish op not supported in MPS'    )
					def 				__UpperCAmelCase		(		self     :Optional[Any]    ):
										super().test_from_save_pretrained_variant()

					@unittest.skipIf(torch_device == 'mps'      ,  'mish op not supported in MPS'    )
					def 				__UpperCAmelCase		(		self     :Optional[int]    ):
										super().test_model_from_pretrained()

					@unittest.skipIf(torch_device == 'mps'      ,  'mish op not supported in MPS'    )
					def 				__UpperCAmelCase		(		self     :Dict    ):
										# UNetRL is a value-function is different output shape
										lowercase   ,      lowercase						   =       self.prepare_init_args_and_inputs_for_common()
										lowercase						   =       self.model_class(**lowercase__    )
										model.to(lowercase__    )
										model.eval()

										with torch.no_grad():
															lowercase						   =       model(**lowercase__    )

															if isinstance(lowercase__      ,  lowercase__    ):
																				lowercase						   =       output.sample

										self.assertIsNotNone(lowercase__    )
										lowercase						   =       torch.Size((inputs_dict['sample'].shape[0], 1)    )
										self.assertEqual(output.shape      ,  lowercase__      ,  'Input and output shapes do not match'    )

					def 				__UpperCAmelCase		(		self     :Optional[int]    ):
										pass

					def 				__UpperCAmelCase		(		self     :Any    ):
										pass

					def 				__UpperCAmelCase		(		self     :Any    ):
										lowercase						   =       {
										    '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',
										}
										lowercase						   =       self.dummy_input
										return init_dict, inputs_dict

					@unittest.skipIf(torch_device == 'mps'      ,  'mish op not supported in MPS'    )
					def 				__UpperCAmelCase		(		self     :Union[str, Any]    ):
										lowercase   ,      lowercase						   =       UNetaDModel.from_pretrained(
										    'bglick13/hopper-medium-v2-value-function-hor32'      ,  output_loading_info=lowercase__      ,  subfolder='value_function'    )
										self.assertIsNotNone(lowercase__    )
										self.assertEqual(len(vf_loading_info['missing_keys']    )      ,  0    )

										value_function.to(lowercase__    )
										lowercase						   =       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     :List[Any]    ):
										lowercase   ,      lowercase						   =       UNetaDModel.from_pretrained(
										    'bglick13/hopper-medium-v2-value-function-hor32'      ,  output_loading_info=lowercase__      ,  subfolder='value_function'    )
										torch.manual_seed(0    )
										if torch.cuda.is_available():
															torch.cuda.manual_seed_all(0    )

										lowercase						   =       value_function.config.in_channels
										lowercase						   =       14
										lowercase						   =       torch.randn((1, seq_len, num_features)    ).permute(
										    0      ,  2      ,  1    )  # match original, we can update values and remove
										lowercase						   =       torch.full((num_features,)      ,  0    )

										with torch.no_grad():
															lowercase						   =       value_function(lowercase__      ,  lowercase__    ).sample

										# fmt: off
										lowercase						   =       torch.tensor([165.25] * seq_len    )
										# fmt: on
										self.assertTrue(torch.allclose(lowercase__      ,  lowercase__      ,  rtol=1E-3    )    )



					def 				__UpperCAmelCase		(		self     :str    ):
										# Not implemented yet for this UNet
										pass











from collections.abc import Callable

__magic_name__       =     Callable[[float], float]  # type alias for a real -> real function





def 		__snake_case						(    _UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase       ,			):
					"""simple docstring"""






					lowercase						   =       float(_UpperCAmelCase				)
					while True:
										try:
															lowercase						   =       prev_guess - function(_UpperCAmelCase				) / derivative(_UpperCAmelCase				)
										except ZeroDivisionError:
															raise ZeroDivisionError('Could not find root'				) from None
										if abs(prev_guess - next_guess				) < 10**-5:
															return next_guess
										lowercase						   =       next_guess





def 		__snake_case						(    _UpperCAmelCase				):
					"""simple docstring"""






					return (x**3) - (2 * x) - 5





def 		__snake_case						(    _UpperCAmelCase				):
					"""simple docstring"""






					return 3 * (x**2) - 2


if __name__ == "__main__":
		print(newton(f, fa, 3))






