



'''simple docstring'''






from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
from .unet_ad_blocks import get_down_block, get_mid_block, get_out_block, get_up_block

@dataclass
class        SCREAMING_SNAKE_CASE(							_snake_case    ):







 """simple docstring"""






 lowerCamelCase__    				= 42

class        SCREAMING_SNAKE_CASE(							_snake_case							,	_snake_case    ):







 """simple docstring"""






 @register_to_config
 def __init__(		self    :  Any	,    __snake_case    :  Any = 65536	,    __snake_case    :  Optional[Any] = None	,    __snake_case    :  List[str] = 2	,    __snake_case    :  Any = 2	,    __snake_case    :  Optional[Any] = 0	,    __snake_case    :  str = "fourier"	,    __snake_case    :  Optional[int] = True	,    __snake_case    :  List[Any] = False	,    __snake_case    :  Any = 0.0	,    __snake_case    :  List[str] = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D")	,    __snake_case    :  str = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip")	,    __snake_case    :  str = "UNetMidBlock1D"	,    __snake_case    :  Tuple = None	,    __snake_case    :  Dict = (32, 32, 64)	,    __snake_case    :  Dict = None	,    __snake_case    :  Union[str, Any] = 8	,    __snake_case    :  List[Any] = 1	,    __snake_case    :  int = False	,    )    ->	List[str]:
        super().__init__()
        UpperCAmelCase				: Optional[Any]											=	sample_size

        # time
        if time_embedding_type == "fourier":
               UpperCAmelCase				: List[Any]											=	GaussianFourierProjection(
                   embedding_size=8	,    set_W_to_weight=snake_case_	,    log=snake_case_	,    flip_sin_to_cos=snake_case_	)
               UpperCAmelCase				: int											=	2 * block_out_channels[0]
        elif time_embedding_type == "positional":
               UpperCAmelCase				: Optional[Any]											=	Timesteps(
                   block_out_channels[0]	,    flip_sin_to_cos=snake_case_	,    downscale_freq_shift=snake_case_	)
               UpperCAmelCase				: List[Any]											=	block_out_channels[0]

        if use_timestep_embedding:
               UpperCAmelCase				: Dict											=	block_out_channels[0] * 4
               UpperCAmelCase				: str											=	TimestepEmbedding(
                   in_channels=snake_case_	,    time_embed_dim=snake_case_	,    act_fn=snake_case_	,    out_dim=block_out_channels[0]	,    )

        UpperCAmelCase				: int											=	nn.ModuleList([]	)
        UpperCAmelCase				: Any											=	None
        UpperCAmelCase				: str											=	nn.ModuleList([]	)
        UpperCAmelCase				: Tuple											=	None

        # down
        UpperCAmelCase				: str											=	in_channels
        for i, down_block_type in enumerate(snake_case_	):
               UpperCAmelCase				: Optional[int]											=	output_channel
               UpperCAmelCase				: List[Any]											=	block_out_channels[i]

               if i == 0:
                      input_channel += extra_in_channels

               UpperCAmelCase				: Union[str, Any]											=	i == len(snake_case_	) - 1

               UpperCAmelCase				: Optional[Any]											=	get_down_block(
                   snake_case_	,    num_layers=snake_case_	,    in_channels=snake_case_	,    out_channels=snake_case_	,    temb_channels=block_out_channels[0]	,    add_downsample=not is_final_block or downsample_each_block	,    )
               self.down_blocks.append(snake_case_	)

        # mid
        UpperCAmelCase				: Any											=	get_mid_block(
            snake_case_	,    in_channels=block_out_channels[-1]	,    mid_channels=block_out_channels[-1]	,    out_channels=block_out_channels[-1]	,    embed_dim=block_out_channels[0]	,    num_layers=snake_case_	,    add_downsample=snake_case_	,    )

        # up
        UpperCAmelCase				: int											=	list(reversed(snake_case_	)	)
        UpperCAmelCase				: List[Any]											=	reversed_block_out_channels[0]
        if out_block_type is None:
               UpperCAmelCase				: Optional[Any]											=	out_channels
        else:
               UpperCAmelCase				: List[str]											=	block_out_channels[0]

        for i, up_block_type in enumerate(snake_case_	):
               UpperCAmelCase				: Any											=	output_channel
               UpperCAmelCase				: Tuple											=	(
                   reversed_block_out_channels[i + 1] if i < len(snake_case_	) - 1 else final_upsample_channels
               )

               UpperCAmelCase				: str											=	i == len(snake_case_	) - 1

               UpperCAmelCase				: Union[str, Any]											=	get_up_block(
                   snake_case_	,    num_layers=snake_case_	,    in_channels=snake_case_	,    out_channels=snake_case_	,    temb_channels=block_out_channels[0]	,    add_upsample=not is_final_block	,    )
               self.up_blocks.append(snake_case_	)
               UpperCAmelCase				: Tuple											=	output_channel

        # out
        UpperCAmelCase				: List[str]											=	norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4	,    32	)
        UpperCAmelCase				: Any											=	get_out_block(
            out_block_type=snake_case_	,    num_groups_out=snake_case_	,    embed_dim=block_out_channels[0]	,    out_channels=snake_case_	,    act_fn=snake_case_	,    fc_dim=block_out_channels[-1] // 4	,    )







 def   A			(		self    :  List[str]	,    __snake_case    :  str	,    __snake_case    :  Optional[Any]	,    __snake_case    :  Any = True	,    )    ->	str:
        UpperCAmelCase				: Optional[int]											=	timestep
        if not torch.is_tensor(snake_case_	):
               UpperCAmelCase				: List[Any]											=	torch.tensor([timesteps]	,    dtype=torch.long	,    device=sample.device	)
        elif torch.is_tensor(snake_case_	) and len(timesteps.shape	) == 0:
               UpperCAmelCase				: str											=	timesteps[None].to(sample.device	)

        UpperCAmelCase				: Tuple											=	self.time_proj(snake_case_	)
        if self.config.use_timestep_embedding:
               UpperCAmelCase				: Any											=	self.time_mlp(snake_case_	)
        else:
               UpperCAmelCase				: Tuple											=	timestep_embed[..., None]
               UpperCAmelCase				: Any											=	timestep_embed.repeat([1, 1, sample.shape[2]]	).to(sample.dtype	)
               UpperCAmelCase				: int											=	timestep_embed.broadcast_to((sample.shape[:1] + timestep_embed.shape[1:])	)

        # 2. down
        UpperCAmelCase				: Optional[Any]											=	()
        for downsample_block in self.down_blocks:
               UpperCAmelCase				: str											=	downsample_block(hidden_states=snake_case_	,    temb=snake_case_	)
               down_block_res_samples += res_samples

        # 3. mid
        if self.mid_block:
               UpperCAmelCase				: List[str]											=	self.mid_block(snake_case_	,    snake_case_	)

        # 4. up
        for i, upsample_block in enumerate(self.up_blocks	):
               UpperCAmelCase				: Any											=	down_block_res_samples[-1:]
               UpperCAmelCase				: Dict											=	down_block_res_samples[:-1]
               UpperCAmelCase				: Dict											=	upsample_block(snake_case_	,    res_hidden_states_tuple=snake_case_	,    temb=snake_case_	)

        # 5. post-process
        if self.out_block:
               UpperCAmelCase				: List[Any]											=	self.out_block(snake_case_	,    snake_case_	)

        if not return_dict:
               return (sample,)

        return UNetaDOutput(sample=snake_case_	)










'''simple docstring'''






from typing import List, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase__:     Optional[Any]										=    logging.get_logger(__name__)

UpperCamelCase__:     Tuple										=    {
    "huggingface/time-series-transformer-tourism-monthly": (
        "https://huggingface.co/huggingface/time-series-transformer-tourism-monthly/resolve/main/config.json"
    ),
    # See all TimeSeriesTransformer models at https://huggingface.co/models?filter=time_series_transformer
}

class        SCREAMING_SNAKE_CASE(							A__    ):







	"""simple docstring"""






	lowerCamelCase__    				= """time_series_transformer"""
	lowerCamelCase__    				= {
	    """hidden_size""": """d_model""",
	    """num_attention_heads""": """encoder_attention_heads""",
	    """num_hidden_layers""": """encoder_layers""",
	}







	def __init__(		self    :  str	,    __snake_case    :  Optional[int] = None	,    __snake_case    :  Optional[int] = None	,    __snake_case    :  str = "student_t"	,    __snake_case    :  str = "nll"	,    __snake_case    :  int = 1	,    __snake_case    :  List[int] = [1, 2, 3, 4, 5, 6, 7]	,    __snake_case    :  Optional[Union[str, bool]] = "mean"	,    __snake_case    :  int = 0	,    __snake_case    :  int = 0	,    __snake_case    :  int = 0	,    __snake_case    :  int = 0	,    __snake_case    :  Optional[List[int]] = None	,    __snake_case    :  Optional[List[int]] = None	,    __snake_case    :  int = 32	,    __snake_case    :  int = 32	,    __snake_case    :  int = 2	,    __snake_case    :  int = 2	,    __snake_case    :  int = 2	,    __snake_case    :  int = 2	,    __snake_case    :  bool = True	,    __snake_case    :  str = "gelu"	,    __snake_case    :  int = 64	,    __snake_case    :  float = 0.1	,    __snake_case    :  float = 0.1	,    __snake_case    :  float = 0.1	,    __snake_case    :  float = 0.1	,    __snake_case    :  float = 0.1	,    __snake_case    :  int = 100	,    __snake_case    :  float = 0.02	,    __snake_case    :  Optional[Any]=True	,    **__snake_case    :  List[Any]	,    )    ->	List[str]:
								# time series specific configuration
								UpperCAmelCase				: List[Any]											=	prediction_length
								UpperCAmelCase				: List[Any]											=	context_length or prediction_length
								UpperCAmelCase				: Tuple											=	distribution_output
								UpperCAmelCase				: Optional[Any]											=	loss
								UpperCAmelCase				: Tuple											=	input_size
								UpperCAmelCase				: Optional[int]											=	num_time_features
								UpperCAmelCase				: Dict											=	lags_sequence
								UpperCAmelCase				: Any											=	scaling
								UpperCAmelCase				: Tuple											=	num_dynamic_real_features
								UpperCAmelCase				: Any											=	num_static_real_features
								UpperCAmelCase				: Optional[int]											=	num_static_categorical_features
								if cardinality and num_static_categorical_features > 0:
															if len(__snake_case	) != num_static_categorical_features:
																						raise ValueError(
																						    '''The cardinality should be a list of the same length as `num_static_categorical_features`'''	)
															UpperCAmelCase				: Any											=	cardinality
								else:
															UpperCAmelCase				: Optional[Any]											=	[0]
								if embedding_dimension and num_static_categorical_features > 0:
															if len(__snake_case	) != num_static_categorical_features:
																						raise ValueError(
																						    '''The embedding dimension should be a list of the same length as `num_static_categorical_features`'''	)
															UpperCAmelCase				: Optional[int]											=	embedding_dimension
								else:
															UpperCAmelCase				: Optional[Any]											=	[min(50	,    (cat + 1) // 2	) for cat in self.cardinality]
								UpperCAmelCase				: List[Any]											=	num_parallel_samples

								# Transformer architecture configuration
								UpperCAmelCase				: int											=	input_size * len(__snake_case	) + self._number_of_features
								UpperCAmelCase				: int											=	d_model
								UpperCAmelCase				: str											=	encoder_attention_heads
								UpperCAmelCase				: str											=	decoder_attention_heads
								UpperCAmelCase				: List[str]											=	encoder_ffn_dim
								UpperCAmelCase				: Any											=	decoder_ffn_dim
								UpperCAmelCase				: Any											=	encoder_layers
								UpperCAmelCase				: str											=	decoder_layers

								UpperCAmelCase				: Optional[int]											=	dropout
								UpperCAmelCase				: Union[str, Any]											=	attention_dropout
								UpperCAmelCase				: List[Any]											=	activation_dropout
								UpperCAmelCase				: Optional[int]											=	encoder_layerdrop
								UpperCAmelCase				: Dict											=	decoder_layerdrop

								UpperCAmelCase				: Tuple											=	activation_function
								UpperCAmelCase				: Dict											=	init_std

								UpperCAmelCase				: Optional[int]											=	use_cache

								super().__init__(is_encoder_decoder=__snake_case	,    **__snake_case	)







	@property
	def   A			(		self    :  Union[str, Any]	)    ->	int:
								return (
								    sum(self.embedding_dimension	)
								    + self.num_dynamic_real_features
								    + self.num_time_features
								    + self.num_static_real_features
								    + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
								)





