
"""simple docstring"""





from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)






@flax.struct.dataclass
class        _UpperCAmelCase(       lowerCamelCase      ):
		lowercase__         = 42






@flax_register_to_config
class        _UpperCAmelCase(       nn.Module ,	lowerCamelCase ,	lowerCamelCase      ):
		lowercase__         = 32
		lowercase__         = 4
		lowercase__         = 4
		lowercase__         = (
		    "CrossAttnDownBlock2D",
		    "CrossAttnDownBlock2D",
		    "CrossAttnDownBlock2D",
		    "DownBlock2D",
		)
		lowercase__         = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
		lowercase__         = False
		lowercase__         = (3_20, 6_40, 12_80, 12_80)
		lowercase__         = 2
		lowercase__         = 8
		lowercase__         = None
		lowercase__         = 12_80
		lowercase__         = 0.0
		lowercase__         = False
		lowercase__         = jnp.floataa
		lowercase__         = True
		lowercase__         = 0
		lowercase__         = False



		def 		UpperCAmelCase				(   self							,						__a)			->			FrozenDict:






							'''simple docstring'''
							# init input tensors
							_UpperCamelCase									=    (1, self.in_channels, self.sample_size, self.sample_size)
							_UpperCamelCase									=    jnp.zeros(__a							,						dtype=jnp.floataa)
							_UpperCamelCase									=    jnp.ones((1,)							,						dtype=jnp.intaa)
							_UpperCamelCase									=    jnp.zeros((1, 1, self.cross_attention_dim)							,						dtype=jnp.floataa)

							_UpperCamelCase					,    _UpperCamelCase									=    jax.random.split(__a)
							_UpperCamelCase									=    {'''params''': params_rng, '''dropout''': dropout_rng}

							return self.init(__a							,						__a							,						__a							,						__a)["params"]



		def 		UpperCAmelCase				(   self)			->			str:






							'''simple docstring'''
							_UpperCamelCase									=    self.block_out_channels
							_UpperCamelCase									=    block_out_channels[0] * 4

							if self.num_attention_heads is not None:
												raise ValueError(
												    '''At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19.''')

							# If `num_attention_heads` is not defined (which is the case for most models)
							# it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
							# The reason for this behavior is to correct for incorrectly named variables that were introduced
							# when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
							# Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
							# which is why we correct for the naming here.
							_UpperCamelCase									=    self.num_attention_heads or self.attention_head_dim

							# input
							_UpperCamelCase									=    nn.Conv(
							    block_out_channels[0]							,						kernel_size=(3, 3)							,						strides=(1, 1)							,						padding=((1, 1), (1, 1))							,						dtype=self.dtype							,						)

							# time
							_UpperCamelCase									=    FlaxTimesteps(
							    block_out_channels[0]							,						flip_sin_to_cos=self.flip_sin_to_cos							,						freq_shift=self.config.freq_shift)
							_UpperCamelCase									=    FlaxTimestepEmbedding(__a							,						dtype=self.dtype)

							_UpperCamelCase									=    self.only_cross_attention
							if isinstance(__a							,						__a):
												_UpperCamelCase									=    (only_cross_attention,) * len(self.down_block_types)

							if isinstance(__a							,						__a):
												_UpperCamelCase									=    (num_attention_heads,) * len(self.down_block_types)

							# down
							_UpperCamelCase									=    []
							_UpperCamelCase									=    block_out_channels[0]
							for i, down_block_type in enumerate(self.down_block_types):
												_UpperCamelCase									=    output_channel
												_UpperCamelCase									=    block_out_channels[i]
												_UpperCamelCase									=    i == len(__a) - 1

												if down_block_type == "CrossAttnDownBlock2D":
																	_UpperCamelCase									=    FlaxCrossAttnDownBlockaD(
																	    in_channels=__a							,						out_channels=__a							,						dropout=self.dropout							,						num_layers=self.layers_per_block							,						num_attention_heads=num_attention_heads[i]							,						add_downsample=not is_final_block							,						use_linear_projection=self.use_linear_projection							,						only_cross_attention=only_cross_attention[i]							,						use_memory_efficient_attention=self.use_memory_efficient_attention							,						dtype=self.dtype							,						)
												else:
																	_UpperCamelCase									=    FlaxDownBlockaD(
																	    in_channels=__a							,						out_channels=__a							,						dropout=self.dropout							,						num_layers=self.layers_per_block							,						add_downsample=not is_final_block							,						dtype=self.dtype							,						)

												down_blocks.append(__a)
							_UpperCamelCase									=    down_blocks

							# mid
							_UpperCamelCase									=    FlaxUNetMidBlockaDCrossAttn(
							    in_channels=block_out_channels[-1]							,						dropout=self.dropout							,						num_attention_heads=num_attention_heads[-1]							,						use_linear_projection=self.use_linear_projection							,						use_memory_efficient_attention=self.use_memory_efficient_attention							,						dtype=self.dtype							,						)

							# up
							_UpperCamelCase									=    []
							_UpperCamelCase									=    list(reversed(__a))
							_UpperCamelCase									=    list(reversed(__a))
							_UpperCamelCase									=    list(reversed(__a))
							_UpperCamelCase									=    reversed_block_out_channels[0]
							for i, up_block_type in enumerate(self.up_block_types):
												_UpperCamelCase									=    output_channel
												_UpperCamelCase									=    reversed_block_out_channels[i]
												_UpperCamelCase									=    reversed_block_out_channels[min(i + 1							,						len(__a) - 1)]

												_UpperCamelCase									=    i == len(__a) - 1

												if up_block_type == "CrossAttnUpBlock2D":
																	_UpperCamelCase									=    FlaxCrossAttnUpBlockaD(
																	    in_channels=__a							,						out_channels=__a							,						prev_output_channel=__a							,						num_layers=self.layers_per_block + 1							,						num_attention_heads=reversed_num_attention_heads[i]							,						add_upsample=not is_final_block							,						dropout=self.dropout							,						use_linear_projection=self.use_linear_projection							,						only_cross_attention=only_cross_attention[i]							,						use_memory_efficient_attention=self.use_memory_efficient_attention							,						dtype=self.dtype							,						)
												else:
																	_UpperCamelCase									=    FlaxUpBlockaD(
																	    in_channels=__a							,						out_channels=__a							,						prev_output_channel=__a							,						num_layers=self.layers_per_block + 1							,						add_upsample=not is_final_block							,						dropout=self.dropout							,						dtype=self.dtype							,						)

												up_blocks.append(__a)
												_UpperCamelCase									=    output_channel
							_UpperCamelCase									=    up_blocks

							# out
							_UpperCamelCase									=    nn.GroupNorm(num_groups=32							,						epsilon=1e-5)
							_UpperCamelCase									=    nn.Conv(
							    self.out_channels							,						kernel_size=(3, 3)							,						strides=(1, 1)							,						padding=((1, 1), (1, 1))							,						dtype=self.dtype							,						)




		def __call__(   self							,						__a							,						__a							,						__a							,						__a=None							,						__a=None							,						__a = True							,						__a = False							,						)			->			Union[FlaxUNetaDConditionOutput, Tuple]:






							'''simple docstring'''
							# 1. time
							if not isinstance(__a							,						jnp.ndarray):
												_UpperCamelCase									=    jnp.array([timesteps]							,						dtype=jnp.intaa)
							elif isinstance(__a							,						jnp.ndarray) and len(timesteps.shape) == 0:
												_UpperCamelCase									=    timesteps.astype(dtype=jnp.floataa)
												_UpperCamelCase									=    jnp.expand_dims(__a							,						0)

							_UpperCamelCase									=    self.time_proj(__a)
							_UpperCamelCase									=    self.time_embedding(__a)

							# 2. pre-process
							_UpperCamelCase									=    jnp.transpose(__a							,						(0, 2, 3, 1))
							_UpperCamelCase									=    self.conv_in(__a)

							# 3. down
							_UpperCamelCase									=    (sample,)
							for down_block in self.down_blocks:
												if isinstance(__a							,						__a):
																	_UpperCamelCase					,    _UpperCamelCase									=    down_block(__a							,						__a							,						__a							,						deterministic=not train)
												else:
																	_UpperCamelCase					,    _UpperCamelCase									=    down_block(__a							,						__a							,						deterministic=not train)
												down_block_res_samples += res_samples

							if down_block_additional_residuals is not None:
												_UpperCamelCase									=    ()

												for down_block_res_sample, down_block_additional_residual in zip(
												    __a							,						__a):
																	down_block_res_sample += down_block_additional_residual
																	new_down_block_res_samples += (down_block_res_sample,)

												_UpperCamelCase									=    new_down_block_res_samples

							# 4. mid
							_UpperCamelCase									=    self.mid_block(__a							,						__a							,						__a							,						deterministic=not train)

							if mid_block_additional_residual is not None:
												sample += mid_block_additional_residual

							# 5. up
							for up_block in self.up_blocks:
												_UpperCamelCase									=    down_block_res_samples[-(self.layers_per_block + 1) :]
												_UpperCamelCase									=    down_block_res_samples[: -(self.layers_per_block + 1)]
												if isinstance(__a							,						__a):
																	_UpperCamelCase									=    up_block(
																	    __a							,						temb=__a							,						encoder_hidden_states=__a							,						res_hidden_states_tuple=__a							,						deterministic=not train							,						)
												else:
																	_UpperCamelCase									=    up_block(__a							,						temb=__a							,						res_hidden_states_tuple=__a							,						deterministic=not train)

        # 6. post-process
							_UpperCamelCase									=    self.conv_norm_out(__a)
							_UpperCamelCase									=    nn.silu(__a)
							_UpperCamelCase									=    self.conv_out(__a)
							_UpperCamelCase									=    jnp.transpose(__a							,						(0, 3, 1, 2))

							if not return_dict:
												return (sample,)

							return FlaxUNetaDConditionOutput(sample=__a)








"""simple docstring"""





from __future__ import annotations

from functools import lru_cache
from math import ceil

_a								=						100

_a								=						set(range(3, NUM_PRIMES, 2))
primes.add(2)
_a								=						42

for prime in range(3, ceil(NUM_PRIMES**0.5), 2):
				if prime not in primes:
								continue
				primes.difference_update(set(range(prime * prime, NUM_PRIMES, prime)))



@lru_cache(maxsize=1_00      )
def  lowerCamelCase__						(     __snake_case      )				->						set[int]:


					"""simple docstring"""





					if number_to_partition < 0:
										return set()
					elif number_to_partition == 0:
										return {1}

					_UpperCamelCase									=    set()
					_UpperCamelCase									=    42
					_UpperCamelCase									=    42

					for prime in primes:
										if prime > number_to_partition:
															continue
										for sub in partition(number_to_partition - prime      ):
															ret.add(sub * prime      )

					return ret



def  lowerCamelCase__						(     __snake_case = 50_00      )				->						int | None:


					"""simple docstring"""





					for number_to_partition in range(1,	__snake_case      ):
										if len(partition(__snake_case      )      ) > number_unique_partitions:
															return number_to_partition
					return None


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






