
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   lowercase__  (       self):


   '''simple docstring'''




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


 def   lowercase__  (       self):


   '''simple docstring'''




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

   lowercase__	: Optional[int]     		=   """A painting of a squirrel eating a burger"""
   lowercase__	: Optional[Any]     		=   jax.device_count()
   lowercase__	: Optional[int]     		=   num_samples * [prompt]
   lowercase__	: Tuple     		=   sd_pipe.prepare_inputs(SCREAMING_SNAKE_CASE_)

   lowercase__	: Union[str, Any]     		=   replicate(SCREAMING_SNAKE_CASE_)
   lowercase__	: Optional[int]     		=   shard(SCREAMING_SNAKE_CASE_)

   lowercase__	: Union[str, Any]     		=   jax.random.PRNGKey(0)
   lowercase__	: Tuple     		=   jax.random.split(SCREAMING_SNAKE_CASE_       ,	jax.device_count())

   lowercase__	: Optional[int]     		=   sd_pipe(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	num_inference_steps=25       ,	jit=SCREAMING_SNAKE_CASE_)[0]
   assert images.shape == (jax.device_count(), 1, 7_68, 7_68, 3)

   lowercase__	: List[str]     		=   images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:])
   lowercase__	: Optional[Any]     		=   images[0, 2_53:2_56, 2_53:2_56, -1]

   lowercase__	: Dict     		=   jnp.asarray(jax.device_get(image_slice.flatten()))
   lowercase__	: Dict     		=   jnp.array([0.4_2_3_8, 0.4_4_1_4, 0.4_3_9_5, 0.4_4_5_3, 0.4_6_2_9, 0.4_5_9_0, 0.4_5_3_1, 0.4_5_5_0_8, 0.4_5_1_2])
   print(f'output_slice: {output_slice}')
   assert jnp.abs(output_slice - expected_slice).max() < 1E-2



 def   lowercase__  (       self):


   '''simple docstring'''




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

   lowercase__	: List[Any]     		=   """A painting of a squirrel eating a burger"""
   lowercase__	: List[Any]     		=   jax.device_count()
   lowercase__	: str     		=   num_samples * [prompt]
   lowercase__	: Optional[Any]     		=   sd_pipe.prepare_inputs(SCREAMING_SNAKE_CASE_)

   lowercase__	: List[str]     		=   replicate(SCREAMING_SNAKE_CASE_)
   lowercase__	: Optional[Any]     		=   shard(SCREAMING_SNAKE_CASE_)

   lowercase__	: str     		=   jax.random.PRNGKey(0)
   lowercase__	: Any     		=   jax.random.split(SCREAMING_SNAKE_CASE_       ,	jax.device_count())

   lowercase__	: Tuple     		=   sd_pipe(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	num_inference_steps=25       ,	jit=SCREAMING_SNAKE_CASE_)[0]
   assert images.shape == (jax.device_count(), 1, 7_68, 7_68, 3)

   lowercase__	: Union[str, Any]     		=   images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:])
   lowercase__	: List[Any]     		=   images[0, 2_53:2_56, 2_53:2_56, -1]

   lowercase__	: List[Any]     		=   jnp.asarray(jax.device_get(image_slice.flatten()))
   lowercase__	: Union[str, Any]     		=   jnp.array([0.4_3_3_6, 0.4_2_9_6_9, 0.4_4_5_3, 0.4_1_9_9, 0.4_2_9_7, 0.4_5_3_1, 0.4_4_3_4, 0.4_4_3_4, 0.4_2_9_7])
   print(f'output_slice: {output_slice}')
   assert jnp.abs(output_slice - expected_slice).max() < 1E-2








lowerCamelCase__     :		dict[tuple[int, int, int], int]        = {}

def 		UpperCamelCase						(   lowercase_			,			lowercase_			,			lowercase_ )		->						int:


  '''simple docstring'''






  if late == 3 or absent == 2:
    return 0

  # if we have no days left, and have not failed any other rules,
  # we have a prize string
  if days == 0:
    return 1

  # No easy solution, so now we need to do the recursive calculation

  # First, check if the combination is already in the cache, and
  # if yes, return the stored value from there since we already
  # know the number of possible prize strings from this point on
  lowercase__	: Tuple     		=   (days, absent, late)
  if key in cache:
    return cache[key]

  # now we calculate the three possible ways that can unfold from
  # this point on, depending on our attendance today

  # 1) if we are late (but not absent), the "absent" counter stays as
  # it is, but the "late" counter increases by one
  lowercase__	: Union[str, Any]     		=   _calculate(days - 1			,			lowercase_			,			late + 1 )

  # 2) if we are absent, the "absent" counter increases by 1, and the
  # "late" counter resets to 0
  lowercase__	: List[str]     		=   _calculate(days - 1			,			absent + 1			,			0 )

  # 3) if we are on time, this resets the "late" counter and keeps the
  # absent counter
  lowercase__	: Dict     		=   _calculate(days - 1			,			lowercase_			,			0 )

  lowercase__	: List[str]     		=   state_late + state_absent + state_ontime

  lowercase__	: List[Any]     		=   prizestrings
  return prizestrings

def 		UpperCamelCase						(   lowercase_ = 30 )		->						int:


  '''simple docstring'''






  return _calculate(lowercase_			,			absent=0			,			late=0 )


if __name__ == "__main__":
     print(solution())






