

'''simple docstring'''

import torch

from diffusers import DPMSolverSDEScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import require_torchsde

from .test_schedulers import SchedulerCommonTest





@require_torchsde
class 							_snake_case (       SCREAMING_SNAKE_CASE__						):
   _A       :     List[Any]       =					(DPMSolverSDEScheduler,)
   _A       :     List[str]       =					1_0
   def     __UpperCamelCase						( self			:      Any     ,**SCREAMING_SNAKE_CASE__			:      str							):
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					{
         "num_train_timesteps": 1_100,
         "beta_start": 0.0_001,
         "beta_end": 0.02,
         "beta_schedule": "linear",
         "noise_sampler_seed": 0,
     }

     config.update(**snake_case__							)
     return config
   def     __UpperCamelCase						( self			:      Tuple							):
     for timesteps in [10, 50, 100, 1_000]:
       self.check_over_configs(num_train_timesteps=snake_case__							)
   def     __UpperCamelCase						( self			:      Union[str, Any]							):
     for beta_start, beta_end in zip([0.00_001, 0.0_001, 0.001]     ,[0.0_002, 0.002, 0.02]							):
       self.check_over_configs(beta_start=snake_case__     ,beta_end=snake_case__							)
   def     __UpperCamelCase						( self			:      Union[str, Any]							):
     for schedule in ["linear", "scaled_linear"]:
       self.check_over_configs(beta_schedule=snake_case__							)
   def     __UpperCamelCase						( self			:      Optional[Any]							):
     for prediction_type in ["epsilon", "v_prediction"]:
       self.check_over_configs(prediction_type=snake_case__							)
   def     __UpperCamelCase						( self			:      Optional[int]							):
     SCREAMING_SNAKE_CASE:int						      =					self.scheduler_classes[0]
     SCREAMING_SNAKE_CASE:List[Any]						      =					self.get_scheduler_config()
     SCREAMING_SNAKE_CASE:Tuple						      =					scheduler_class(**snake_case__							)

     scheduler.set_timesteps(self.num_inference_steps							)

     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					self.dummy_model()
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					self.dummy_sample_deter * scheduler.init_noise_sigma
     SCREAMING_SNAKE_CASE:List[Any]						      =					sample.to(snake_case__							)

     for i, t in enumerate(scheduler.timesteps							):
       SCREAMING_SNAKE_CASE:str						      =					scheduler.scale_model_input(snake_case__     ,snake_case__							)

       SCREAMING_SNAKE_CASE:List[str]						      =					model(snake_case__     ,snake_case__							)

       SCREAMING_SNAKE_CASE:Dict						      =					scheduler.step(snake_case__     ,snake_case__     ,snake_case__							)
       SCREAMING_SNAKE_CASE:Tuple						      =					output.prev_sample

     SCREAMING_SNAKE_CASE:List[str]						      =					torch.sum(torch.abs(snake_case__							)							)
     SCREAMING_SNAKE_CASE:int						      =					torch.mean(torch.abs(snake_case__							)							)

     if torch_device in ["mps"]:
       assert abs(result_sum.item() - 167.47_821_044_921_875							) < 1e-2
       assert abs(result_mean.item() - 0.2_178_705_964_565_277							) < 1e-3
     elif torch_device in ["cuda"]:
       assert abs(result_sum.item() - 171.59_352_111_816_406							) < 1e-2
       assert abs(result_mean.item() - 0.22_342_906_892_299_652							) < 1e-3
     else:
       assert abs(result_sum.item() - 162.52_383_422_851_562							) < 1e-2
       assert abs(result_mean.item() - 0.211_619_570_851_326							) < 1e-3
   def     __UpperCamelCase						( self			:      Union[str, Any]							):
     SCREAMING_SNAKE_CASE:str						      =					self.scheduler_classes[0]
     SCREAMING_SNAKE_CASE:Dict						      =					self.get_scheduler_config(prediction_type="v_prediction"							)
     SCREAMING_SNAKE_CASE:Dict						      =					scheduler_class(**snake_case__							)

     scheduler.set_timesteps(self.num_inference_steps							)

     SCREAMING_SNAKE_CASE:Optional[Any]						      =					self.dummy_model()
     SCREAMING_SNAKE_CASE:Dict						      =					self.dummy_sample_deter * scheduler.init_noise_sigma
     SCREAMING_SNAKE_CASE:Optional[Any]						      =					sample.to(snake_case__							)

     for i, t in enumerate(scheduler.timesteps							):
       SCREAMING_SNAKE_CASE:Optional[int]						      =					scheduler.scale_model_input(snake_case__     ,snake_case__							)

       SCREAMING_SNAKE_CASE:List[str]						      =					model(snake_case__     ,snake_case__							)

       SCREAMING_SNAKE_CASE:Optional[int]						      =					scheduler.step(snake_case__     ,snake_case__     ,snake_case__							)
       SCREAMING_SNAKE_CASE:str						      =					output.prev_sample

     SCREAMING_SNAKE_CASE:int						      =					torch.sum(torch.abs(snake_case__							)							)
     SCREAMING_SNAKE_CASE:Optional[Any]						      =					torch.mean(torch.abs(snake_case__							)							)

     if torch_device in ["mps"]:
       assert abs(result_sum.item() - 124.77_149_200_439_453							) < 1e-2
       assert abs(result_mean.item() - 0.16_226_289_014_816_284							) < 1e-3
     elif torch_device in ["cuda"]:
       assert abs(result_sum.item() - 128.1_663_360_595_703							) < 1e-2
       assert abs(result_mean.item() - 0.16_688_326_001_167_297							) < 1e-3
     else:
       assert abs(result_sum.item() - 119.8_487_548_828_125							) < 1e-2
       assert abs(result_mean.item() - 0.1_560_530_662_536_621							) < 1e-3
   def     __UpperCamelCase						( self			:      Union[str, Any]							):
     SCREAMING_SNAKE_CASE:Union[str, Any]						      =					self.scheduler_classes[0]
     SCREAMING_SNAKE_CASE:str						      =					self.get_scheduler_config()
     SCREAMING_SNAKE_CASE:List[Any]						      =					scheduler_class(**snake_case__							)

     scheduler.set_timesteps(self.num_inference_steps     ,device=snake_case__							)

     SCREAMING_SNAKE_CASE:Optional[Any]						      =					self.dummy_model()
     SCREAMING_SNAKE_CASE:List[Any]						      =					self.dummy_sample_deter.to(snake_case__							) * scheduler.init_noise_sigma

     for t in scheduler.timesteps:
       SCREAMING_SNAKE_CASE:Optional[int]						      =					scheduler.scale_model_input(snake_case__     ,snake_case__							)

       SCREAMING_SNAKE_CASE:str						      =					model(snake_case__     ,snake_case__							)

       SCREAMING_SNAKE_CASE:Any						      =					scheduler.step(snake_case__     ,snake_case__     ,snake_case__							)
       SCREAMING_SNAKE_CASE:Tuple						      =					output.prev_sample

     SCREAMING_SNAKE_CASE:Tuple						      =					torch.sum(torch.abs(snake_case__							)							)
     SCREAMING_SNAKE_CASE:Optional[Any]						      =					torch.mean(torch.abs(snake_case__							)							)

     if torch_device in ["mps"]:
       assert abs(result_sum.item() - 167.46_957_397_460_938							) < 1e-2
       assert abs(result_mean.item() - 0.21_805_934_607_982_635							) < 1e-3
     elif torch_device in ["cuda"]:
       assert abs(result_sum.item() - 171.59_353_637_695_312							) < 1e-2
       assert abs(result_mean.item() - 0.22_342_908_382_415_771							) < 1e-3
     else:
       assert abs(result_sum.item() - 162.52_383_422_851_562							) < 1e-2
       assert abs(result_mean.item() - 0.211_619_570_851_326							) < 1e-3







   def     __UpperCamelCase						( self			:      Tuple							):
     SCREAMING_SNAKE_CASE:Any						      =					self.scheduler_classes[0]
     SCREAMING_SNAKE_CASE:Tuple						      =					self.get_scheduler_config()
     SCREAMING_SNAKE_CASE:Optional[int]						      =					scheduler_class(**snake_case__     ,use_karras_sigmas=snake_case__							)

     scheduler.set_timesteps(self.num_inference_steps     ,device=snake_case__							)

     SCREAMING_SNAKE_CASE:Optional[Any]						      =					self.dummy_model()
     SCREAMING_SNAKE_CASE:Any						      =					self.dummy_sample_deter.to(snake_case__							) * scheduler.init_noise_sigma
     SCREAMING_SNAKE_CASE:List[Any]						      =					sample.to(snake_case__							)

     for t in scheduler.timesteps:
       SCREAMING_SNAKE_CASE:str						      =					scheduler.scale_model_input(snake_case__     ,snake_case__							)

       SCREAMING_SNAKE_CASE:List[str]						      =					model(snake_case__     ,snake_case__							)

       SCREAMING_SNAKE_CASE:str						      =					scheduler.step(snake_case__     ,snake_case__     ,snake_case__							)
       SCREAMING_SNAKE_CASE:str						      =					output.prev_sample

     SCREAMING_SNAKE_CASE:List[str]						      =					torch.sum(torch.abs(snake_case__							)							)
     SCREAMING_SNAKE_CASE:List[str]						      =					torch.mean(torch.abs(snake_case__							)							)

     if torch_device in ["mps"]:
       assert abs(result_sum.item() - 176.66_974_135_742_188							) < 1e-2
       assert abs(result_mean.item() - 0.23_003_872_730_981_811							) < 1e-2
     elif torch_device in ["cuda"]:
       assert abs(result_sum.item() - 177.63_653_564_453_125							) < 1e-2
       assert abs(result_mean.item() - 0.23_003_872_730_981_811							) < 1e-2
     else:
       assert abs(result_sum.item() - 170.3_135_223_388_672							) < 1e-2
       assert abs(result_mean.item() - 0.23_003_872_730_981_811							) < 1e-2



'''simple docstring'''

def    A_  ( snake_case ):
  SCREAMING_SNAKE_CASE:Dict						      =					(1 + 24 * n) ** 0.5
  return ((1 + root) / 6) % 1 == 0





def    A_  ( snake_case = 5000 ):
  SCREAMING_SNAKE_CASE:int						      =					[(i * (3 * i - 1)) // 2 for i in range(1      ,       snake_case )]
  for i, pentagonal_i in enumerate(snake_case ):
    for j in range(snake_case      ,       len(snake_case ) ):
      SCREAMING_SNAKE_CASE:int						      =					pentagonal_nums[j]
      SCREAMING_SNAKE_CASE:Any						      =					pentagonal_i + pentagonal_j
      SCREAMING_SNAKE_CASE:int						      =					pentagonal_j - pentagonal_i
      if is_pentagonal(snake_case ) and is_pentagonal(snake_case ):
        return b

  return -1


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






