from __future__ import annotations

import math





def 				UpperCAmelCase						(    lowercase							):







				"""simple docstring"""



				if 1 < number < 4:
								# 2 and 3 are primes
								return True
				elif number < 2 or number % 2 == 0 or number % 3 == 0:
								# Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
								return False

				# All primes number are in format of 6k +/- 1
				for i in range(5							,			int(math.sqrt(lowercase							) + 1							)							,			6							):
								if number % i == 0 or number % (i + 2) == 0:
												return False
				return True





def 				UpperCAmelCase						(    lowercase							):







				"""simple docstring"""



				__lowercase 			=							str(lowercase							)
				__lowercase 			=							[n]
				for i in range(1							,			len(lowercase							)							):
								list_nums.append(int(str_num[i:]							)							)
								list_nums.append(int(str_num[:-i]							)							)
				return list_nums





def 				UpperCAmelCase						(    lowercase							):







				"""simple docstring"""



				if len(str(lowercase							)							) > 3:
								if not is_prime(int(str(lowercase							)[-3:]							)							) or not is_prime(int(str(lowercase							)[:3]							)							):
												return False
				return True





def 				UpperCAmelCase						(    lowercase = 11							):







				"""simple docstring"""



				__lowercase 			=							[]
				__lowercase 			=							13
				while len(lowercase							) != count:
								if validate(lowercase							):
												__lowercase 			=							list_truncated_nums(lowercase							)
												if all(is_prime(lowercase							) for i in list_nums							):
																list_truncated_primes.append(lowercase							)
								num += 2
				return list_truncated_primes





def 				UpperCAmelCase						(    ):







				"""simple docstring"""



				return sum(compute_truncated_primes(11							)							)


if __name__ == "__main__":
		print(F'''{sum(compute_truncated_primes(1_1)) = }''')
import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest




class    _UpperCamelCase      (						_UpperCAmelCase ):


			"""simple docstring"""






			__a  : Optional[Any]			 =						(KDPMaDiscreteScheduler,)
			__a  : Dict			 =						10





			def        _SCREAMING_SNAKE_CASE			(   self      ,   **lowerCAmelCase__						)     ->					int:

							'''simple docstring'''


							__lowercase 			=							{
							    '''num_train_timesteps''': 11_00,
							    '''beta_start''': 0.0001,
							    '''beta_end''': 0.02,
							    '''beta_schedule''': '''linear''',
							}

							config.update(**lowerCAmelCase__						)
							return config





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

							'''simple docstring'''


							for timesteps in [10, 50, 1_00, 10_00]:
											self.check_over_configs(num_train_timesteps=lowerCAmelCase__						)





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Any:

							'''simple docstring'''


							for beta_start, beta_end in zip([0.0_0001, 0.0001, 0.001]      ,   [0.0002, 0.002, 0.02]						):
											self.check_over_configs(beta_start=lowerCAmelCase__      ,   beta_end=lowerCAmelCase__						)





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

							'''simple docstring'''


							for schedule in ["linear", "scaled_linear"]:
											self.check_over_configs(beta_schedule=lowerCAmelCase__						)





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

							'''simple docstring'''


							for prediction_type in ["epsilon", "v_prediction"]:
											self.check_over_configs(prediction_type=lowerCAmelCase__						)





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Tuple:

							'''simple docstring'''


							__lowercase 			=							self.scheduler_classes[0]
							__lowercase 			=							self.get_scheduler_config(prediction_type='''v_prediction'''						)
							__lowercase 			=							scheduler_class(**lowerCAmelCase__						)

							scheduler.set_timesteps(self.num_inference_steps						)

							__lowercase 			=							self.dummy_model()
							__lowercase 			=							self.dummy_sample_deter * scheduler.init_noise_sigma
							__lowercase 			=							sample.to(lowerCAmelCase__						)

							for i, t in enumerate(scheduler.timesteps						):
											__lowercase 			=							scheduler.scale_model_input(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							model(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							scheduler.step(lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__						)
											__lowercase 			=							output.prev_sample

							__lowercase 			=							torch.sum(torch.abs(lowerCAmelCase__						)						)
							__lowercase 			=							torch.mean(torch.abs(lowerCAmelCase__						)						)

							if torch_device in ["cpu", "mps"]:
											assert abs(result_sum.item() - 4.6934E-07						) < 1E-2
											assert abs(result_mean.item() - 6.1112E-10						) < 1E-3
							else:
											# CUDA
											assert abs(result_sum.item() - 4.693_4286_5017_0972E-07						) < 1E-2
											assert abs(result_mean.item() - 0.0002						) < 1E-3





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					str:

							'''simple docstring'''


							if torch_device == "mps":
											return
							__lowercase 			=							self.scheduler_classes[0]
							__lowercase 			=							self.get_scheduler_config()
							__lowercase 			=							scheduler_class(**lowerCAmelCase__						)

							scheduler.set_timesteps(self.num_inference_steps						)

							__lowercase 			=							self.dummy_model()
							__lowercase 			=							self.dummy_sample_deter * scheduler.init_noise_sigma
							__lowercase 			=							sample.to(lowerCAmelCase__						)

							for i, t in enumerate(scheduler.timesteps						):
											__lowercase 			=							scheduler.scale_model_input(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							model(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							scheduler.step(lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__						)
											__lowercase 			=							output.prev_sample

							__lowercase 			=							torch.sum(torch.abs(lowerCAmelCase__						)						)
							__lowercase 			=							torch.mean(torch.abs(lowerCAmelCase__						)						)

							if torch_device in ["cpu", "mps"]:
											assert abs(result_sum.item() - 20.4125						) < 1E-2
											assert abs(result_mean.item() - 0.0266						) < 1E-3
							else:
											# CUDA
											assert abs(result_sum.item() - 20.4125						) < 1E-2
											assert abs(result_mean.item() - 0.0266						) < 1E-3





			def        _SCREAMING_SNAKE_CASE			(   self						)     ->					Dict:

							'''simple docstring'''


							if torch_device == "mps":
											return
							__lowercase 			=							self.scheduler_classes[0]
							__lowercase 			=							self.get_scheduler_config()
							__lowercase 			=							scheduler_class(**lowerCAmelCase__						)

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

							__lowercase 			=							self.dummy_model()
							__lowercase 			=							self.dummy_sample_deter.to(lowerCAmelCase__						) * scheduler.init_noise_sigma

							for t in scheduler.timesteps:
											__lowercase 			=							scheduler.scale_model_input(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							model(lowerCAmelCase__      ,   lowerCAmelCase__						)

											__lowercase 			=							scheduler.step(lowerCAmelCase__      ,   lowerCAmelCase__      ,   lowerCAmelCase__						)
											__lowercase 			=							output.prev_sample

							__lowercase 			=							torch.sum(torch.abs(lowerCAmelCase__						)						)
							__lowercase 			=							torch.mean(torch.abs(lowerCAmelCase__						)						)

							if str(lowerCAmelCase__						).startswith('''cpu'''						):
											# The following sum varies between 148 and 156 on mps. Why?
											assert abs(result_sum.item() - 20.4125						) < 1E-2
											assert abs(result_mean.item() - 0.0266						) < 1E-3
							else:
											# CUDA
											assert abs(result_sum.item() - 20.4125						) < 1E-2
											assert abs(result_mean.item() - 0.0266						) < 1E-3