| """TinyCast model assembly. |
| |
| model(past_values, scale_factor, prediction_length, batch_first).quantile_outputs |
| """ |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from .config import TinyCastConfig |
|
|
|
|
| @dataclass |
| class PredictionOutput: |
| quantile_outputs: torch.Tensor |
| prediction_outputs: torch.Tensor = None |
|
|
|
|
| class TinyCastBackbone(nn.Module): |
| """Inner model: per-window min-max norm -> dilated-conv core -> denorm.""" |
|
|
| def __init__(self, config: TinyCastConfig): |
| super().__init__() |
| self.config = config |
| from .backbone import DilatedConvBackbone |
| from .normalization import WindowMinMax |
|
|
| self.core = DilatedConvBackbone( |
| seq_len=int(config.seq_len), |
| p_out=int(config.output_token_len), |
| n_quantiles=int(config.num_quantiles), |
| d=int(config.conv_dim), |
| n_layers=int(config.n_layers), |
| kernel=int(config.kernel_size), |
| ffn_mult=float(config.ffn_mult), |
| top_k_periods=int(config.top_k_periods), |
| significance_alpha=float(config.significance_alpha), |
| n_harmonics=int(config.n_harmonics), |
| pool_kind=str(config.pool_kind), |
| causal=bool(config.causal), |
| phase_bins=int(config.phase_bins), |
| decoder_depth=int(config.decoder_depth), |
| separable_conv=bool(config.separable_conv), |
| share_ffn=bool(config.share_ffn), |
| future_conv=bool(config.future_conv), |
| future_conv_layers=int(config.future_conv_layers), |
| future_conv_seed=int(config.future_conv_seed), |
| ) |
| self.norm = WindowMinMax(eps_clamp=1e-5) |
|
|
| def encode( |
| self, |
| past_values: torch.Tensor, |
| batch_first: bool = True, |
| scale_factor: "torch.Tensor | float | None" = None, |
| horizon: "int | None" = None, |
| ) -> "tuple[torch.Tensor, torch.Tensor, torch.Tensor]": |
| """per-window min-max norm -> dilated-conv core -> normalized y.""" |
| x = past_values |
| if not batch_first: |
| x = x.transpose(0, 1) |
| if x.dim() == 2: |
| x = x.unsqueeze(-1) |
|
|
| x_normed, x_min, x_range = self.norm.transform(x) |
| nan_mask = (~torch.isnan(x)).to(x.dtype) |
| y_norm = self.core( |
| x_normed, nan_mask=nan_mask, scale_factor=scale_factor, |
| horizon=horizon, |
| ) |
| return y_norm, x_min, x_range |
|
|
|
|
| class TinyCastForPrediction(nn.Module): |
| """Top-level model container. |
| |
| Usage: |
| model = TinyCastForPrediction(TinyCastConfig()) |
| out = model(past_values=x, scale_factor=sf, prediction_length=pl, |
| batch_first=False) |
| quantiles = out.quantile_outputs # (B, Q, pred_len, 1) |
| |
| NOTE: this ``forward`` is the SINGLE-SHOT arbitrary-horizon path. The |
| deployed GIFT-Eval inference uses the AR-rollout predictor |
| (``tinycast.predictor.TinyCastPredictor``), which calls ``self.model.encode`` |
| in 48-step chunks. For horizons <= 48 the two paths coincide. |
| """ |
|
|
| def __init__(self, config: TinyCastConfig): |
| super().__init__() |
| self.config = config |
| self.model = TinyCastBackbone(config) |
|
|
| def forward(self, past_values, scale_factor=None, prediction_length=None, |
| batch_first=None): |
| if scale_factor is None: |
| scale_factor = 1.0 |
| if batch_first is None: |
| batch_first = True |
| ctx = past_values |
| if not batch_first: |
| ctx = ctx.transpose(0, 1) |
| if ctx.dim() == 2: |
| ctx = ctx.unsqueeze(-1) |
| ctx_len_max = int(self.config.seq_len) |
| ctx_in = ctx[:, -ctx_len_max:, :] |
| p_out_native = int(self.model.core.p_out) |
| pred_len = int(prediction_length) if prediction_length else p_out_native |
| y_norm, x_min, x_range = self.model.encode( |
| ctx_in, batch_first=True, scale_factor=scale_factor, |
| horizon=pred_len, |
| ) |
| if y_norm.dim() == 2: |
| y_norm = y_norm.unsqueeze(-1) |
| y_pred = y_norm * x_range + x_min |
| quantile_outputs = y_pred.permute(0, 2, 1).unsqueeze(-1) |
| return PredictionOutput(quantile_outputs=quantile_outputs) |
|
|
