"""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      ):
				A__ :  jnp.ndarray






@flax_register_to_config
class   __UpperCAmelCase		(      nn.Module , __lowerCAmelCase , __lowerCAmelCase      ):
				A__ :  int    				=			3_2
				A__ :  int    				=			4
				A__ :  int    				=			4
				A__ :  Tuple[str]    				=			(
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "DownBlock2D",
				)
				A__ :  Tuple[str]    				=			("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
				A__ :  Union[bool, Tuple[bool]]    				=			False
				A__ :  Tuple[int]    				=			(3_2_0, 6_4_0, 1_2_8_0, 1_2_8_0)
				A__ :  int    				=			2
				A__ :  Union[int, Tuple[int]]    				=			8
				A__ :  Optional[Union[int, Tuple[int]]]    				=			None
				A__ :  int    				=			1_2_8_0
				A__ :  float    				=			0.0
				A__ :  bool    				=			False
				A__ :  jnp.dtype    				=			jnp.floataa
				A__ :  bool    				=			True
				A__ :  int    				=			0
				A__ :  bool    				=			False

				def 						_a      ( self      ,					_lowerCamelCase ):
										# init input tensors
										lowerCamelCase__   =(1, self.in_channels, self.sample_size, self.sample_size)
										lowerCamelCase__   =jnp.zeros(_lowerCamelCase      ,					dtype=jnp.floataa )
										lowerCamelCase__   =jnp.ones((1,)      ,					dtype=jnp.intaa )
										lowerCamelCase__   =jnp.zeros((1, 1, self.cross_attention_dim)      ,					dtype=jnp.floataa )

										lowerCamelCase__   , lowerCamelCase__   =jax.random.split(_lowerCamelCase )
										lowerCamelCase__   ={"params": params_rng, "dropout": dropout_rng}

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

				def 						_a      ( self ):
										lowerCamelCase__   =self.block_out_channels
										lowerCamelCase__   =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.
										lowerCamelCase__   =self.num_attention_heads or self.attention_head_dim

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

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

										lowerCamelCase__   =self.only_cross_attention
										if isinstance(_lowerCamelCase      ,					_lowerCamelCase ):
																lowerCamelCase__   =(only_cross_attention,) * len(self.down_block_types )

										if isinstance(_lowerCamelCase      ,					_lowerCamelCase ):
																lowerCamelCase__   =(num_attention_heads,) * len(self.down_block_types )

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

																if down_block_type == "CrossAttnDownBlock2D":
																						lowerCamelCase__   =FlaxCrossAttnDownBlockaD(
																						    in_channels=_lowerCamelCase      ,					out_channels=_lowerCamelCase      ,					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:
																						lowerCamelCase__   =FlaxDownBlockaD(
																						    in_channels=_lowerCamelCase      ,					out_channels=_lowerCamelCase      ,					dropout=self.dropout      ,					num_layers=self.layers_per_block      ,					add_downsample=not is_final_block      ,					dtype=self.dtype      ,					)

																down_blocks.append(_lowerCamelCase )
										lowerCamelCase__   =down_blocks

										# mid
										lowerCamelCase__   =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
										lowerCamelCase__   =[]
										lowerCamelCase__   =list(reversed(_lowerCamelCase ) )
										lowerCamelCase__   =list(reversed(_lowerCamelCase ) )
										lowerCamelCase__   =list(reversed(_lowerCamelCase ) )
										lowerCamelCase__   =reversed_block_out_channels[0]
										for i, up_block_type in enumerate(self.up_block_types ):
																lowerCamelCase__   =output_channel
																lowerCamelCase__   =reversed_block_out_channels[i]
																lowerCamelCase__   =reversed_block_out_channels[min(i + 1      ,					len(_lowerCamelCase ) - 1 )]

																lowerCamelCase__   =i == len(_lowerCamelCase ) - 1

																if up_block_type == "CrossAttnUpBlock2D":
																						lowerCamelCase__   =FlaxCrossAttnUpBlockaD(
																						    in_channels=_lowerCamelCase      ,					out_channels=_lowerCamelCase      ,					prev_output_channel=_lowerCamelCase      ,					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:
																						lowerCamelCase__   =FlaxUpBlockaD(
																						    in_channels=_lowerCamelCase      ,					out_channels=_lowerCamelCase      ,					prev_output_channel=_lowerCamelCase      ,					num_layers=self.layers_per_block + 1      ,					add_upsample=not is_final_block      ,					dropout=self.dropout      ,					dtype=self.dtype      ,					)

																up_blocks.append(_lowerCamelCase )
																lowerCamelCase__   =output_channel
										lowerCamelCase__   =up_blocks

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






				def __call__( self      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase=None      ,					_lowerCamelCase=None      ,					_lowerCamelCase = True      ,					_lowerCamelCase = False      ,					):
										# 1. time
										if not isinstance(_lowerCamelCase      ,					jnp.ndarray ):
																lowerCamelCase__   =jnp.array([timesteps]      ,					dtype=jnp.intaa )
										elif isinstance(_lowerCamelCase      ,					jnp.ndarray ) and len(timesteps.shape ) == 0:
																lowerCamelCase__   =timesteps.astype(dtype=jnp.floataa )
																lowerCamelCase__   =jnp.expand_dims(_lowerCamelCase      ,					0 )

										lowerCamelCase__   =self.time_proj(_lowerCamelCase )
										lowerCamelCase__   =self.time_embedding(_lowerCamelCase )

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

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

										if down_block_additional_residuals is not None:
																lowerCamelCase__   =()

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

																lowerCamelCase__   =new_down_block_res_samples

										# 4. mid
										lowerCamelCase__   =self.mid_block(_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					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:
																lowerCamelCase__   =down_block_res_samples[-(self.layers_per_block + 1) :]
																lowerCamelCase__   =down_block_res_samples[: -(self.layers_per_block + 1)]
																if isinstance(_lowerCamelCase      ,					_lowerCamelCase ):
																						lowerCamelCase__   =up_block(
																						    _lowerCamelCase      ,					temb=_lowerCamelCase      ,					encoder_hidden_states=_lowerCamelCase      ,					res_hidden_states_tuple=_lowerCamelCase      ,					deterministic=not train      ,					)
																else:
																						lowerCamelCase__   =up_block(_lowerCamelCase      ,					temb=_lowerCamelCase      ,					res_hidden_states_tuple=_lowerCamelCase      ,					deterministic=not train )

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

										if not return_dict:
																return (sample,)

										return FlaxUNetaDConditionOutput(sample=_lowerCamelCase )



"""simple docstring"""


def      lowerCamelCase_					(   __lowerCAmelCase							)				->   "list[int]":




						'''simple docstring'''



						if upper_limit < 0:
												raise ValueError("Limit for the Catalan sequence must be ≥ 0"							)

						lowerCamelCase__   =[0] * (upper_limit + 1)

						# Base case: C(0) = C(1) = 1
						lowerCamelCase__   =1
						if upper_limit > 0:
												lowerCamelCase__   =1

						# Recurrence relation: C(i) = sum(C(j).C(i-j-1)), from j = 0 to i
						for i in range(2  ,       upper_limit + 1							):
												for j in range(__lowerCAmelCase							):
																		catalan_list[i] += catalan_list[j] * catalan_list[i - j - 1]

						return catalan_list


if __name__ == "__main__":
	print('\n********* Catalan Numbers Using Dynamic Programming ************\n')
	print('\n*** Enter -1 at any time to quit ***')
	print('\nEnter the upper limit (≥ 0) for the Catalan number sequence: ', end='')
	try:
		while True:
			a				=int(input().strip())
			if N < 0:
				print('\n********* Goodbye!! ************')
				break
			else:
				print(F'''The Catalan numbers from 0 through {N} are:''')
				print(catalan_numbers(N))
				print('Try another upper limit for the sequence: ', end='')
	except (NameError, ValueError):
		print('\n********* Invalid input, goodbye! ************\n')

	import doctest

	doctest.testmod()


