
"""simple docstring"""
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       UpperCamelCase   (				unittest.TestCase	):
     def 					__SCREAMING_SNAKE_CASE				(  self					):




       """simple docstring"""


       super().tearDown()
       gc.collect()
     def 					__SCREAMING_SNAKE_CASE				(  self					):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE :					Union[str, Any]				     =			FlaxStableDiffusionPipeline.from_pretrained(
           "stabilityai/stable-diffusion-2"	,  revision="bf16"	,  dtype=jnp.bfloataa	,  )

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

       _SCREAMING_SNAKE_CASE :					List[str]				     =			replicate(__SCREAMING_SNAKE_CASE					)
       _SCREAMING_SNAKE_CASE :					Optional[Any]				     =			shard(__SCREAMING_SNAKE_CASE					)

       _SCREAMING_SNAKE_CASE :					Union[str, Any]				     =			jax.random.PRNGKey(0					)
       _SCREAMING_SNAKE_CASE :					Union[str, Any]				     =			jax.random.split(__SCREAMING_SNAKE_CASE	,  jax.device_count()					)

       _SCREAMING_SNAKE_CASE :					Any				     =			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, 768, 768, 3)

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

       _SCREAMING_SNAKE_CASE :					Tuple				     =			jnp.asarray(jax.device_get(image_slice.flatten()					)					)
       _SCREAMING_SNAKE_CASE :					Tuple				     =			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 					__SCREAMING_SNAKE_CASE				(  self					):




       """simple docstring"""


       _SCREAMING_SNAKE_CASE :					Any				     =			"stabilityai/stable-diffusion-2"
       _SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE :					Optional[Any]				     =			FlaxDPMSolverMultistepScheduler.from_pretrained(__SCREAMING_SNAKE_CASE	,  subfolder="scheduler"					)
       _SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE :					Optional[Any]				     =			FlaxStableDiffusionPipeline.from_pretrained(
           __SCREAMING_SNAKE_CASE	,  scheduler=__SCREAMING_SNAKE_CASE	,  revision="bf16"	,  dtype=jnp.bfloataa	,  )
       _SCREAMING_SNAKE_CASE :					str				     =			scheduler_params

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

       _SCREAMING_SNAKE_CASE :					Dict				     =			replicate(__SCREAMING_SNAKE_CASE					)
       _SCREAMING_SNAKE_CASE :					Any				     =			shard(__SCREAMING_SNAKE_CASE					)

       _SCREAMING_SNAKE_CASE :					int				     =			jax.random.PRNGKey(0					)
       _SCREAMING_SNAKE_CASE :					Union[str, Any]				     =			jax.random.split(__SCREAMING_SNAKE_CASE	,  jax.device_count()					)

       _SCREAMING_SNAKE_CASE :					Any				     =			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, 768, 768, 3)

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

       _SCREAMING_SNAKE_CASE :					Tuple				     =			jnp.asarray(jax.device_get(image_slice.flatten()					)					)
       _SCREAMING_SNAKE_CASE :					List[str]				     =			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





"""simple docstring"""




from itertools import product



def     __magic_name__		(					UpperCamelCase     :		int							,		UpperCamelCase     :		int     )     ->						list[int]:
    a__		=			sides_number
    a__		=			max_face_number * dice_number
    a__		=			[0] * (max_total + 1)

    a__		=			1
    a__		=			range(UpperCamelCase							,		max_face_number + 1     )
    for dice_numbers in product(UpperCamelCase							,		repeat=UpperCamelCase     ):
        a__		=			sum(UpperCamelCase     )
        totals_frequencies[total] += 1

    return totals_frequencies



def     __magic_name__		(					)     ->						float:
    a__		=			total_frequency_distribution(
        sides_number=4							,		dice_number=9     )
    a__		=			total_frequency_distribution(
        sides_number=6							,		dice_number=6     )

    a__		=			0
    a__		=			9
    a__		=			4 * 9
    a__		=			6
    for peter_total in range(UpperCamelCase							,		max_peter_total + 1     ):
        peter_wins_count += peter_totals_frequencies[peter_total] * sum(
            colin_totals_frequencies[min_colin_total:peter_total]     )

    a__		=			(4**9) * (6**6)
    a__		=			peter_wins_count / total_games_number

    a__		=			round(UpperCamelCase							,		ndigits=7     )

    return rounded_peter_win_probability


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


