


import itertools
import math


def 		lowercase_  (     _UpperCamelCase  ):


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


def 		lowercase_  (     ):


			'''simple docstring'''







			__lowercase 					=   2
			while True:
						if is_prime(_UpperCamelCase  ):
									yield num
						num += 1


def 		lowercase_  (     _UpperCamelCase = 1_00_01  ):


			'''simple docstring'''







			return next(itertools.islice(prime_generator()       ,			nth - 1       ,			_UpperCamelCase  )  )


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






import unittest

from diffusers import FlaxAutoencoderKL
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax

from .test_modeling_common_flax import FlaxModelTesterMixin


if is_flax_available():
	import jax
@require_flax
class 	lowerCamelCase_       (				lowerCAmelCase__     ,     unittest.TestCase				):



					'''simple docstring'''







					__UpperCAmelCase            =				FlaxAutoencoderKL






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







								'''simple docstring'''

								__lowercase 					=   4
								__lowercase 					=   3
								__lowercase 					=   (3_2, 3_2)

								__lowercase 					=   jax.random.PRNGKey(0						)
								__lowercase 					=   jax.random.uniform(snake_case_ ,       ((batch_size, num_channels) + sizes)						)

								return {"sample": image, "prng_key": prng_key}






					def A      (				self						)   ->			Any:







								'''simple docstring'''

								__lowercase 					=   {
								    '''block_out_channels''': [3_2, 6_4],
								    '''in_channels''': 3,
								    '''out_channels''': 3,
								    '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
								    '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
								    '''latent_channels''': 4,
								}
								__lowercase 					=   self.dummy_input
								return init_dict, inputs_dict


