

"""simple docstring"""




import math

import flax.linen as nn
import jax.numpy as jnp
def 	lowerCAmelCase    (       __UpperCamelCase				,     __UpperCamelCase				,     __UpperCamelCase = 1				,     __UpperCamelCase = 1				,     __UpperCamelCase = 1.0e4				,     __UpperCamelCase = False				,     __UpperCamelCase = 1.0				,     ):
 '''simple docstring'''





 assert timesteps.ndim == 1, "Timesteps should be a 1d-array"
 assert embedding_dim % 2 == 0, F"Embedding dimension {embedding_dim} should be even"
 UpperCAmelCase__       :	Tuple        =       float(embedding_dim // 2   )
 UpperCAmelCase__       :	List[str]        =       math.log(max_timescale / min_timescale   ) / (num_timescales - freq_shift)
 UpperCAmelCase__       :	Union[str, Any]        =       min_timescale * jnp.exp(jnp.arange(snake_case__				,     dtype=jnp.floataa   ) * -log_timescale_increment   )
 UpperCAmelCase__       :	List[Any]        =       jnp.expand_dims(snake_case__				,     1   ) * jnp.expand_dims(snake_case__				,     0   )

 # scale embeddings
 UpperCAmelCase__       :	List[Any]        =       scale * emb

 if flip_sin_to_cos:
  UpperCAmelCase__       :	str        =       jnp.concatenate([jnp.cos(snake_case__   ), jnp.sin(snake_case__   )]				,     axis=1   )
 else:
  UpperCAmelCase__       :	List[str]        =       jnp.concatenate([jnp.sin(snake_case__   ), jnp.cos(snake_case__   )]				,     axis=1   )
 UpperCAmelCase__       :	Any        =       jnp.reshape(snake_case__				,     [jnp.shape(snake_case__   )[0], embedding_dim]   )
 return signal







class   __lowercase      (     nn.Module       ):
  snake_case_												=							3_2
  snake_case_												=							jnp.floataa

  @nn.compact
  def __call__(	self		:      List[Any]       ,A		:      Dict     ):






   '''simple docstring'''
   UpperCAmelCase__       :	List[str]        =       nn.Dense(self.time_embed_dim       ,dtype=self.dtype       ,name="""linear_1"""     )(a__     )
   UpperCAmelCase__       :	Any        =       nn.silu(a__     )
   UpperCAmelCase__       :	Optional[Any]        =       nn.Dense(self.time_embed_dim       ,dtype=self.dtype       ,name="""linear_2"""     )(a__     )
   return temb







class   __lowercase      (     nn.Module       ):
  snake_case_												=							3_2
  snake_case_												=							False
  snake_case_												=							1

  @nn.compact
  def __call__(	self		:      str       ,A		:      Optional[int]     ):






   '''simple docstring'''
   return get_sinusoidal_embeddings(
       a__       ,embedding_dim=self.dim       ,flip_sin_to_cos=self.flip_sin_to_cos       ,freq_shift=self.freq_shift     )













from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


lowerCamelCase__       		=  logging.get_logger(__name__)

lowerCamelCase__       		=  {
    """microsoft/focalnet-tiny""": """https://huggingface.co/microsoft/focalnet-tiny/resolve/main/config.json""",
}




class       snake_case__ (				lowercase_						,     lowercase_):

       '''simple docstring'''
       lowerCamelCase				:					List[Any]								=			"focalnet"



       def __init__(				self      ,    a__=2_24      ,    a__=4      ,    a__=3      ,    a__=96      ,    a__=False      ,    a__=[1_92, 3_84, 7_68, 7_68]      ,    a__=[2, 2, 6, 2]      ,    a__=[2, 2, 2, 2]      ,    a__=[3, 3, 3, 3]      ,    a__="gelu"      ,    a__=4.0      ,    a__=0.0      ,    a__=0.1      ,    a__=False      ,    a__=1e-4      ,    a__=False      ,    a__=False      ,    a__=False      ,    a__=0.02      ,    a__=1e-5      ,    a__=32      ,    a__=None      ,    a__=None      ,    **a__      ,    )			->   Tuple:






            '''simple docstring'''





            super().__init__(**a__     )

            __snake_case      :Optional[Any] 							=    image_size
            __snake_case      :Dict 							=    patch_size
            __snake_case      :Tuple 							=    num_channels
            __snake_case      :List[str] 							=    embed_dim
            __snake_case      :Any 							=    use_conv_embed
            __snake_case      :List[Any] 							=    hidden_sizes
            __snake_case      :str 							=    depths
            __snake_case      :List[Any] 							=    focal_levels
            __snake_case      :int 							=    focal_windows
            __snake_case      :Any 							=    hidden_act
            __snake_case      :List[str] 							=    mlp_ratio
            __snake_case      :str 							=    hidden_dropout_prob
            __snake_case      :List[str] 							=    drop_path_rate
            __snake_case      :int 							=    use_layerscale
            __snake_case      :int 							=    layerscale_value
            __snake_case      :Tuple 							=    use_post_layernorm
            __snake_case      :Dict 							=    use_post_layernorm_in_modulation
            __snake_case      :Union[str, Any] 							=    normalize_modulator
            __snake_case      :List[str] 							=    initializer_range
            __snake_case      :Optional[Any] 							=    layer_norm_eps
            __snake_case      :Tuple 							=    encoder_stride
            __snake_case      :List[str] 							=    ["""stem"""] + [F'''stage{idx}''' for idx in range(1      ,    len(self.depths     ) + 1     )]
            __snake_case							,			__snake_case      :str 							=    get_aligned_output_features_output_indices(
                out_features=a__      ,    out_indices=a__      ,    stage_names=self.stage_names     )


