dsipts.models.chronos2 package

Submodules

dsipts.models.chronos2.config module

class dsipts.models.chronos2.config.Chronos2CoreConfig(d_model: int = 512, d_kv: int = 64, d_ff: int = 2048, num_layers: int = 6, num_heads: int = 8, dropout_rate: float = 0.1, layer_norm_epsilon: float = 1e-06, initializer_factor: float = 0.05, feed_forward_proj: str = 'relu', vocab_size: int = 2, pad_token_id: int = 0, rope_theta: float = 10000.0, attn_implementation: Literal['eager', 'sdpa'] | None = None, **kwargs)

Bases: PretrainedConfig

HF transformers-style pretrained model config for Chronos-2.0, based on T5Config.

Parameters:

• d_model – Size model’s hidden states, by default 512

• d_kv – Size of the key, query, value projections per attention head, by default 64

• d_ff – Size of the intermediate feed forward layers, by default 2048

• num_layers – Number of hidden layers in the encoder, by default 6

• num_heads – Number of attention heads for each attention layer, by default 8

• dropout_rate – The ratio for all dropout layers, by default 0.1

• layer_norm_epsilon – The epsilon used by the layer normalization layers, by default 1e-6

• initializer_factor – A factor for initializing all weight matrices, by default 0.05

• feed_forward_proj – Type of feed forward layer to be used, by default “relu”

• vocab_size – Size of vocabulary for special tokens, by default 2

• pad_token_id – Token ID for padding/missing value token, by default 0

• rope_theta – The base theta for rotary position embedding (RoPE), by default 10000.0

• attn_implementation – The attention implementation to use. Options: “eager” or “sdpa”, by default None (uses “sdpa”)

attribute_map: dict[str, str] = {'head_dim': 'd_kv', 'hidden_size': 'd_model', 'num_attention_heads': 'num_heads', 'num_hidden_layers': 'num_layers'}

model_type: str = 't5'

class dsipts.models.chronos2.config.Chronos2ForecastingConfig(context_length: int, output_patch_size: int, input_patch_size: int, input_patch_stride: int, quantiles: List[float], use_reg_token: bool = False, use_arcsinh: bool = False, max_output_patches: int = 1, time_encoding_scale: int | None = None)

Bases: object

context_length: int

classmethod editable_fields() → list[str]

Fields that maybe modified during the fine-tuning stage.

input_patch_size: int

input_patch_stride: int

max_output_patches: int = 1

output_patch_size: int

quantiles: List[float]

time_encoding_scale: int | None = None

use_arcsinh: bool = False

use_reg_token: bool = False

dsipts.models.chronos2.layers module

class dsipts.models.chronos2.layers.AttentionOutput(hidden_states: torch.Tensor | None = None, attn_weights: torch.Tensor | None = None)

Bases: ModelOutput

attn_weights: Tensor | None = None

hidden_states: Tensor | None = None

class dsipts.models.chronos2.layers.Chronos2Encoder(config: Chronos2CoreConfig)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(inputs_embeds: Tensor, *, group_ids: Tensor, attention_mask: Tensor | None = None, position_ids: Tensor | None = None, output_attentions: bool = False) → Chronos2EncoderOutput

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dsipts.models.chronos2.layers.Chronos2EncoderBlock(config)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(hidden_states: Tensor, *, position_ids: Tensor, attention_mask: Tensor, group_time_mask: Tensor, output_attentions: bool = False) → Chronos2EncoderBlockOutput

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dsipts.models.chronos2.layers.Chronos2EncoderBlockOutput(hidden_states: torch.Tensor | None = None, time_self_attn_weights: torch.Tensor | None = None, group_self_attn_weights: torch.Tensor | None = None)

Bases: ModelOutput

group_self_attn_weights: Tensor | None = None

hidden_states: Tensor | None = None

time_self_attn_weights: Tensor | None = None

