




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__   :






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											return self.get_dummy_input()






						@property
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""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     __SCREAMING_SNAKE_CASE(       self		,      _A=True		,      _A=False		,      _A=False		,      _A=False		,      ):




											"""simple docstring"""



											__lowerCAmelCase           =	4
											__lowerCAmelCase           =	3_2
											__lowerCAmelCase           =	(3_2, 3_2)

											__lowerCAmelCase           =	torch.manual_seed(0      )
											__lowerCAmelCase           =	torch.device(_A      )
											__lowerCAmelCase           =	(batch_size, num_channels) + sizes
											__lowerCAmelCase           =	randn_tensor(_A		,      generator=_A		,      device=_A      )
											__lowerCAmelCase           =	{"hidden_states": hidden_states}

											if include_temb:
																__lowerCAmelCase           =	1_2_8
																__lowerCAmelCase           =	randn_tensor((batch_size, temb_channels)		,      generator=_A		,      device=_A      )

											if include_res_hidden_states_tuple:
																__lowerCAmelCase           =	torch.manual_seed(1      )
																__lowerCAmelCase           =	(randn_tensor(_A		,      generator=_A		,      device=_A      ),)

											if include_encoder_hidden_states:
																__lowerCAmelCase           =	floats_tensor((batch_size, 3_2, 3_2)      ).to(_A      )

											if include_skip_sample:
																__lowerCAmelCase           =	randn_tensor(((batch_size, 3) + sizes)		,      generator=_A		,      device=_A      )

											return dummy_input






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	{
											    "in_channels": 3_2,
											    "out_channels": 3_2,
											    "temb_channels": 1_2_8,
											}
											if self.block_type == "up":
																__lowerCAmelCase           =	3_2

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

											__lowerCAmelCase           =	self.dummy_input
											return init_dict, inputs_dict






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase       ,							__lowerCAmelCase           =	self.prepare_init_args_and_inputs_for_common()
											__lowerCAmelCase           =	self.block_class(**_A      )
											unet_block.to(_A      )
											unet_block.eval()

											with torch.no_grad():
																__lowerCAmelCase           =	unet_block(**_A      )

											if isinstance(_A		,      _A      ):
																__lowerCAmelCase           =	output[0]

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

											__lowerCAmelCase           =	output[0, -1, -3:, -3:]
											__lowerCAmelCase           =	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     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase       ,							__lowerCAmelCase           =	self.prepare_init_args_and_inputs_for_common()
											__lowerCAmelCase           =	self.block_class(**_A      )
											model.to(_A      )
											model.train()
											__lowerCAmelCase           =	model(**_A      )

											if isinstance(_A		,      _A      ):
																__lowerCAmelCase           =	output[0]

											__lowerCAmelCase           =	torch.device(_A      )
											__lowerCAmelCase           =	randn_tensor(output.shape		,      device=_A      )
											__lowerCAmelCase           =	torch.nn.functional.mse_loss(_A		,      _A      )
											loss.backward()










import math







def      _a				(						SCREAMING_SNAKE_CASE_							:    int					):

					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(SCREAMING_SNAKE_CASE_					) + 1					)				,				6					):
										if number % i == 0 or number % (i + 2) == 0:
															return False
					return True







def      _a				(						SCREAMING_SNAKE_CASE_							:    float = 0.1					):
					__lowerCAmelCase           =	3
					__lowerCAmelCase           =	3

					while primes / (2 * j - 1) >= ratio:
										for i in range(j * j + j + 1				,				(j + 2) * (j + 2)				,				j + 1					):
															primes += is_prime(SCREAMING_SNAKE_CASE_					)
										j += 2
					return j


if __name__ == "__main__":
					import doctest

					doctest.testmod()




