





import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
      import jax
      import jax.numpy as jnp
      from flax.jax_utils import replicate
      from flax.training.common_utils import shard






@slow
@require_flax
class        snake_case_    (		unittest.TestCase						):






  def 				UpperCAmelCase__						(							self			:					int      )->Dict:






   '''simple docstring'''

   super().tearDown()
   gc.collect()






  def 				UpperCAmelCase__						(							self			:					Optional[int]      )->List[Any]:






   '''simple docstring'''

   __lowerCAmelCase ,					__lowerCAmelCase					:	int      =       FlaxStableDiffusionPipeline.from_pretrained(
       """stabilityai/stable-diffusion-2"""       ,  revision="""bf16"""       ,  dtype=jnp.bfloataa       ,  )

   __lowerCAmelCase					:	Dict      =       """A painting of a squirrel eating a burger"""
   __lowerCAmelCase					:	Any      =       jax.device_count()
   __lowerCAmelCase					:	List[str]      =       num_samples * [prompt]
   __lowerCAmelCase					:	Tuple      =       sd_pipe.prepare_inputs(_snake_case      )

   __lowerCAmelCase					:	Union[str, Any]      =       replicate(_snake_case      )
   __lowerCAmelCase					:	str      =       shard(_snake_case      )

   __lowerCAmelCase					:	Optional[Any]      =       jax.random.PRNGKey(0      )
   __lowerCAmelCase					:	List[str]      =       jax.random.split(_snake_case       ,  jax.device_count()      )

   __lowerCAmelCase					:	Optional[int]      =       sd_pipe(_snake_case       ,  _snake_case       ,  _snake_case       ,  num_inference_steps=25       ,  jit=_snake_case      )[0]
   assert images.shape == (jax.device_count(), 1, 768, 768, 3)

   __lowerCAmelCase					:	List[Any]      =       images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]      )
   __lowerCAmelCase					:	List[str]      =       images[0, 253:256, 253:256, -1]

   __lowerCAmelCase					:	Optional[Any]      =       jnp.asarray(jax.device_get(image_slice.flatten()      )      )
   __lowerCAmelCase					:	Optional[int]      =       jnp.array([0.4_238, 0.4_414, 0.4_395, 0.4_453, 0.4_629, 0.4_590, 0.4_531, 0.45_508, 0.4_512]      )
   print(F'''output_slice: {output_slice}'''      )
   assert jnp.abs(output_slice - expected_slice      ).max() < 1E-2






  def 				UpperCAmelCase__						(							self			:					Dict      )->Optional[Any]:






   '''simple docstring'''

   __lowerCAmelCase					:	Union[str, Any]      =       """stabilityai/stable-diffusion-2"""
   __lowerCAmelCase ,					__lowerCAmelCase					:	str      =       FlaxDPMSolverMultistepScheduler.from_pretrained(_snake_case       ,  subfolder="""scheduler"""      )
   __lowerCAmelCase ,					__lowerCAmelCase					:	List[Any]      =       FlaxStableDiffusionPipeline.from_pretrained(
       _snake_case       ,  scheduler=_snake_case       ,  revision="""bf16"""       ,  dtype=jnp.bfloataa       ,  )
   __lowerCAmelCase					:	Dict      =       scheduler_params

   __lowerCAmelCase					:	int      =       """A painting of a squirrel eating a burger"""
   __lowerCAmelCase					:	Any      =       jax.device_count()
   __lowerCAmelCase					:	Any      =       num_samples * [prompt]
   __lowerCAmelCase					:	Dict      =       sd_pipe.prepare_inputs(_snake_case      )

   __lowerCAmelCase					:	Any      =       replicate(_snake_case      )
   __lowerCAmelCase					:	List[str]      =       shard(_snake_case      )

   __lowerCAmelCase					:	List[str]      =       jax.random.PRNGKey(0      )
   __lowerCAmelCase					:	Any      =       jax.random.split(_snake_case       ,  jax.device_count()      )

   __lowerCAmelCase					:	int      =       sd_pipe(_snake_case       ,  _snake_case       ,  _snake_case       ,  num_inference_steps=25       ,  jit=_snake_case      )[0]
   assert images.shape == (jax.device_count(), 1, 768, 768, 3)

   __lowerCAmelCase					:	List[str]      =       images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]      )
   __lowerCAmelCase					:	Optional[int]      =       images[0, 253:256, 253:256, -1]

   __lowerCAmelCase					:	List[Any]      =       jnp.asarray(jax.device_get(image_slice.flatten()      )      )
   __lowerCAmelCase					:	Optional[int]      =       jnp.array([0.4_336, 0.42_969, 0.4_453, 0.4_199, 0.4_297, 0.4_531, 0.4_434, 0.4_434, 0.4_297]      )
   print(F'''output_slice: {output_slice}'''      )
   assert jnp.abs(output_slice - expected_slice      ).max() < 1E-2






from math import isqrt






def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:int   )      ->   list[int]:
 __lowerCAmelCase					:	Tuple      =       [True] * max_number
 for i in range(2   ,   isqrt(max_number - 1   ) + 1   ):
  if is_prime[i]:
   for j in range(i**2   ,   SCREAMING_SNAKE_CASE   ,   SCREAMING_SNAKE_CASE   ):
    __lowerCAmelCase					:	Tuple      =       False

 return [i for i in range(2   ,   SCREAMING_SNAKE_CASE   ) if is_prime[i]]






def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:int = 10**8   )      ->   int:
 __lowerCAmelCase					:	int      =       calculate_prime_numbers(max_number // 2   )

 __lowerCAmelCase					:	List[Any]      =       0
 __lowerCAmelCase					:	List[str]      =       0
 __lowerCAmelCase					:	List[Any]      =       len(SCREAMING_SNAKE_CASE   ) - 1
 while left <= right:
  while prime_numbers[left] * prime_numbers[right] >= max_number:
   right -= 1
  semiprimes_count += right - left + 1
  left += 1

 return semiprimes_count


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