class dsipts.models.chronos2.layers.Chronos2EncoderOutput(last_hidden_state: torch.Tensor | None = None, all_time_self_attn_weights: tuple[torch.Tensor, ...] | None = None, all_group_self_attn_weights: tuple[torch.Tensor, ...] | None = None)

Bases: ModelOutput

all_group_self_attn_weights: tuple[Tensor, ...] | None = None

all_time_self_attn_weights: tuple[Tensor, ...] | None = None

last_hidden_state: Tensor | None = None

class dsipts.models.chronos2.layers.Chronos2LayerNorm(hidden_size: int, eps: float = 1e-06)

Bases: Module

Construct a layernorm module in the T5 style. No bias and no subtraction of mean.

forward(hidden_states)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dsipts.models.chronos2.layers.FeedForward(config)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(hidden_states: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

mlp: Module

class dsipts.models.chronos2.layers.GroupSelfAttention(config)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(hidden_states: Tensor, attention_mask: Tensor, output_attentions: bool = False) → AttentionOutput

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dsipts.models.chronos2.layers.InstanceNorm(eps: float = 1e-05, use_arcsinh: bool = False)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor, loc_scale: tuple[Tensor, Tensor] | None = None) → tuple[Tensor, tuple[Tensor, Tensor]]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

inverse(x: Tensor, loc_scale: tuple[Tensor, Tensor]) → Tensor

class dsipts.models.chronos2.layers.MHA(config, use_rope: bool = True)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

d_model: int

dropout: float

forward(hidden_states: Tensor, mask: Tensor, encoder_states: Tensor | None = None, position_ids: Tensor | None = None, output_attentions: bool = False) → AttentionOutput

Multi-head attention forward pass.

Parameters:

• hidden_states – Input tensor of shape [batch_size, seq_len, d_model]

• mask – Attention mask tensor of shape [batch_size, num_heads, q_len, kv_len]

• encoder_states – Encoder states for cross-attention. Defaults to None.

• position_ids – Position IDs for RoPE. Defaults to None.

• output_attentions – Whether to return attention weights. Defaults to False.

Returns:

Contains:

• hidden_states : Output tensor of shape [batch_size, seq_len, d_model]

• attn_weights : Attention weights if output_attentions=True

Return type:

AttentionOutput

inner_dim: int

kv_proj_dim: int

n_heads: int

class dsipts.models.chronos2.layers.MLP(config)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(hidden_states: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dsipts.models.chronos2.layers.Patch(patch_size: int, patch_stride: int)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dsipts.models.chronos2.layers.ResidualBlock(in_dim: int, h_dim: int, out_dim: int, act_fn_name: str, dropout_p: float = 0.0, use_layer_norm: bool = False)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class dsipts.models.chronos2.layers.RoPE(dim: int, base: float = 10000)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

static apply_rotary_pos_emb(q: Tensor, k: Tensor, cos: Tensor, sin: Tensor, unsqueeze_dim: int = 1) → tuple[Tensor, Tensor]

Applies Rotary Position Embedding to the query and key tensors.

Parameters:

• q (torch.Tensor) – The query tensor.

• k (torch.Tensor) – The key tensor.

• cos (torch.Tensor) – The cosine part of the rotary embedding.

• sin (torch.Tensor) – The sine part of the rotary embedding.

• unsqueeze_dim (int, optional, defaults to 1) – The ‘unsqueeze_dim’ argument specifies the dimension along which to unsqueeze cos[position_ids] and sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.

Returns:

tuple(torch.Tensor) comprising of the query and key tensors rotated using the Rotary Position Embedding.

forward(x: Tensor, position_ids: Tensor) → tuple[Tensor, Tensor]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

inv_freq: Tensor

static rotate_half(x)

Rotates half the hidden dims of the input.

class dsipts.models.chronos2.layers.TimeSelfAttention(config)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(hidden_states: Tensor, attention_mask: Tensor, position_ids: Tensor, output_attentions: bool = False) → AttentionOutput

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Module contents