From cf1129506a395b550da5a27e7939060e1c184bc9 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 13:46:44 +0000 Subject: [PATCH 1/9] feat(t0): add t0-alpha and t0-beta inference Port The Forecasting Company's t0 family from tfc-t0 0.6.0 as Tsfm.Forecasting.T0 (Apache-2.0): causal scaler, patch encoder, alternating time/group transformer with xPos rotary embedding, residual decoder and monotone quantile head. Each checkpoint's scaler convention is read from its config; t0-beta's differs from alpha's. T0Forecaster takes the TimesFM shape (ForecastJoint with known-future flags, Forecast over OHLCVA). Both now implement a shared IQuantileForecaster in the core package, so one call site drives either. Unlike TimesFM, any context length is accepted and up to 1024 steps are decoded in one pass. Verified stage by stage against the reference on both checkpoints, including the quickstart notebook's electricity-price and M4 examples. Weights are fetched at pinned, checksummed revisions; both exceed the nuget.org size limit, so neither is packaged. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01F1dTUhvq2UnFWdBmH7AA8G --- .github/workflows/ci.yml | 5 + .gitignore | 4 + CHANGELOG.md | 36 +++ NOTICE | 36 +++ README.md | 87 ++++- Tsfm.Forecasting.slnx | 3 + reference/dump_t0_parity.py | 156 +++++++++ .../GroceryCovariates.csproj | 1 + samples/GroceryCovariates/Program.cs | 53 +++- scripts/fetch-t0-checkpoint.sh | 41 +++ src/Tsfm.Forecasting.T0/T0Forecaster.cs | 234 ++++++++++++++ src/Tsfm.Forecasting.T0/T0Model.cs | 173 ++++++++++ src/Tsfm.Forecasting.T0/T0Modules.cs | 296 ++++++++++++++++++ src/Tsfm.Forecasting.T0/T0Preprocess.cs | 155 +++++++++ .../Tsfm.Forecasting.T0.csproj | 24 ++ .../TimesFmForecaster.cs | 4 + src/Tsfm.Forecasting/IQuantileForecaster.cs | 40 +++ tests/Tsfm.Forecasting.T0.Parity/Program.cs | 154 +++++++++ .../Tsfm.Forecasting.T0.Parity.csproj | 16 + .../CheckpointTests.cs | 97 ++++++ .../InvariantTests.cs | 215 +++++++++++++ .../Tsfm.Forecasting.T0.Tests.csproj | 19 ++ 22 files changed, 1827 insertions(+), 22 deletions(-) create mode 100644 reference/dump_t0_parity.py create mode 100755 scripts/fetch-t0-checkpoint.sh create mode 100644 src/Tsfm.Forecasting.T0/T0Forecaster.cs create mode 100644 src/Tsfm.Forecasting.T0/T0Model.cs create mode 100644 src/Tsfm.Forecasting.T0/T0Modules.cs create mode 100644 src/Tsfm.Forecasting.T0/T0Preprocess.cs create mode 100644 src/Tsfm.Forecasting.T0/Tsfm.Forecasting.T0.csproj create mode 100644 src/Tsfm.Forecasting/IQuantileForecaster.cs create mode 100644 tests/Tsfm.Forecasting.T0.Parity/Program.cs create mode 100644 tests/Tsfm.Forecasting.T0.Parity/Tsfm.Forecasting.T0.Parity.csproj create mode 100644 tests/Tsfm.Forecasting.T0.Tests/CheckpointTests.cs create mode 100644 tests/Tsfm.Forecasting.T0.Tests/InvariantTests.cs create mode 100644 tests/Tsfm.Forecasting.T0.Tests/Tsfm.Forecasting.T0.Tests.csproj diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 45ff753..1f5affd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -49,6 +49,11 @@ jobs: - name: TimesFM tests (self-skipping without weights) run: dotnet test tests/Tsfm.Forecasting.TimesFm.IntegrationTests -c Release --nologo + # The invariant tests run on a tiny random-init model and need no weights; the + # checkpoint tests self-skip, as the ~1.4 GB of t0 weights are not fetched here. + - name: t0 tests + run: dotnet test tests/Tsfm.Forecasting.T0.Tests -c Release --nologo + # Packing embeds the checkpoints, so this job needs them. Platform-neutral. pack: runs-on: ubuntu-latest diff --git a/.gitignore b/.gitignore index 009ee29..fcc80fa 100644 --- a/.gitignore +++ b/.gitignore @@ -20,3 +20,7 @@ reference/parity.safetensors # Python environment for the TimesFM parity dump. .venv/ + +# t0 parity fixtures and the public datasets they are built from. +reference/*-parity.safetensors +reference/data/ diff --git a/CHANGELOG.md b/CHANGELOG.md index 89af825..be56638 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,41 @@ # Changelog +## Unreleased + +### Tsfm.Forecasting.T0 (Apache-2.0) — new + +- Inference for The Forecasting Company's **t0-alpha** (102M) and **t0-beta** (256M), ported + from `tfc-t0` 0.6.0: causal scaler, patch encoder, alternating time/group transformer + with xPos rotary embedding, residual decoder, monotone quantile head. +- `T0Forecaster` takes the TimesFM shape: `ForecastJoint` over variate-major series with + known-future flags and a target index, and `Forecast` over OHLCVA for forward returns. + Unlike TimesFM, any context length is accepted (left-padded onto the patch grid and + masked), and up to 1024 steps are decoded in one pass. +- Each checkpoint's scaler convention is read from its `config.json`. t0-beta's differs + from alpha's, and running it under alpha's loads cleanly and forecasts worse. +- Verified stage by stage against the reference, both checkpoints, including the + quickstart notebook's electricity-price and M4 examples (`tests/Tsfm.Forecasting.T0.Parity`, + `reference/dump_t0_parity.py`). +- Weights are Apache-2.0 but not packaged: each exceeds nuget.org's 250 MB limit. + `scripts/fetch-t0-checkpoint.sh` fetches pinned revisions and checks their SHA-256. +- Not ported: autoregressive rollout past 1024 steps, quantile interpolation and tail + extrapolation to levels off the trained grid, and the optional symmetric averaging, + clipping and integer snapping. + +### Tsfm.Forecasting + +- `IQuantileForecaster`: the shared `ForecastJoint` / `Forecast` / `Quantiles` shape. + `TimesFmForecaster` and `T0Forecaster` implement it, so one call site drives either. + Additive; TimesFM behaviour is unchanged. + +### Samples and testing + +- `samples/GroceryCovariates` runs on TimesFM or t0 (`-- t0-beta`) through the interface. +- `Tsfm.Forecasting.T0.Tests` checks structural invariants on a tiny random-init model, so + they run in CI without weights: a target never reads its own future, known-future rows + never read targets, time attention is causal, any context length is accepted. Behaviour + on the published checkpoints is covered too, gated on their presence. + ## 0.1.3 — 2026-09-02 ### Joint forecasting with known-future covariates diff --git a/NOTICE b/NOTICE index 15d1dd1..6db2ace 100644 --- a/NOTICE +++ b/NOTICE @@ -41,3 +41,39 @@ TimesFM WEIGHTS ARE NOT DISTRIBUTED AND CANNOT BE. The published checkpoints are under the TimesFM Non-Commercial License v1.0 — research and evaluation only, no production or revenue-generating use. Obtaining them is the user's own act under Google's terms, and nothing here grants any right in them. + +──────────────────────────────────────────────────────────────────────────────── +Tsfm.Forecasting.T0 — Apache-2.0, NOT MIT +──────────────────────────────────────────────────────────────────────────────── + +That package contains software derived from tfc-t0 +(https://github.com/theforecastingcompany/tfc-t0), Copyright 2026 The Forecasting +Company, licensed under the Apache License, Version 2.0. Its NOTICE carries these +attributions, reproduced here: + + This product includes software developed at Datadog (https://www.datadoghq.com/), + Copyright 2025 Datadog, Inc., licensed under the Apache License 2.0. + Source: https://github.com/DataDog/toto + + This product includes software derived from Chronos-2 by Amazon.com, Inc. or + its affiliates, Copyright Amazon.com, Inc. All Rights Reserved, licensed under + the Apache License 2.0. + Source: https://github.com/amazon-science/chronos-forecasting + +Changes to the derived work, as required by Apache-2.0 section 4(b): + + - The inference path is translated from Python/PyTorch to C#/TorchSharp: the + causal scaler, patch alignment, attention-mask construction, the patch + encoder, the time/group transformer stack with xPos rotary embedding, the + residual decoder and the monotone quantile head. + - Rotary embedding is computed directly rather than through + rotary-embedding-torch, with the same interleaved pairing and xPos scale. + - Attention returns zero for a query with no permitted key, explicitly rather + than through the reference's safe softmax. + - Training, fine-tuning, the MLX runtime, autoregressive rollout beyond one + pass, quantile interpolation and tail extrapolation, and the optional + symmetric-averaging, clipping and snapping adjustments are not ported. + +No upstream source files and no weights are redistributed. The t0-alpha and +t0-beta checkpoints are Apache-2.0, published by The Forecasting Company, and +are fetched by scripts/fetch-t0-checkpoint.sh at pinned revisions. diff --git a/README.md b/README.md index 7b0f10c..73345b3 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,8 @@ # Tsfm.Forecasting Native .NET inference for time-series foundation models, on CPU or GPU via -TorchSharp. Currently [Kronos](https://github.com/shiyu-coder/Kronos) and [TimesFM 3.0](https://github.com/google-research/timesfm). +TorchSharp. Currently [Kronos](https://github.com/shiyu-coder/Kronos), [TimesFM 3.0](https://github.com/google-research/timesfm) +and [t0](https://huggingface.co/theforecastingcompany/t0-beta) (t0-alpha, t0-beta). [![ci](https://github.com/markhammond/tsfm-forecasting/actions/workflows/ci.yml/badge.svg)](https://github.com/markhammond/tsfm-forecasting/actions/workflows/ci.yml) @@ -16,12 +17,13 @@ NuGet packages: | Tsfm.Forecasting.Kronos.Weights.Small [![NuGet](https://img.shields.io/nuget/v/Tsfm.Forecasting.Kronos.Weights.Small.svg)](https://www.nuget.org/packages/Tsfm.Forecasting.Kronos.Weights.Small) | MIT | Kronos-small (24.7M) + Tokenizer-base | | Tsfm.Forecasting.Kronos.Weights.Mini [![NuGet](https://img.shields.io/nuget/v/Tsfm.Forecasting.Kronos.Weights.Mini.svg)](https://www.nuget.org/packages/Tsfm.Forecasting.Kronos.Weights.Mini) | MIT | Kronos-mini (4.1M) + Tokenizer-2k | | Tsfm.Forecasting.TimesFm [![NuGet](https://img.shields.io/nuget/v/Tsfm.Forecasting.TimesFm.svg)](https://www.nuget.org/packages/Tsfm.Forecasting.TimesFm) | **Apache-2.0** | TimesFM 3.0 — **no weights** | +| Tsfm.Forecasting.T0 [![NuGet](https://img.shields.io/nuget/v/Tsfm.Forecasting.T0.svg)](https://www.nuget.org/packages/Tsfm.Forecasting.T0) | **Apache-2.0** | t0-alpha, t0-beta — open weights, fetched separately | Refer to respective package for license and conditions of use (if any). ## Why this exists -Neither model is a text LLM, so existing .NET inference runtimes cannot load them. +None of these models is a text LLM, so existing .NET inference runtimes cannot load them. Kronos pairs a Binary Spherical Quantization autoencoder with a decoder that has a hierarchical two-subtoken embedding, a dependency-aware cross-attention layer and two conditional heads. GGUF has no representation for any of that, which rules out @@ -103,6 +105,54 @@ costs no more than a point forecast. Every OHLCVA channel is supplied as a varia marked a target: preprocessing masks future covariate slots only for target variates, so a channel left non-target would be handed its own future values. +## t0 — open weights, fetched + +[t0](https://arxiv.org/abs/2609.24559) from The Forecasting Company is a decoder-style patch +transformer alternating causal time attention with variate (group) attention, with native +known-future covariates. Both checkpoints are **Apache-2.0**. They are fetched rather than +packaged only because each exceeds nuget.org's 250 MB limit. + +| Checkpoint | Params | Quantiles | Scaler eps | +|---|---|---|---| +| t0-alpha | 102M | 5 (0.1 … 0.9) | 0.1, variance offset | +| t0-beta | 256M | 21 (0.01 … 0.99) | 0.01, std clamp | + +``` +dotnet add package Tsfm.Forecasting.T0 +dotnet add package TorchSharp-cpu +./scripts/fetch-t0-checkpoint.sh # both; or name one: t0-alpha | t0-beta +``` + +```csharp +using Tsfm.Forecasting.T0; + +var t0 = T0Forecaster.Load("checkpoints/t0-beta", new Device("mps")); + +// Same shape as TimesFM: variate-major series, known-future flags, target index. +double[,] q = t0.ForecastJoint(series, variates: 3, contextBars: 512, horizon: 24, + knownFuture: [false, true, true]); // [step, quantile] +double[,] r = t0.Forecast(ohlcva, horizon: 4); // forward returns +``` + +Both forecasters implement `IQuantileForecaster`, so one call site drives either; +`Quantiles` names the columns, whose count differs per checkpoint. `samples/GroceryCovariates` +runs on either (`-- t0-beta`). + +What differs from TimesFM: + +- **Any context length.** The context is left-padded onto the 32-step patch grid and the + padding masked, so nothing need fill whole patches. No ceiling is imposed; the reference + evaluates at 8192. +- **Up to 1024 steps per pass.** Every horizon patch is decoded at once. The reference + continues past 1024 by an autoregressive rollout over quantile paths, which is not ported. +- **Known-future series cannot read the targets.** Group attention lets a known-future row + attend only to other known-future rows, so information flows from covariates into + targets and never back. They are scaled by global rather than causal statistics. +- **Each checkpoint names its own scaler convention.** t0-beta's `config.json` carries + `scaler_eps` and `scaler_eps_mode`; alpha's predates them, so their absence means alpha's + values. Running beta under alpha's convention loads cleanly and silently degrades the + forecast — which is why they are read rather than assumed. + ## Choose your own backend Neither model package propagates a native runtime, because the right one is platform- @@ -128,10 +178,10 @@ by placing a tensor rather than trusting the name. the NuGet package is never used. Without it, loading fails with a message claiming the backend reference is missing — which it is not. -**Neither model has a key-value cache.** Attention is recomputed over the whole context -at every decode step, matching both references. A cache would change the arithmetic and +**No model has a key-value cache.** Attention is recomputed over the whole context +at every decode step, matching the references. A cache would change the arithmetic and forfeit checkable parity. For Kronos this costs `horizon` full passes per rollout, so -long horizons are expensive; TimesFM emits its whole horizon in one pass and is +long horizons are expensive; TimesFM and t0 emit their whole horizon in one pass and are unaffected. ## Parity @@ -156,6 +206,19 @@ and let every position attend to its own future; and the reference leaves `rescale_logits` false, making its logits `QK^T * sqrt(d)` rather than the conventional `QK^T / sqrt(d)`. +**t0.** Every stage — scaling, attention masks, patch embedding, each layer, the +quantile head, the rescale — agrees with `tfc-t0` 0.6.0 to float32 rounding: masks and +scaling are bit-exact, layers within 6e-7 relative, forecasts within 5e-7. Each stage is +fed the reference's previous output, so a fault cannot hide behind an earlier one. The +end-to-end cases include the quickstart notebook's own examples: on the electricity-price +series, t0-beta's median misses the truth by 6.522 without covariates and 2.067 with them +(notebook: 6.528 without, recorded on different hardware), and the port reproduces +the reference to the same digits. Three details would each have produced plausible +forecasts if missed: rotary pairs are interleaved where TimesFM's are split in halves; the +QK norm precedes the rotation, where TimesFM's follows it; and a padding patch has no key +to attend to, so attention must return zero there, not NaN, or every row reading it is +poisoned. + ```bash ./scripts/fetch-checkpoints.sh # Kronos dotnet test tests/Tsfm.Forecasting.Kronos.Tests -c Release @@ -165,6 +228,12 @@ dotnet test tests/Tsfm.Forecasting.Kronos.Tests -c Release python3 -m venv .venv && ./.venv/bin/pip install torch safetensors numpy huggingface_hub ./.venv/bin/python reference/dump_timesfm_parity.py dotnet run --project tests/Tsfm.Forecasting.TimesFm.Parity -c Release + +./scripts/fetch-t0-checkpoint.sh # t0, Apache-2.0 +./.venv/bin/pip install "tfc-t0[torch]==0.6.0" pandas pyarrow +./.venv/bin/python reference/dump_t0_parity.py checkpoints/t0-alpha +./.venv/bin/python reference/dump_t0_parity.py checkpoints/t0-beta +dotnet run --project tests/Tsfm.Forecasting.T0.Parity -c Release -- t0-alpha t0-beta ``` ## Attribution @@ -179,5 +248,11 @@ packages, which record the revision each was built from. Google LLC, Apache-2.0. No Google source files and no weights are redistributed here; both are fetched at pinned revisions by the scripts above. -This project is an independent port and is not affiliated with or endorsed by either. +**t0.** Derived from [tfc-t0](https://github.com/theforecastingcompany/tfc-t0), +Copyright 2026 The Forecasting Company, Apache-2.0, which in turn builds on Datadog's +[Toto](https://github.com/DataDog/toto) and Amazon's +[Chronos-2](https://github.com/amazon-science/chronos-forecasting), both Apache-2.0. No +weights or upstream sources are redistributed; checkpoints are fetched at pinned revisions. + +This project is an independent port and is not affiliated with or endorsed by any of them. See NOTICE for the full attribution and the statement of changes. diff --git a/Tsfm.Forecasting.slnx b/Tsfm.Forecasting.slnx index f22472b..b8645c4 100644 --- a/Tsfm.Forecasting.slnx +++ b/Tsfm.Forecasting.slnx @@ -4,12 +4,15 @@ + + + diff --git a/reference/dump_t0_parity.py b/reference/dump_t0_parity.py new file mode 100644 index 0000000..e425823 --- /dev/null +++ b/reference/dump_t0_parity.py @@ -0,0 +1,156 @@ +"""Dump reference tensors so the .NET t0 port can be bisected stage by stage. + +Runs the published tfc-t0 package (pinned below) against a local checkpoint and records +every stage boundary of one forward pass — scaling, patch embedding, selected layers, the +head, the rescale — plus end-to-end `predict` outputs. A single end-to-end comparison +cannot tell an interleaved-vs-halves RoPE from a misplaced norm; the first diverging stage +can. + +The end-to-end cases include the t0 quickstart notebook's electricity-price and M4-hourly +examples, so the port reproduces the published sample, not just synthetic input. + + pip install "tfc-t0[torch]==0.6.0" pandas pyarrow + python reference/dump_t0_parity.py checkpoints/t0-alpha + python reference/dump_t0_parity.py checkpoints/t0-beta +""" + +import importlib.metadata +import pathlib +import sys +import urllib.request + +import numpy as np +import torch +from safetensors.torch import save_file + +from t0 import T0Forecaster, TimeSeries +from t0.mask import build_attention_masks +from t0.rollout import prepare_rollout_buffer +from t0.types import round_up + +ROOT = pathlib.Path(__file__).resolve().parent.parent +EXPECTED_VERSION = "0.6.0" + +# Not a multiple of the patch: exercises left padding, and a horizon that ends inside its +# last patch, leaving known-future rows PAD there. +T, H = 150, 40 +DATA = "https://autogluon.s3.amazonaws.com/datasets/timeseries" + + +def synthetic() -> tuple[torch.Tensor, torch.Tensor]: + """Two jointly forecast targets with distinct scale and drift, plus two known-future + covariates; one NaN in each kind so the missing-value path is exercised.""" + g = torch.Generator().manual_seed(20261019) + t = torch.arange(T + H, dtype=torch.float32) + ctx = torch.stack([ + 100.0 + 5.0 * torch.sin(t[:T] / 7.0) + 0.02 * t[:T], + -3.0 + 0.25 * torch.cos(t[:T] / 5.0), + ]) + 0.05 * torch.randn(2, T, generator=g) + fut = torch.stack([ + 20.0 + 2.0 * torch.sin(t / 7.0 + 0.3), + (t % 24 < 6).float(), + ]) + ctx[0, 40] = float("nan") + fut[0, 90] = float("nan") + return ctx, fut + + +def stages(model: T0Forecaster, ctx: torch.Tensor, fut: torch.Tensor) -> dict[str, torch.Tensor]: + p = model.patch_size + series = TimeSeries.from_array(ctx[None], fut[None]) + buffer = prepare_rollout_buffer(series, H, T, p) + window = buffer.time_slice(0, round_up(T, p) + round_up(H, p)) + out = { + "buf_values": window.variates, "buf_mask": window.mask.float(), + "buf_groups": window.group_ids.float(), "buf_types": window.variate_type.float(), + } + scaled, ls = model.scaler.scale_input(window) + out |= {"scaled": scaled.variates, "loc": ls.loc, "scale": ls.scale} + + patch = model.patcher.patch + values, mask, vtype, groups = (patch(f) for f in (scaled.variates, scaled.mask, scaled.variate_type, scaled.group_ids)) + h = model.patch_encoder(values, mask, vtype) + out["patch_embedding"] = h + time_mask, group_mask = build_attention_masks(groups, vtype, mask) + out |= {"time_mask": time_mask.float(), "group_mask": group_mask.float()} + for i, layer in enumerate(model.transformer.layers): + h = layer(h, time_attn_mask=time_mask, group_attn_mask=group_mask) + if i in (0, 1, 2, len(model.transformer.layers) - 1): + out[f"layer{i}_output"] = h + h = model.transformer.out_norm(h) + out["transformer_output"] = h + decoded = model.decoder(h) + out["decoder_output"] = decoded + q = model.head(decoded.reshape(*decoded.shape[:-1], p, model.head.n_quantiles)) + out["head_output"] = q + out["rescaled"] = model.scaler.rescale_predictions(q, ls, p) + return out + + +def fetch(name: str) -> pathlib.Path: + cache = ROOT / "reference" / "data" + cache.mkdir(exist_ok=True) + path = cache / name.replace("/", "_") + if not path.exists(): + urllib.request.urlretrieve(f"{DATA}/{name}", path) + return path + + +def notebook_cases() -> dict[str, torch.Tensor]: + """The quickstart notebook's inputs, exactly as it builds them.""" + import pandas as pd + + m4 = pd.read_csv(fetch("m4_hourly/train.csv")) + sid = m4["item_id"].iloc[0] + full = m4.loc[m4["item_id"] == sid, "target"].to_numpy(dtype=np.float32) + train = pd.read_parquet(fetch("electricity_price/train.parquet")) + test = pd.read_parquet(fetch("electricity_price/test.parquet")) + cov_cols = [c for c in test.columns if c not in ("id", "timestamp", "target")] + cov = np.concatenate( + [train[cov_cols].to_numpy(dtype=np.float32)[-512:].T, test[cov_cols].to_numpy(dtype=np.float32).T], axis=1 + ) + return { + "m4_ctx": torch.as_tensor(full[:512]), + "m4_truth": torch.as_tensor(full[512:560]), + "ep_ctx": torch.as_tensor(train["target"].to_numpy(dtype=np.float32)[-512:]), + "ep_cov": torch.as_tensor(cov), + "ep_truth": torch.as_tensor(test["target"].to_numpy(dtype=np.float32)), + } + + +def main() -> None: + installed = importlib.metadata.version("tfc-t0") + if installed != EXPECTED_VERSION: + print(f"warning: tfc-t0 {installed} installed; the port was verified against {EXPECTED_VERSION}") + ckpt = pathlib.Path(sys.argv[1]).resolve() + model = T0Forecaster.from_pretrained(str(ckpt)).eval() + levels = list(model.config.quantile_levels) + torch.set_grad_enabled(False) + + ctx, fut = synthetic() + out = {"in_ctx": ctx, "in_fut": fut} + out |= stages(model, ctx, fut) + # End to end through predict, at the trained levels so no interpolation intervenes. + out["predict_joint"] = model.predict(ctx[None], H, levels, future_covariates=fut[None]).quantiles[0] + out["predict_targets_only"] = model.predict(ctx[None], H, levels).quantiles[0] + # A univariate context that is already patch-aligned, with no padding at all. + out["predict_aligned"] = model.predict(ctx[0, 64:128][None], 32, levels).quantiles[0] + + nb = notebook_cases() + out |= nb + out["m4_predict"] = model.predict(nb["m4_ctx"][None], 48, levels).quantiles[0] + out["ep_predict"] = model.predict(nb["ep_ctx"][None], 24, levels).quantiles[0] + out["ep_predict_cov"] = model.predict( + nb["ep_ctx"][None], 24, levels, future_covariates=nb["ep_cov"][None]).quantiles[0] + mid = levels.index(0.5) + for name in ("ep_predict", "ep_predict_cov"): + err = (out[name][:, mid] - nb["ep_truth"]).abs().mean().item() + print(f" {name:16s} |median-truth| = {err:.4f}") + + dest = ROOT / "reference" / f"{ckpt.name}-parity.safetensors" + save_file({k: v.contiguous().float() for k, v in out.items()}, str(dest)) + print(f" wrote {dest} ({len(out)} tensors)") + + +if __name__ == "__main__": + main() diff --git a/samples/GroceryCovariates/GroceryCovariates.csproj b/samples/GroceryCovariates/GroceryCovariates.csproj index f63bca3..2623283 100644 --- a/samples/GroceryCovariates/GroceryCovariates.csproj +++ b/samples/GroceryCovariates/GroceryCovariates.csproj @@ -12,6 +12,7 @@ + diff --git a/samples/GroceryCovariates/Program.cs b/samples/GroceryCovariates/Program.cs index 861845b..dcb1a84 100644 --- a/samples/GroceryCovariates/Program.cs +++ b/samples/GroceryCovariates/Program.cs @@ -7,10 +7,16 @@ // residuals. TimesFM 3.0 instead attends over the covariates as variates, so the // mechanism differs and the outputs should not be expected to match. // -// dotnet run --project samples/GroceryCovariates +// Either model runs through the same IQuantileForecaster call; only the layout differs, +// because TimesFM reads the known future by rolling whole patches and t0 does not. +// +// dotnet run --project samples/GroceryCovariates # TimesFM 3.0 +// dotnet run --project samples/GroceryCovariates -- t0-beta # or t0-alpha using TorchSharp; +using Tsfm.Forecasting; +using Tsfm.Forecasting.T0; using Tsfm.Forecasting.TimesFm; using static TorchSharp.torch; @@ -33,21 +39,33 @@ catch { device = new Device(DeviceType.CPU); } var root = Path.GetFullPath(Path.Combine(AppContext.BaseDirectory, "../../../../..")); -var forecaster = TimesFmForecaster.Load( - Path.Combine(root, "checkpoints", "timesfm-3.0-pytorch"), device); +var model = args.Length > 0 ? args[0] : "timesfm-3.0-pytorch"; +var checkpoint = Path.Combine(root, "checkpoints", model); -// The model patches 32 steps at a time and reads the known future by rolling whole -// patches forward, so the context must fill one patch and the horizon must sit in the -// next. Seven days is far below that granularity: the observed week is left-padded and -// masked to fill a patch, and the forecast week occupies the one after it. -var patch = forecaster.Config.InputPatchLen; -var context = patch; // one whole patch of context -var total = context + patch; // one further patch holds the horizon -var pad = context - observed; // left padding, masked as absent +IQuantileForecaster forecaster; +int context, total, pad; +if (model.StartsWith("t0")) +{ + // t0 left-pads a short context onto its patch grid itself, and masks the padding. + forecaster = T0Forecaster.Load(checkpoint, device); + (context, total, pad) = (observed, observed + horizon, 0); + Console.WriteLine($"{model} on {device.type}: {observed} days observed, {horizon} forecast\n"); +} +else +{ + // TimesFM patches 32 steps at a time and reads the known future by rolling whole + // patches forward, so the context must fill one patch and the horizon must sit in the + // next. Seven days is far below that granularity: the observed week is left-padded to + // fill a patch, and the forecast week occupies the one after it. + var timesFm = TimesFmForecaster.Load(checkpoint, device); + var patch = timesFm.Config.InputPatchLen; + (forecaster, context, total, pad) = (timesFm, patch, 2 * patch, patch - observed); + Console.WriteLine($"{model} on {device.type}: {observed} days observed, {horizon} forecast; " + + $"padded to {context} context + {patch} horizon patch (patch length {patch})\n"); +} -var quantiles = forecaster.Config.Quantiles; -Console.WriteLine($"{device.type}: {observed} days observed, {horizon} forecast; " - + $"padded to {context} context + {patch} horizon patch (patch length {patch})\n"); +var quantiles = forecaster.Quantiles; +var (lo, mid, hi) = (Level(0.1), Level(0.5), Level(0.9)); foreach (var (name, sales, promotion) in products) { @@ -68,10 +86,13 @@ knownFuture: [false, true, true, true], targetVariate: 0); Console.WriteLine($"{name} (observed: {string.Join(", ", sales)})"); - Console.WriteLine($" day promo temp p10 p50 p90"); + Console.WriteLine($" day promo temp p{quantiles[lo] * 100:F0} p{quantiles[mid] * 100:F0} p{quantiles[hi] * 100:F0}"); for (var h = 0; h < horizon; h++) Console.WriteLine($" {h + 1,3} {(promotion[observed + h] > 0 ? "yes" : " -"),5}" + $" {temperature[observed + h],4:F1}" - + $" {q[h, 0],6:F1} {q[h, quantiles.Length / 2],6:F1} {q[h, quantiles.Length - 1],6:F1}"); + + $" {q[h, lo],6:F1} {q[h, mid],6:F1} {q[h, hi],6:F1}"); Console.WriteLine(); } + +// The column nearest a level: the grids differ (TimesFM 9, t0-alpha 5, t0-beta 21 levels). +int Level(double level) => Enumerable.Range(0, quantiles.Count).MinBy(i => Math.Abs(quantiles[i] - level)); diff --git a/scripts/fetch-t0-checkpoint.sh b/scripts/fetch-t0-checkpoint.sh new file mode 100755 index 0000000..4323809 --- /dev/null +++ b/scripts/fetch-t0-checkpoint.sh @@ -0,0 +1,41 @@ +#!/usr/bin/env bash +# Fetch the t0 checkpoints: t0-alpha (~407 MB) and t0-beta (~1.0 GB). +# +# The weights are Apache-2.0, published by The Forecasting Company. They are not bundled +# into a NuGet package only because both exceed nuget.org's 250 MB limit. +# +# Revisions are pinned rather than tracking main, and each payload is checked against its +# LFS SHA-256: a silently refreshed checkpoint would change every downstream forecast while +# leaving the consuming code identical, and a truncated one loads until a read runs off +# the end. +# +# ./scripts/fetch-t0-checkpoint.sh # both +# ./scripts/fetch-t0-checkpoint.sh t0-alpha # one +set -euo pipefail + +ROOT="$(cd "$(dirname "$0")/.." && pwd)/checkpoints" +# macOS ships shasum, not sha256sum. +if command -v sha256sum > /dev/null; then SHA256="sha256sum"; else SHA256="shasum -a 256"; fi + +fetch () { + local name=$1 rev=$2 sha=$3 + local dir="$ROOT/$name" + local base="https://huggingface.co/theforecastingcompany/$name/resolve/$rev" + mkdir -p "$dir" + for f in config.json LICENSE NOTICE README.md model.safetensors; do + if [ -s "$dir/$f" ]; then echo " have $name/$f"; continue; fi + echo " fetch $name/$f" + curl -fsSL --retry 3 -o "$dir/$f.part" "$base/$f" + mv "$dir/$f.part" "$dir/$f" + done + echo "$sha $dir/model.safetensors" | $SHA256 -c --quiet - \ + || { echo " $name/model.safetensors does not match the pinned revision $rev"; exit 1; } + echo " ok $name @ ${rev:0:12}" +} + +want="${1:-all}" +[[ $want == all || $want == t0-alpha ]] && fetch t0-alpha \ + 741927847fe02367624749bb9290b2f09f128e7d 16c030d3fd70f06dc4238e9a8356e9b5a631d07f80f1bc76ba539991aed5897f +[[ $want == all || $want == t0-beta ]] && fetch t0-beta \ + 941e522547ec63f0f2aad57089fa0f5b8498bdf8 a0fd8abd51275dd30afd21f4892c763ebd45e2cb1811eff26eed43f89a543f7d +exit 0 diff --git a/src/Tsfm.Forecasting.T0/T0Forecaster.cs b/src/Tsfm.Forecasting.T0/T0Forecaster.cs new file mode 100644 index 0000000..686679e --- /dev/null +++ b/src/Tsfm.Forecasting.T0/T0Forecaster.cs @@ -0,0 +1,234 @@ +using static TorchSharp.torch; + +namespace Tsfm.Forecasting.T0; + +/// +/// Quantile forecasts from a t0 checkpoint, in the same shape as +/// TimesFmForecaster: a joint forecast over related series, any of which may be +/// known into the future, and a forward-return forecast from OHLCVA bars. +/// +/// Unlike TimesFM, nothing need fill whole patches: a context of any length is +/// left-padded onto the patch grid, and the padding is masked out of attention. +/// +/// Every series not marked known is a forecast target, and targets are forecast +/// jointly — so a channel cannot be handed its own future. Known-future series are +/// conditioned on but not forecast, and by construction cannot read the targets: their +/// group attention reads only other known-future rows. +/// +/// Inference is not thread safe. +public sealed class T0Forecaster(T0Model model, T0Config config, Device device) : IQuantileForecaster +{ + /// Longest horizon decoded in one pass. The reference continues beyond it by an + /// autoregressive rollout over quantile paths, which is not ported. + public const int MaxHorizon = 1024; + + private const int CloseChannel = 3; // open, high, low, close, volume, amount + private const int Channels = 6; + + public T0Config Config => config; + + public T0Model Model => model; + + /// + public IReadOnlyList Quantiles => config.Quantiles; + + public static T0Forecaster Load(string checkpointDir, Device device) + { + var net = T0Model.Load(checkpointDir, device); + return new T0Forecaster(net, net.Config, device); + } + + /// + /// Joint forecast over several series, any of which may be known into the future. + /// + /// Series are scaled independently — causally for targets, globally for + /// known-future series — so a price in the thousands beside a {0,1} flag needs no + /// manual scaling. There is no variate identity or ordering; the model infers each + /// series' relevance from the supplied window alone. + /// + /// Variate-major [variates x total], so series v at + /// step t is series[v * total + t] where total = series.Length / + /// variates. Steps from on are read only for known + /// variates, and only the first of them. NaN is a missing + /// observation. + /// Series supplied. + /// Observed steps, any positive number. + /// Steps ahead to return, 1... + /// Per variate: true where its future region holds real values + /// the model may read. Empty means none. A variate marked known is NOT a forecast + /// target. + /// Which series the returned quantiles describe. + /// [horizon, Quantiles.Count] in the units of the target series. + public double[,] ForecastJoint( + ReadOnlySpan series, int variates, int contextBars, int horizon, + ReadOnlySpan knownFuture = default, int targetVariate = 0) + { + if (variates < 1) + throw new ArgumentOutOfRangeException(nameof(variates), "at least one series is required"); + if (series.Length % variates != 0) + throw new ArgumentException( + $"series length {series.Length} is not divisible by {variates} variates", nameof(series)); + var total = series.Length / variates; + ValidateHorizon(horizon); + if (contextBars < 1 || contextBars > total) + throw new ArgumentOutOfRangeException(nameof(contextBars), + $"contextBars must lie in 1..{total}; got {contextBars}"); + if (targetVariate < 0 || targetVariate >= variates) + throw new ArgumentOutOfRangeException(nameof(targetVariate)); + if (!knownFuture.IsEmpty && knownFuture.Length != variates) + throw new ArgumentException( + $"knownFuture must hold {variates} flags or be empty", nameof(knownFuture)); + if (!knownFuture.IsEmpty && knownFuture[targetVariate]) + throw new ArgumentException( + "the target variate cannot be known into the future", nameof(knownFuture)); + var anyKnown = !knownFuture.IsEmpty && knownFuture.Contains(true); + if (anyKnown && horizon > total - contextBars) + throw new ArgumentOutOfRangeException(nameof(horizon), + $"horizon {horizon} exceeds the {total - contextBars} known steps supplied after the context"); + + // Targets first, then known-future rows: the reference's row order. Attention is + // indifferent to it, but summation order is not, and parity is checked bit-close. + var order = new int[variates]; + var known = new bool[variates]; + var rows = 0; + for (var v = 0; v < variates; v++) + if (knownFuture.IsEmpty || !knownFuture[v]) order[rows++] = v; + for (var v = 0; v < variates; v++) + if (!knownFuture.IsEmpty && knownFuture[v]) { known[rows] = true; order[rows++] = v; } + + var layout = new Layout(config.PatchSize, variates, contextBars, horizon); + var targetRow = Array.IndexOf(order, targetVariate); + for (var r = 0; r < variates; r++) + { + var src = series.Slice(order[r] * total, total); + layout.Row(r, src[..contextBars], known[r] ? src.Slice(contextBars, horizon) : default); + } + + var head = Run(layout, targetRow); + var result = new double[horizon, config.NumQuantiles]; + for (var h = 0; h < horizon; h++) + for (var q = 0; q < config.NumQuantiles; q++) + result[h, q] = head[h * config.NumQuantiles + q]; + return result; + } + + /// + /// Forward-return quantiles for each step up to . + /// + /// Row-major [L x 6]: open, high, low, close, volume, + /// amount. Any L. Every channel is a target, forecast jointly, so none is handed its + /// own future values. + /// Steps ahead to return, 1.., all from + /// one forward pass. + /// [horizon, numQuantiles] of cumulative return relative to the anchor close. + public double[,] Forecast(ReadOnlySpan ohlcva, int horizon) + { + if (ohlcva.Length == 0 || ohlcva.Length % Channels != 0) + throw new ArgumentException($"ohlcva must hold whole bars of {Channels} channels", nameof(ohlcva)); + ValidateHorizon(horizon); + var bars = ohlcva.Length / Channels; + + var layout = new Layout(config.PatchSize, Channels, bars, horizon); + var column = new float[bars]; + for (var c = 0; c < Channels; c++) + { + for (var t = 0; t < bars; t++) column[t] = ohlcva[t * Channels + c]; + layout.Row(c, column, default); + } + + var head = Run(layout, CloseChannel); + var anchorClose = ohlcva[(bars - 1) * Channels + CloseChannel]; + var result = new double[horizon, config.NumQuantiles]; + for (var h = 0; h < horizon; h++) + for (var q = 0; q < config.NumQuantiles; q++) + result[h, q] = head[h * config.NumQuantiles + q] / anchorClose - 1.0; + return result; + } + + private static void ValidateHorizon(int horizon) + { + if (horizon < 1 || horizon > MaxHorizon) + throw new ArgumentOutOfRangeException(nameof(horizon), $"horizon must lie in 1..{MaxHorizon}"); + } + + /// One forward pass; the target row's first horizon steps as + /// [horizon * quantiles], in data units. + private float[] Run(Layout layout, int targetRow) + { + using var scope = NewDisposeScope(); + using var _ = no_grad(); + + long[] dims = [layout.Rows, layout.Width]; + var values = tensor(layout.Values, dims).to(device); + var mask = tensor(layout.Mask, dims).to(device); + var groups = tensor(layout.Groups, dims).to(device); + var types = tensor(layout.Types, dims).to(device); + + var (scaled, loc, scale) = T0Preprocess.ScaleInput(values, mask, groups, types, config); + var quantiles = model.Backbone(scaled, mask, groups, types); + var data = T0Preprocess.Rescale(quantiles, loc, scale, config.PatchSize, config.ScalerUseArcsinh); + + // The output at patch i forecasts patch i + 1, so the horizon starts at the last + // context patch. + var p = config.PatchSize; + var forecast = data[targetRow] + .narrow(0, layout.ContextPatches - 1, layout.ForecastPatches) + .reshape(layout.ForecastPatches * p, config.NumQuantiles) + .narrow(0, 0, layout.Horizon); + forecast = nan_to_num(forecast, 0.0, 0.0, 0.0); // the reference zeroes non-finite output + return forecast.cpu().contiguous().data().ToArray(); + } + + /// + /// The per-step input the model reads: the context left-padded onto the patch grid, then + /// the horizon rounded up to whole patches. Over the horizon a target is withheld and a + /// known-future row carries its values; padding is inert (value 0, group and role -1). + /// + private sealed class Layout + { + public readonly int Rows, Width, PadLeft, Context, Horizon, ContextPatches, ForecastPatches; + public readonly float[] Values; + public readonly long[] Mask, Groups, Types; + + public Layout(int patch, int rows, int context, int horizon) + { + (Rows, Context, Horizon) = (rows, context, horizon); + PadLeft = (patch - context % patch) % patch; + ForecastPatches = (horizon + patch - 1) / patch; + ContextPatches = (PadLeft + context) / patch; + Width = PadLeft + context + ForecastPatches * patch; + Values = new float[rows * Width]; + Mask = new long[rows * Width]; + Groups = new long[rows * Width]; + Types = new long[rows * Width]; + } + + /// Row index. + /// Observed steps. + /// Known future values, or empty for a target. + public void Row(int row, ReadOnlySpan context, ReadOnlySpan future) + { + var isFuture = !future.IsEmpty; + var o = row * Width; + for (var i = 0; i < Width; i++) + { + var pad = i < PadLeft; + Groups[o + i] = pad ? -1 : 0; + Types[o + i] = pad ? -1 : isFuture ? VariateType.Future : VariateType.Target; + if (pad) { Mask[o + i] = MaskType.Pad; continue; } + + var t = i - PadLeft; + if (t < Context) Write(o + i, context[t]); + else if (!isFuture) Mask[o + i] = MaskType.Withheld; + else if (t - Context < Horizon) Write(o + i, future[t - Context]); + else Mask[o + i] = MaskType.Pad; // the horizon's last patch, beyond what was asked + } + } + + private void Write(int at, float value) + { + if (float.IsNaN(value)) { Mask[at] = MaskType.Missing; return; } + Values[at] = value; + } + } +} diff --git a/src/Tsfm.Forecasting.T0/T0Model.cs b/src/Tsfm.Forecasting.T0/T0Model.cs new file mode 100644 index 0000000..8b2a3be --- /dev/null +++ b/src/Tsfm.Forecasting.T0/T0Model.cs @@ -0,0 +1,173 @@ +using System.Text.Json; + +using TorchSharp; + +using static TorchSharp.torch; +using static TorchSharp.torch.nn; + +namespace Tsfm.Forecasting.T0; + +/// +/// Checkpoint hyper-parameters, read from the published config.json. +/// +/// Absent fields take t0-alpha's values, as the reference does: alpha's config +/// predates them. t0-beta states scaler_eps 0.01 under std_clamp; run +/// under alpha's defaults (0.1, variance_offset) it loads and forecasts without +/// complaint, just worse — so these are read, never assumed. +/// +public sealed record T0Config( + long EmbedDim, int NumLayers, long NumHeads, long MlpHiddenDim, int PatchSize, int GroupEveryN, + double[] Quantiles, bool ScalerUseArcsinh, double ScalerEps, bool ScalerStdClamp, + bool XposHalvesLayout, bool ScalerZeroNoValue) +{ + public int NumQuantiles => Quantiles.Length; + + public static T0Config Load(string path) => Parse(File.ReadAllText(path)); + + public static T0Config Parse(string json) + { + using var doc = JsonDocument.Parse(json); + var r = doc.RootElement; + string? Text(string name) => r.TryGetProperty(name, out var e) && e.ValueKind == JsonValueKind.String + ? e.GetString() : null; + bool Flag(string name, bool fallback) => r.TryGetProperty(name, out var e) ? e.GetBoolean() : fallback; + + var epsMode = Text("scaler_eps_mode") ?? "variance_offset"; + var layout = Text("xpos_scale_layout") ?? "interleaved"; + if (epsMode is not ("variance_offset" or "std_clamp")) + throw new InvalidDataException($"unknown scaler_eps_mode '{epsMode}'"); + if (layout is not ("interleaved" or "halves")) + throw new InvalidDataException($"unknown xpos_scale_layout '{layout}'"); + + return new T0Config( + r.GetProperty("embed_dim").GetInt64(), + r.GetProperty("num_layers").GetInt32(), + r.GetProperty("num_heads").GetInt64(), + r.GetProperty("mlp_hidden_dim").GetInt64(), + r.GetProperty("patch_size").GetInt32(), + r.GetProperty("group_every_n").GetInt32(), + [.. r.GetProperty("quantile_levels").EnumerateArray().Select(q => q.GetDouble())], + Flag("scaler_use_arcsinh", true), + r.TryGetProperty("scaler_eps", out var eps) ? eps.GetDouble() : 0.1, + epsMode == "std_clamp", + layout == "halves", + Flag("scaler_zero_no_value", true)); + } +} + +/// +/// t0: patch encoder, alternating time/group transformer, and a residual decoder whose +/// output is made monotone across quantiles. +/// +/// Next-patch prediction: the output at patch i forecasts patch i + 1. +/// Every patch of a known horizon is decoded in the same pass, so a predictive interval over +/// up to steps costs one forward. +/// +public sealed class T0Model : Module +{ + private readonly T0Config _cfg; + + public PatchEncoder patch_encoder; + public Transformer transformer; + public ResidualBlock decoder; + + public T0Model(T0Config cfg) : base(nameof(T0Model)) + { + _cfg = cfg; + patch_encoder = new PatchEncoder(cfg.EmbedDim, cfg.PatchSize); + transformer = new Transformer(cfg); + decoder = new ResidualBlock(cfg.EmbedDim, cfg.EmbedDim, cfg.PatchSize * cfg.NumQuantiles); + RegisterComponents(); + } + + public T0Config Config => _cfg; + + /// Split a per-step array (rows, width) into patches. + public Tensor Patch(Tensor x) => x.view(x.shape[0], x.shape[1] / _cfg.PatchSize, _cfg.PatchSize); + + /// + /// Per-patch quantiles in scaled space. + /// + /// Scaled values (rows, width), width a multiple of the patch. + /// per step. + /// Series id per step; rows sharing one attend to each other. + /// -1 on padding. + /// per step, -1 on padding. + /// (rows, patches, patch, quantiles), non-decreasing along the last axis. + public Tensor Backbone(Tensor scaled, Tensor mask, Tensor groupIds, Tensor variateType) + { + using var scope = NewDisposeScope(); + var (values, m, g, t) = (Patch(scaled), Patch(mask), Patch(groupIds), Patch(variateType)); + var embedded = patch_encoder.forward(values, m, t); + var (timeMask, groupMask) = T0Preprocess.AttentionMasks(g, t, m); + var h = transformer.Forward(embedded, timeMask, groupMask); + var decoded = decoder.forward(h); + decoded = decoded.view(decoded.shape[0], decoded.shape[1], _cfg.PatchSize, _cfg.NumQuantiles); + return Monotone(decoded).MoveToOuterDisposeScope(); + } + + /// The quantile head: the first column free, each later one the previous plus + /// a softplus, so quantiles cannot cross. + public static Tensor Monotone(Tensor x) + { + using var scope = NewDisposeScope(); + var q = x.shape[^1]; + var first = x.narrow(-1, 0, 1); + if (q == 1) return first.MoveToOuterDisposeScope(); + var rest = first + functional.softplus(x.narrow(-1, 1, q - 1)).cumsum(-1); + return cat([first, rest], -1).MoveToOuterDisposeScope(); + } + + /// Bind published weights by their PyTorch names. Every parameter must be + /// matched: an unbound tensor keeps its random init and still forecasts plausibly. + public static T0Model Load(string checkpointDir, Device device) + { + var cfg = T0Config.Load(Path.Combine(checkpointDir, "config.json")); + using var stream = File.OpenRead(Path.Combine(checkpointDir, "model.safetensors")); + return Load(cfg, stream, device, checkpointDir); + } + + /// + public static T0Model Load(T0Config cfg, Stream weights, Device device, string label = "") + { + var net = new T0Model(cfg); + using var st = Safetensors.Load(weights, label); + + // The head's levels ride in the checkpoint as a buffer. A config that disagrees with + // them would label every column wrongly while computing it correctly. + if (st.Has("head.quantile_levels")) + { + using var levels = st.Get("head.quantile_levels"); + var stored = levels.data().ToArray(); + if (stored.Length != cfg.NumQuantiles + || stored.Where((l, i) => Math.Abs(l - cfg.Quantiles[i]) > 1e-6).Any()) + throw new InvalidDataException( + $"{label}: config quantiles [{string.Join(", ", cfg.Quantiles)}] disagree with " + + $"the checkpoint's [{string.Join(", ", stored)}]"); + } + + var bound = 0; + var missing = new List(); + using (no_grad()) + { + foreach (var (name, param) in net.named_parameters()) + { + if (!st.Has(name)) { missing.Add(name); continue; } + using var t = st.Get(name); + if (!t.shape.SequenceEqual(param.shape)) + throw new InvalidDataException( + $"{name}: checkpoint {string.Join('x', t.shape)} vs model {string.Join('x', param.shape)}"); + param.copy_(t); + bound++; + } + } + if (missing.Count > 0) + throw new InvalidDataException( + $"{missing.Count} parameters absent from {label}, e.g. {string.Join(", ", missing.Take(5))}"); + + Console.Error.WriteLine($"bound {bound} parameters from {label}"); + net.to(device); + net.eval(); + return net; + } +} diff --git a/src/Tsfm.Forecasting.T0/T0Modules.cs b/src/Tsfm.Forecasting.T0/T0Modules.cs new file mode 100644 index 0000000..46d5771 --- /dev/null +++ b/src/Tsfm.Forecasting.T0/T0Modules.cs @@ -0,0 +1,296 @@ +using TorchSharp.Modules; + +using static TorchSharp.torch; +using static TorchSharp.torch.nn; + +namespace Tsfm.Forecasting.T0; + +/// RMS norm, weight named scale, eps 1e-8 — not TimesFM's 1e-6. +public sealed class RmsNorm(long dims, double eps = 1e-8) : Module(nameof(RmsNorm)) +{ + public Parameter scale = new(ones(dims)); + + public override Tensor forward(Tensor x) + { + using var s = NewDisposeScope(); + return (x / (x * x).mean([-1L], keepdim: true).add(eps).sqrt() * scale).MoveToOuterDisposeScope(); + } +} + +/// Two biased linear layers with a ReLU between. +public sealed class Mlp : Module +{ + public Linear hidden_layer, output_layer; + + public Mlp(long input, long hidden, long output) : base(nameof(Mlp)) + { + hidden_layer = Linear(input, hidden); + output_layer = Linear(hidden, output); + RegisterComponents(); + } + + public override Tensor forward(Tensor x) + { + using var s = NewDisposeScope(); + return output_layer.forward(functional.relu(hidden_layer.forward(x))).MoveToOuterDisposeScope(); + } +} + +/// mlp(x) + residual_layer(x), every projection biased. +public sealed class ResidualBlock : Module +{ + public Mlp mlp; + public Linear residual_layer; + + public ResidualBlock(long input, long hidden, long output) : base(nameof(ResidualBlock)) + { + mlp = new Mlp(input, hidden, output); + residual_layer = Linear(input, output); + RegisterComponents(); + } + + public override Tensor forward(Tensor x) + { + using var s = NewDisposeScope(); + return (mlp.forward(x) + residual_layer.forward(x)).MoveToOuterDisposeScope(); + } +} + +/// +/// Embeds each patch as [values ‖ within-patch time ‖ validity], then adds a learned +/// embedding of the row's role (target, historical, known-future). +/// +public sealed class PatchEncoder : Module +{ + private readonly int _patch; + + public ResidualBlock projection; + public Embedding type_embeddings; + + public PatchEncoder(long embedDim, int patchSize) : base(nameof(PatchEncoder)) + { + _patch = patchSize; + projection = new ResidualBlock(3 * patchSize, embedDim, embedDim); + type_embeddings = Embedding(3, embedDim); + RegisterComponents(); + } + + /// Scaled values (rows, patches, patch). + /// Patched . + /// Patched , -1 on padding. + public override Tensor forward(Tensor values, Tensor mask, Tensor variateType) + { + using var s = NewDisposeScope(); + var (v, p) = (values.shape[0], values.shape[1]); + var time = (arange(_patch, values.dtype, values.device) / _patch).expand(v, p, _patch); + var validity = mask.eq(MaskType.Valid).to(values.dtype); + var embedded = projection.forward(cat([values, time, validity], -1)); + // Left padding must not retype a covariate as a target: the role comes from the first + // real step, and an all-padding patch (-1) is clamped to 0 and masked downstream. + var role = T0Preprocess.ReducePatchMetadata(variateType, mask).clamp_min(0); + return (embedded + type_embeddings.forward(role)).MoveToOuterDisposeScope(); + } +} + +/// +/// Rotary embedding over patches with xPos scaling: queries are scaled up and keys down by +/// a per-frequency factor that decays with distance from the window's centre. +/// +/// Interleaved, not split-half: pairs are adjacent dimensions (2i, 2i+1), as +/// rotary-embedding-torch forms them. TimesFM pairs the halves; reusing its RoPE here would +/// load and run without complaint. +/// +public sealed class XposRotary(long headDim, bool halvesLayout, double theta = 10000.0, double scaleBase = 512.0) +{ + /// Cos, sin and xPos scale for positions, each + /// (length, headDim). Computed per call: the window length varies. + public (Tensor Cos, Tensor Sin, Tensor Scale) Tables(long length, Device device) + { + using var s = NewDisposeScope(); + var idx = arange(0, headDim, 2, ScalarType.Int64, device); + var exponent = idx.to(ScalarType.Float32) / headDim; + var inv = 1.0 / full_like(exponent, theta).pow(exponent); + var pos = arange(length, ScalarType.Float32, device); + var angles = (pos.unsqueeze(-1) * inv).repeat_interleave(2, -1); + + var basis = (idx + 0.4 * headDim) / (1.4 * headDim); + var power = (pos - Math.Floor((length - 1) / 2.0)) / scaleBase; + var half = basis.pow(power.unsqueeze(-1)); + var scale = halvesLayout ? cat([half, half], -1) : half.repeat_interleave(2, -1); + + return (angles.cos().MoveToOuterDisposeScope(), angles.sin().MoveToOuterDisposeScope(), + scale.MoveToOuterDisposeScope()); + } + + /// Rotate (…, seq, headDim) by the tables; is the + /// xPos factor, inverted for keys. + public static Tensor Apply(Tensor t, Tensor cos, Tensor sin, Tensor scale) + { + using var s = NewDisposeScope(); + var pairs = t.reshape([.. t.shape[..^1], t.shape[^1] / 2, 2]); + var rotated = stack([-pairs.select(-1, 1), pairs.select(-1, 0)], -1).flatten(-2); // (-x2, x1) + return (t * cos * scale + rotated * sin * scale).MoveToOuterDisposeScope(); + } +} + +/// +/// Multi-head self-attention over the middle axis of (batch, seq, embed), with a +/// fused biased QKV projection and per-head RMS norm on queries and keys. With rotary +/// tables supplied, they are applied AFTER the norm — TimesFM applies them before. +/// +public sealed class SelfAttention : Module +{ + private readonly long _heads, _headDim; + + public Linear wQKV, wO; + public RmsNorm q_norm, k_norm; + + public SelfAttention(long embedDim, long heads) : base(nameof(SelfAttention)) + { + (_heads, _headDim) = (heads, embedDim / heads); + wQKV = Linear(embedDim, 3 * embedDim); + wO = Linear(embedDim, embedDim); + q_norm = new RmsNorm(_headDim); + k_norm = new RmsNorm(_headDim); + RegisterComponents(); + } + + /// (batch, seq, embed) + /// (batch, 1, seq, seq), true where attention is allowed. + /// Rotary tables, or null for unordered (variate) attention. + public Tensor Forward(Tensor x, Tensor mask, (Tensor Cos, Tensor Sin, Tensor Scale)? rope) + { + using var s = NewDisposeScope(); + var (b, n) = (x.shape[0], x.shape[1]); + + // "(three h d)" on the projection: three-major, then heads. + var qkv = wQKV.forward(x).view(b, n, 3, _heads, _headDim).permute(2, 0, 3, 1, 4); + var q = q_norm.forward(qkv[0]); + var k = k_norm.forward(qkv[1]); + var v = qkv[2]; + + if (rope is var (cos, sin, scale)) + { + q = XposRotary.Apply(q, cos, sin, scale); + k = XposRotary.Apply(k, cos, sin, scale.pow(-1)); // keys take the inverse + } + + var o = functional.scaled_dot_product_attention(q, k, v, attn_mask: mask); + // A query with no permitted key (a padding patch) must yield zero, as the reference's + // safe softmax does, not NaN: as a value it would poison every row that reads it. + o = where(mask.any(-1, keepdim: true), o, zeros_like(o)); + + o = o.transpose(1, 2).reshape(b, n, _heads * _headDim); + return wO.forward(o).MoveToOuterDisposeScope(); + } +} + +/// Pre-norm attention with a residual. Group blocks attend across rows at each +/// patch, by swapping the row and patch axes around the attention. +public sealed class AttentionBlock : Module +{ + private readonly bool _group; + + public RmsNorm norm; + public SelfAttention attention; + + public AttentionBlock(long embedDim, long heads, bool group) : base(nameof(AttentionBlock)) + { + _group = group; + norm = new RmsNorm(embedDim); + attention = new SelfAttention(embedDim, heads); + RegisterComponents(); + } + + public Tensor Forward(Tensor x, Tensor mask, (Tensor, Tensor, Tensor)? rope) + { + using var s = NewDisposeScope(); + var h = norm.forward(x); + var o = _group + ? attention.Forward(h.transpose(0, 1), mask, null).transpose(0, 1) + : attention.Forward(h, mask, rope); + return (x + o).MoveToOuterDisposeScope(); + } +} + +/// SwiGLU with the gate in the FIRST half — the xFormers ordering, the reverse of +/// the usual. +public sealed class SwiGlu() : Module(nameof(SwiGlu)) +{ + public override Tensor forward(Tensor x) + { + using var s = NewDisposeScope(); + var halves = x.chunk(2, -1); + return (functional.silu(halves[0]) * halves[1]).MoveToOuterDisposeScope(); + } +} + +/// One attention block — time or group — then a pre-norm SwiGLU feed-forward. +public sealed class TransformerLayer : Module +{ + public bool IsGroup { get; } + + public AttentionBlock attention_block; + public RmsNorm norm; + public Sequential mlp; + + public TransformerLayer(long embedDim, long heads, long hidden, bool group) : base(nameof(TransformerLayer)) + { + IsGroup = group; + attention_block = new AttentionBlock(embedDim, heads, group); + norm = new RmsNorm(embedDim); + // Named as the checkpoint's nn.Sequential indexes them: mlp.0 and mlp.2. + mlp = Sequential( + ("0", Linear(embedDim, 2 * hidden)), + ("1", new SwiGlu()), + ("2", Linear(hidden, embedDim))); + RegisterComponents(); + } + + public Tensor Forward(Tensor x, Tensor timeMask, Tensor groupMask, (Tensor, Tensor, Tensor) rope) + { + using var s = NewDisposeScope(); + var h = IsGroup + ? attention_block.Forward(x, groupMask, null) + : attention_block.Forward(x, timeMask, rope); + return (h + mlp.forward(norm.forward(h))).MoveToOuterDisposeScope(); + } +} + +/// +/// The layer stack: every group_every_n-th layer attends across rows, the rest along +/// time. For the published checkpoints that is two time layers then one group layer, eight +/// times over. +/// +public sealed class Transformer : Module +{ + private readonly XposRotary _rope; + + public ModuleList layers; + public RmsNorm out_norm; + + public Transformer(T0Config cfg) : base(nameof(Transformer)) + { + _rope = new XposRotary(cfg.EmbedDim / cfg.NumHeads, cfg.XposHalvesLayout); + layers = new ModuleList( + [.. Enumerable.Range(0, cfg.NumLayers).Select(i => new TransformerLayer( + cfg.EmbedDim, cfg.NumHeads, cfg.MlpHiddenDim, + group: cfg.GroupEveryN > 0 && (i + 1) % cfg.GroupEveryN == 0))]); + out_norm = new RmsNorm(cfg.EmbedDim); + RegisterComponents(); + } + + public (Tensor, Tensor, Tensor) Rope(long patches, Device device) => _rope.Tables(patches, device); + + /// Patch embeddings (rows, patches, embed). + /// From . + /// From . + public Tensor Forward(Tensor x, Tensor timeMask, Tensor groupMask) + { + using var s = NewDisposeScope(); + var rope = Rope(x.shape[1], x.device); + var h = x; + foreach (var layer in layers) h = layer.Forward(h, timeMask, groupMask, rope); + return out_norm.forward(h).MoveToOuterDisposeScope(); + } +} diff --git a/src/Tsfm.Forecasting.T0/T0Preprocess.cs b/src/Tsfm.Forecasting.T0/T0Preprocess.cs new file mode 100644 index 0000000..8e4227a --- /dev/null +++ b/src/Tsfm.Forecasting.T0/T0Preprocess.cs @@ -0,0 +1,155 @@ +using static TorchSharp.torch; + +namespace Tsfm.Forecasting.T0; + +/// Role of a row. Values match the type embedding indices the checkpoint learned. +public static class VariateType +{ + public const long Target = 0, Historical = 1, Future = 2; +} + +/// Why a cell is not an observation. Withheld marks steps to predict. +public static class MaskType +{ + public const long Valid = 0, Pad = 1, Missing = 2, Censored = 3, Withheld = 4; +} + +/// +/// The scaling and masking that surround the transformer. Every array is per time step, +/// (rows, width), with the width already aligned to whole patches. +/// +/// Targets are standardised by causal running statistics — each step by the mean and +/// deviation of every valid step up to and including it — and known-future rows by their +/// global statistics over context and horizon alike. +/// +public static class T0Preprocess +{ + /// Causal mean and deviation along time, one series per row. + /// Cumulative sums rather than a recurrence, exactly as the reference forms + /// them: the checkpoint was trained on these numerics, and the two eps conventions + /// disagree wherever a local deviation sits near eps. + public static (Tensor Loc, Tensor Scale) CausalStats(Tensor x, Tensor invalid, double eps, bool stdClamp) + { + using var scope = NewDisposeScope(); + var valid = (~invalid).to(x.dtype); + var count = valid.cumsum(-1).clamp_min(1.0); + var masked = where(invalid, zeros_like(x), x); + var means = masked.cumsum(-1) / count; + var shifted = cat([zeros_like(means.narrow(1, 0, 1)), means.narrow(1, 0, means.shape[1] - 1)], -1); + var m2 = ((masked - shifted) * (masked - means) * valid).cumsum(-1).clamp_min(0.0); + var variance = m2 / (count - 1.0).clamp_min(1.0); + var scale = stdClamp ? variance.sqrt().clamp_min(eps) : (variance + eps).sqrt(); + return (means.MoveToOuterDisposeScope(), scale.MoveToOuterDisposeScope()); + } + + /// Global mean and deviation per row, broadcast back over time. The deviation + /// is always clamped at eps, whatever the causal convention. + public static (Tensor Loc, Tensor Scale) GlobalStats(Tensor x, Tensor invalid, double eps) + { + using var scope = NewDisposeScope(); + var valid = (~invalid).to(x.dtype); + var count = RowSum(valid); + var masked = where(invalid, zeros_like(x), x); + var mean = RowSum(masked) / count.clamp_min(1.0); + var squared = where(invalid, zeros_like(x), (x - mean) * (x - mean)); + var std = (RowSum(squared) / count.clamp_min(2.0)).sqrt().clamp_min(eps); + return (mean.expand_as(x).MoveToOuterDisposeScope(), std.expand_as(x).MoveToOuterDisposeScope()); + } + + /// Sum over the last axis by a one-bucket scatter-add, the summation order the + /// reference uses. A plain sum reduces pairwise and differs in the last bits. + private static Tensor RowSum(Tensor x) + { + var index = zeros(x.shape, ScalarType.Int64, x.device); + return zeros([x.shape[0], 1], x.dtype, x.device).scatter_add_(-1, index, x); + } + + /// Standardise every row; returns the scaled values and the statistics used, + /// which needs to invert the model's output. + public static (Tensor Scaled, Tensor Loc, Tensor Scale) ScaleInput( + Tensor x, Tensor mask, Tensor groupIds, Tensor variateType, T0Config cfg) + { + using var scope = NewDisposeScope(); + var invalid = mask.ne(MaskType.Valid) | x.isnan(); + if (!cfg.ScalerZeroNoValue) x = where(x.isnan(), zeros_like(x), x); + + var (causalLoc, causalScale) = CausalStats(x, invalid, cfg.ScalerEps, cfg.ScalerStdClamp); + var (futureLoc, futureScale) = GlobalStats(x, invalid, cfg.ScalerEps); + + // Padding carries group -1 and takes neither statistic. + var real = groupIds.ge(0); + var isFuture = real & variateType.eq(VariateType.Future); + var isCausal = real & (variateType.eq(VariateType.Target) | variateType.eq(VariateType.Historical)); + var loc = where(isFuture, futureLoc, where(isCausal, causalLoc, zeros_like(x))); + var scale = where(isFuture, futureScale, where(isCausal, causalScale, ones_like(x))); + + var scaled = (x - loc) / scale; + if (cfg.ScalerUseArcsinh) scaled = scaled.asinh(); + if (cfg.ScalerZeroNoValue) + { + // A cell holding no observation reads as zero whatever placeholder it carried; + // a censored cell keeps its scaled bound. + var noValue = mask.eq(MaskType.Pad) | mask.eq(MaskType.Missing) | mask.eq(MaskType.Withheld); + scaled = where(noValue, zeros_like(scaled), scaled); + } + return (scaled.MoveToOuterDisposeScope(), loc.MoveToOuterDisposeScope(), scale.MoveToOuterDisposeScope()); + } + + /// Invert scaling on per-patch predictions (rows, patches, patch, quantiles), + /// using the statistics at each patch's last step. + public static Tensor Rescale(Tensor predictions, Tensor loc, Tensor scale, int patchSize, bool arcsinh) + { + using var scope = NewDisposeScope(); + var (rows, patches) = (predictions.shape[0], predictions.shape[1]); + var l = loc.reshape(rows, patches, patchSize).select(-1, patchSize - 1).unsqueeze(-1).unsqueeze(-1); + var s = scale.reshape(rows, patches, patchSize).select(-1, patchSize - 1).unsqueeze(-1).unsqueeze(-1); + var values = arcsinh ? predictions.sinh() : predictions; + return (values * s + l).MoveToOuterDisposeScope(); + } + + /// Collapse patched metadata (rows, patches, patch) to each patch's first + /// non-padding step, or -1 where the patch is all padding. + public static Tensor ReducePatchMetadata(Tensor patched, Tensor patchedMask) + { + using var scope = NewDisposeScope(); + var isReal = patchedMask.ne(MaskType.Pad); + var first = isReal.to(ScalarType.Int32).argmax(-1, keepdim: true); + var reduced = patched.gather(-1, first).squeeze(-1); + return where(isReal.any(-1), reduced, full_like(reduced, -1)).MoveToOuterDisposeScope(); + } + + /// + /// The two attention masks of a forward pass, true where attention is allowed. + /// + /// Time (rows, 1, patches, patches): causal for targets and historicals, + /// bidirectional for known-future rows; a patch reads only its own series, and nothing + /// reads an all-padding patch. + /// Group (patches, 1, rows, rows): rows sharing a group attend to one + /// another, but a known-future row reads only known-future rows — information flows + /// from covariates into targets and never back. + /// + /// + public static (Tensor Time, Tensor Group) AttentionMasks( + Tensor patchedGroupIds, Tensor patchedVariateType, Tensor patchedMask) + { + using var scope = NewDisposeScope(); + var groups = ReducePatchMetadata(patchedGroupIds, patchedMask); // (V, P) + var types = ReducePatchMetadata(patchedVariateType, patchedMask); + var attendable = patchedMask.ne(MaskType.Pad).any(-1); + var valid = groups.ge(0); + var isFuture = types.eq(VariateType.Future); + + var sameSeries = groups.unsqueeze(2).eq(groups.unsqueeze(1)) & valid.unsqueeze(2) & valid.unsqueeze(1); + var positions = arange(groups.shape[1], ScalarType.Int64, groups.device); + var causal = positions.unsqueeze(1).ge(positions.unsqueeze(0)); + var futureQuery = isFuture.unsqueeze(2); + var time = (futureQuery & sameSeries) | (~futureQuery & sameSeries & causal); + time = time & attendable.unsqueeze(1); + + var (ids, ok, fut) = (groups.t(), valid.t(), isFuture.t()); + var group = ids.unsqueeze(2).eq(ids.unsqueeze(1)) & ok.unsqueeze(2) & ok.unsqueeze(1); + group = group & ~(fut.unsqueeze(2) & ~fut.unsqueeze(1)); + + return (time.unsqueeze(1).MoveToOuterDisposeScope(), group.unsqueeze(1).MoveToOuterDisposeScope()); + } +} diff --git a/src/Tsfm.Forecasting.T0/Tsfm.Forecasting.T0.csproj b/src/Tsfm.Forecasting.T0/Tsfm.Forecasting.T0.csproj new file mode 100644 index 0000000..b798734 --- /dev/null +++ b/src/Tsfm.Forecasting.T0/Tsfm.Forecasting.T0.csproj @@ -0,0 +1,24 @@ + + + + true + Tsfm.Forecasting.T0 + Tsfm.Forecasting.T0 + Tsfm.Forecasting.T0 + + Apache-2.0 + Native .NET inference for The Forecasting Company's t0 family (t0-alpha, t0-beta): a decoder-style patch transformer with time and variate attention and known-future covariates. Weights are Apache-2.0 but not bundled; fetch them from Hugging Face. + timeseries;forecasting;foundation-model;t0;covariates;quantile;torchsharp + + + + + + + + + + + diff --git a/src/Tsfm.Forecasting.TimesFm/TimesFmForecaster.cs b/src/Tsfm.Forecasting.TimesFm/TimesFmForecaster.cs index fb1dd95..7a8187a 100644 --- a/src/Tsfm.Forecasting.TimesFm/TimesFmForecaster.cs +++ b/src/Tsfm.Forecasting.TimesFm/TimesFmForecaster.cs @@ -13,12 +13,16 @@ namespace Tsfm.Forecasting.TimesFm; /// /// Inference is not thread safe. public sealed class TimesFmForecaster(TimesFmModel model, TimesFmConfig config, Device device) + : IQuantileForecaster { private const int CloseChannel = 3; // open, high, low, close, volume, amount private const int Channels = 6; public TimesFmConfig Config => config; + /// + public IReadOnlyList Quantiles => config.Quantiles; + public static TimesFmForecaster Load(string checkpointDir, Device device) { var cfg = TimesFmConfig.Load(Path.Combine(checkpointDir, "config.json")); diff --git a/src/Tsfm.Forecasting/IQuantileForecaster.cs b/src/Tsfm.Forecasting/IQuantileForecaster.cs new file mode 100644 index 0000000..ac04a3d --- /dev/null +++ b/src/Tsfm.Forecasting/IQuantileForecaster.cs @@ -0,0 +1,40 @@ +namespace Tsfm.Forecasting; + +/// +/// A model that emits a quantile forecast for one series of a jointly attended set, any of +/// which may be known into the future. Implemented by the TimesFM and t0 adapters, so one +/// call site can drive either. +/// +/// Shapes are shared; constraints are not. Each implementation documents its own +/// limits on context, horizon and variate count — TimesFM, for instance, needs whole +/// patches where t0 pads. +/// +/// Implementations are not thread safe. +public interface IQuantileForecaster +{ + /// Levels of the returned quantile columns, ascending — the checkpoint's own + /// trained grid, so its length varies between checkpoints. + IReadOnlyList Quantiles { get; } + + /// Joint forecast over several series. + /// Variate-major [variates x (contextBars + horizon)]: series + /// v at step t is series[v * (contextBars + horizon) + t]. The future + /// region of a variate not marked known is never read. + /// Series supplied. + /// Observed steps. + /// Steps ahead to return. + /// Per variate: true where its future region holds real values + /// the model may read. Empty means none. A variate marked known is not a target. + /// Which series the returned quantiles describe. + /// [horizon, Quantiles.Count] in the units of the target series. + double[,] ForecastJoint( + ReadOnlySpan series, int variates, int contextBars, int horizon, + ReadOnlySpan knownFuture = default, int targetVariate = 0); + + /// Forward-return quantiles of the close from OHLCVA bars, every channel a + /// target. + /// Row-major [L x 6]: open, high, low, close, volume, amount. + /// Steps ahead to return. + /// [horizon, Quantiles.Count] of cumulative return relative to the anchor close. + double[,] Forecast(ReadOnlySpan ohlcva, int horizon); +} diff --git a/tests/Tsfm.Forecasting.T0.Parity/Program.cs b/tests/Tsfm.Forecasting.T0.Parity/Program.cs new file mode 100644 index 0000000..1e481dc --- /dev/null +++ b/tests/Tsfm.Forecasting.T0.Parity/Program.cs @@ -0,0 +1,154 @@ +// Bisecting parity harness for t0. Compares the .NET port against tensors dumped from the +// published tfc-t0 package by reference/dump_t0_parity.py, one stage at a time, then end +// to end through ForecastJoint — including the quickstart notebook's own examples. +// +// Each stage is fed the REFERENCE output of the stage before it, so a failure localises to +// one stage rather than being inherited from an earlier one. +// +// dotnet run --project tests/Tsfm.Forecasting.T0.Parity -c Release -- t0-alpha t0-beta + +using TorchSharp; + +using Tsfm.Forecasting; +using Tsfm.Forecasting.T0; + +using static TorchSharp.torch; + +var root = Path.GetFullPath(Path.Combine(AppContext.BaseDirectory, "../../../../..")); +var names = args.Length > 0 ? args : ["t0-alpha"]; + +Device device; +try { device = new Device(DeviceType.MPS); using (ones(1).to(device)) { } } +catch { device = new Device(DeviceType.CPU); } +Console.WriteLine($"device={device.type}"); + +var ok = true; +foreach (var name in names) ok &= Run(name); +Console.WriteLine(); +Console.WriteLine(ok ? "PARITY: PASS" : "PARITY: FAIL — the first failing stage above localises the fault"); +return ok ? 0 : 1; + +bool Run(string name) +{ + var fixturePath = Path.Combine(root, "reference", $"{name}-parity.safetensors"); + if (!File.Exists(fixturePath)) + { + Console.Error.WriteLine($"missing {fixturePath}; run reference/dump_t0_parity.py checkpoints/{name}"); + return false; + } + + using var _ = no_grad(); + using var fx = Safetensors.Load(fixturePath); + Tensor E(string n) => fx.Get(n).to(device); + Tensor L(string n) => E(n).to(ScalarType.Int64); + + var forecaster = T0Forecaster.Load(Path.Combine(root, "checkpoints", name), device); + var (model, cfg) = (forecaster.Model, forecaster.Config); + var pass = true; + + Console.WriteLine($"\n── {name}: {cfg.NumLayers} layers x {cfg.EmbedDim}, {cfg.NumQuantiles} quantiles, " + + $"eps {cfg.ScalerEps} {(cfg.ScalerStdClamp ? "std_clamp" : "variance_offset")}"); + Console.WriteLine($"{"stage",-24} {"max|Δ|",12} {"rel",12} note"); + + void Check(string stage, Tensor got, Tensor want, double tol, string note = "") + { + var absMax = (got.to(ScalarType.Float32).to(want.device) - want).abs().max().item(); + var scale = want.abs().max().item(); + var rel = absMax / Math.Max(scale, 1e-12); + var good = rel <= tol || absMax <= 1e-6; + pass &= good; + Console.WriteLine($"{stage,-24} {absMax,12:E3} {rel,12:E3} {(good ? "ok" : "FAIL")} {note}"); + } + + // 1. Scaling: causal statistics for targets, global for known-future rows. + var (values, mask, groups, types) = (E("buf_values"), L("buf_mask"), L("buf_groups"), L("buf_types")); + var (scaled, loc, scale) = T0Preprocess.ScaleInput(values, mask, groups, types, cfg); + Check("loc", loc, E("loc"), 1e-6); + Check("scale", scale, E("scale"), 1e-6); + Check("scaled", scaled, E("scaled"), 1e-5); + + // 2. Masks and patch embedding, from the reference's scaled input. + var (pv, pm, pg, pt) = (model.Patch(E("scaled")), model.Patch(mask), model.Patch(groups), model.Patch(types)); + var (timeMask, groupMask) = T0Preprocess.AttentionMasks(pg, pt, pm); + Check("time_mask", timeMask.to(ScalarType.Float32), E("time_mask"), 0); + Check("group_mask", groupMask.to(ScalarType.Float32), E("group_mask"), 0); + Check("patch_embedding", model.patch_encoder.forward(pv, pm, pt), E("patch_embedding"), 1e-5); + + // 3. Layers, each from the reference's previous output: time, time, group, ..., last. + var layers = model.transformer.layers; + var rope = model.transformer.Rope(pv.shape[1], device); + var h = E("patch_embedding"); + for (var i = 0; i < layers.Count; i++) + { + var refName = $"layer{i}_output"; + h = layers[i].Forward(h, timeMask, groupMask, rope); + if (!fx.Has(refName)) continue; + Check(refName, h, E(refName), 1e-4, layers[i].IsGroup ? "(group)" : "(time)"); + h = E(refName); + } + Check("transformer_output", model.transformer.Forward(E("patch_embedding"), timeMask, groupMask), + E("transformer_output"), 2e-4, "(whole stack)"); + + // 4. Decoder, quantile head and rescale. + var decoded = model.decoder.forward(E("transformer_output")); + Check("decoder_output", decoded, E("decoder_output"), 1e-5); + var head = T0Model.Monotone(E("decoder_output").view( + decoded.shape[0], decoded.shape[1], cfg.PatchSize, cfg.NumQuantiles)); + Check("head_output", head, E("head_output"), 1e-5); + Check("rescaled", T0Preprocess.Rescale(E("head_output"), E("loc"), E("scale"), cfg.PatchSize, + cfg.ScalerUseArcsinh), E("rescaled"), 1e-5); + + // 5. End to end through the public API, from raw series. + var ctx = fx.Get("in_ctx"); // (2, T) + var fut = fx.Get("in_fut"); // (2, T + H) + var (t, horizon) = ((int)ctx.shape[1], (int)(fut.shape[1] - ctx.shape[1])); + var total = t + horizon; + var joint = new float[4 * total]; + for (var v = 0; v < 2; v++) + for (var i = 0; i < t; i++) joint[v * total + i] = ctx[v, i].item(); + for (var v = 0; v < 2; v++) + for (var i = 0; i < total; i++) joint[(2 + v) * total + i] = fut[v, i].item(); + + for (var target = 0; target < 2; target++) + { + Check($"joint[{target}]", Grid(forecaster.ForecastJoint(joint, 4, t, horizon, + [false, false, true, true], target)), E("predict_joint")[target], 1e-4, "(2 targets + 2 known)"); + var plain = joint.AsSpan(0, 2 * total).ToArray(); + Check($"targets_only[{target}]", Grid(forecaster.ForecastJoint(plain, 2, t, horizon, + targetVariate: target)), E("predict_targets_only")[target], 1e-4); + } + var aligned = Enumerable.Range(64, 64).Select(i => ctx[0, i].item()).ToArray(); + Check("aligned", Grid(forecaster.ForecastJoint(aligned, 1, 64, 32)), E("predict_aligned"), 1e-4, + "(no padding)"); + + // 6. The quickstart notebook's examples. + var m4 = Floats(fx.Get("m4_ctx")); + Check("notebook m4 H1", Grid(forecaster.ForecastJoint(m4, 1, m4.Length, 48)), E("m4_predict"), 1e-4); + + var ep = Floats(fx.Get("ep_ctx")); + var cov = fx.Get("ep_cov"); // (2, 512 + 24) + var epTruth = Floats(fx.Get("ep_truth")); + var epPlain = forecaster.ForecastJoint(ep, 1, ep.Length, epTruth.Length); + Check("notebook ep", Grid(epPlain), E("ep_predict"), 1e-4); + + var width = (int)cov.shape[1]; + var epJoint = new float[3 * width]; + ep.CopyTo(epJoint, 0); + Floats(cov[0]).CopyTo(epJoint, width); + Floats(cov[1]).CopyTo(epJoint, 2 * width); + var epCov = forecaster.ForecastJoint(epJoint, 3, ep.Length, epTruth.Length, [false, true, true]); + Check("notebook ep + cov", Grid(epCov), E("ep_predict_cov"), 1e-4); + + var mid = Array.IndexOf(cfg.Quantiles, 0.5); + Console.WriteLine($" electricity price |median-truth|: {MeanAbs(epPlain, epTruth, mid):F4} without covariates, " + + $"{MeanAbs(epCov, epTruth, mid):F4} with"); + return pass; +} + +static float[] Floats(Tensor x) => x.contiguous().data().ToArray(); + +static Tensor Grid(double[,] q) => tensor(q.Cast().Select(d => (float)d).ToArray(), + [q.GetLength(0), q.GetLength(1)]); + +static double MeanAbs(double[,] q, float[] truth, int col) + => Enumerable.Range(0, truth.Length).Average(h => Math.Abs(q[h, col] - truth[h])); diff --git a/tests/Tsfm.Forecasting.T0.Parity/Tsfm.Forecasting.T0.Parity.csproj b/tests/Tsfm.Forecasting.T0.Parity/Tsfm.Forecasting.T0.Parity.csproj new file mode 100644 index 0000000..2364912 --- /dev/null +++ b/tests/Tsfm.Forecasting.T0.Parity/Tsfm.Forecasting.T0.Parity.csproj @@ -0,0 +1,16 @@ + + + + Exe + Tsfm.Forecasting.T0.Parity + + + + + + + + + + + diff --git a/tests/Tsfm.Forecasting.T0.Tests/CheckpointTests.cs b/tests/Tsfm.Forecasting.T0.Tests/CheckpointTests.cs new file mode 100644 index 0000000..37d6783 --- /dev/null +++ b/tests/Tsfm.Forecasting.T0.Tests/CheckpointTests.cs @@ -0,0 +1,97 @@ +using TorchSharp; +using Xunit; +using Xunit.Abstractions; + +using static TorchSharp.torch; + +namespace Tsfm.Forecasting.T0.Tests; + +/// +/// Behaviour of the published checkpoints through . +/// +/// Gated on the checkpoint being present, and so vacuous in CI, which does not fetch +/// ~1.4 GB of weights. Fetch them with scripts/fetch-t0-checkpoint.sh to run these locally; +/// numerical agreement with the reference is the parity harness's job, not this one's. +/// +public class CheckpointTests(ITestOutputHelper output) +{ + private const int Context = 96, Horizon = 24, Total = Context + Horizon, Variates = 3; + + private static string Dir(string name) => Path.Combine( + Path.GetFullPath(Path.Combine(AppContext.BaseDirectory, "../../../../..")), "checkpoints", name); + + private T0Forecaster? Load(string name) + { + if (File.Exists(Path.Combine(Dir(name), "model.safetensors"))) + return T0Forecaster.Load(Dir(name), new Device(DeviceType.CPU)); + output.WriteLine($"{name} absent — vacuous"); + return null; + } + + /// Variate 0 is the target, driven by variate 1 one step ahead; 1 is known + /// into the future; 2 is a calendar flag. + private static float[] Series(float driverFuture) + { + var s = new float[Variates * Total]; + for (var t = 0; t < Total; t++) + { + var driver = t < Context ? 20f + 4f * MathF.Sin(t / 6f) : 20f + driverFuture; + s[1 * Total + t] = driver; + s[2 * Total + t] = t % 7 < 5 ? 1f : 0f; + } + for (var t = 0; t < Context; t++) + s[t] = 100f + 3f * s[Total + Math.Max(t - 1, 0)] + s[2 * Total + t]; + return s; + } + + [Theory] + [InlineData("t0-alpha", 5)] + [InlineData("t0-beta", 21)] + public void QuantileGridComesFromTheCheckpoint(string name, int levels) + { + var f = Load(name); if (f is null) return; + Assert.Equal(levels, f.Quantiles.Count); + Assert.Contains(0.5, f.Quantiles); + } + + [Theory] + [InlineData("t0-alpha")] + [InlineData("t0-beta")] + public void KnownFutureValuesChangeTheForecast(string name) + { + var f = Load(name); if (f is null) return; + var a = f.ForecastJoint(Series(+8f), Variates, Context, Horizon, [false, true, true]); + var b = f.ForecastJoint(Series(-8f), Variates, Context, Horizon, [false, true, true]); + + var mid = Array.IndexOf(f.Config.Quantiles, 0.5); + var lift = Enumerable.Range(0, Horizon).Average(h => a[h, mid] - b[h, mid]); + output.WriteLine($"{name}: mean median lift from a higher driver = {lift:F2}"); + Assert.True(lift > 0, "a target that tracks its driver should rise with the driver's known future"); + } + + [Theory] + [InlineData("t0-alpha")] + [InlineData("t0-beta")] + public void IdenticalInputGivesIdenticalForecast(string name) + { + var f = Load(name); if (f is null) return; + var s = Series(0f); + Assert.Equal(f.ForecastJoint(s, Variates, Context, Horizon, [false, true, true]), + f.ForecastJoint(s, Variates, Context, Horizon, [false, true, true])); + } + + [Theory] + [InlineData("t0-alpha")] + [InlineData("t0-beta")] + public void QuantilesAreOrderedAndFinite(string name) + { + var f = Load(name); if (f is null) return; + var q = f.ForecastJoint(Series(0f), Variates, Context, Horizon, [false, true, true]); + for (var h = 0; h < Horizon; h++) + for (var i = 0; i < q.GetLength(1); i++) + { + Assert.True(double.IsFinite(q[h, i])); + if (i > 0) Assert.True(q[h, i] >= q[h, i - 1]); + } + } +} diff --git a/tests/Tsfm.Forecasting.T0.Tests/InvariantTests.cs b/tests/Tsfm.Forecasting.T0.Tests/InvariantTests.cs new file mode 100644 index 0000000..61125d3 --- /dev/null +++ b/tests/Tsfm.Forecasting.T0.Tests/InvariantTests.cs @@ -0,0 +1,215 @@ +using TorchSharp; +using Xunit; + +using static TorchSharp.torch; + +namespace Tsfm.Forecasting.T0.Tests; + +/// +/// Structural properties that hold for any weights, checked on a tiny randomly initialised +/// model so they run in CI without a checkpoint. Each is a leak or a blind spot that would +/// otherwise forecast plausibly: a target reading its own future, a known-future row +/// reading a target, padding acting as data. +/// +public class InvariantTests +{ + private static readonly T0Config Tiny = new( + EmbedDim: 16, NumLayers: 3, NumHeads: 2, MlpHiddenDim: 32, PatchSize: 4, GroupEveryN: 3, + Quantiles: [0.1, 0.5, 0.9], ScalerUseArcsinh: true, ScalerEps: 0.1, ScalerStdClamp: false, + XposHalvesLayout: false, ScalerZeroNoValue: true); + + private static T0Forecaster Forecaster() + { + manual_seed(7); + var model = new T0Model(Tiny); + model.eval(); + return new T0Forecaster(model, Tiny, new Device(DeviceType.CPU)); + } + + private const int Context = 10, Horizon = 6, Total = Context + Horizon; + + /// Variate 0 a target, 1 a second target, 2 known into the future. + private static float[] Series(Func? tweak = null) + { + var s = new float[3 * Total]; + for (var t = 0; t < Total; t++) + { + s[0 * Total + t] = 50f + 3f * MathF.Sin(t / 2f); + s[1 * Total + t] = 2f * MathF.Cos(t / 3f); + s[2 * Total + t] = t % 4; + } + if (tweak is not null) + for (var v = 0; v < 3; v++) + for (var t = 0; t < Total; t++) + s[v * Total + t] += tweak(v, t); + return s; + } + + private static readonly bool[] Known = [false, false, true]; + + [Fact] + public void TargetFutureRegionIsNeverRead() + { + var f = Forecaster(); + var a = f.ForecastJoint(Series(), 3, Context, Horizon, Known); + var b = f.ForecastJoint(Series((v, t) => v < 2 && t >= Context ? 1e4f : 0f), 3, Context, Horizon, Known); + Assert.Equal(a, b); + } + + [Fact] + public void KnownFutureValuesReachTheTarget() + { + var f = Forecaster(); + var a = f.ForecastJoint(Series(), 3, Context, Horizon, Known); + var b = f.ForecastJoint(Series((v, t) => v == 2 && t >= Context ? 5f : 0f), 3, Context, Horizon, Known); + Assert.NotEqual(a, b); + } + + /// The group mask lets known-future rows read only each other. If it leaked, + /// a known row's representation — and through it the target's — would depend on the + /// target's history in both directions. + [Fact] + public void KnownFutureRowsDoNotReadTargets() + { + manual_seed(7); + using var model = new T0Model(Tiny); + model.eval(); + using var _ = no_grad(); + + // Rows: target, known-future. Width 16 = 4 patches; context 2 patches, horizon 2. + const int w = 16, ctx = 8; + long[] dims = [2, w]; + var mask = zeros(dims, ScalarType.Int64); + mask[0, TensorIndex.Slice(ctx)] = MaskType.Withheld; + var groups = zeros(dims, ScalarType.Int64); + var types = zeros(dims, ScalarType.Int64); + types[1] = VariateType.Future; + + Tensor KnownRowOutput(float targetLevel) + { + var values = stack([full(w, targetLevel), arange(w, ScalarType.Float32)]); + var (scaled, _, _) = T0Preprocess.ScaleInput(values, mask, groups, types, Tiny); + // Perturb in scaled space so the scaler cannot absorb it. + scaled[0, TensorIndex.Slice(0, ctx)] = scaled[0, TensorIndex.Slice(0, ctx)] + targetLevel; + return model.Backbone(scaled, mask, groups, types)[1]; + } + + Assert.True(KnownRowOutput(0f).allclose(KnownRowOutput(3f))); + } + + /// Time attention is causal for targets: an output must not move when only + /// later steps change. + [Fact] + public void TargetOutputsAreCausalInTime() + { + manual_seed(7); + using var model = new T0Model(Tiny); + model.eval(); + using var _ = no_grad(); + + const int w = 16; + long[] dims = [1, w]; + var (mask, groups, types) = (zeros(dims, ScalarType.Int64), zeros(dims, ScalarType.Int64), + zeros(dims, ScalarType.Int64)); + var a = randn(dims); + var b = a.clone(); + b[0, TensorIndex.Slice(12)] = b[0, TensorIndex.Slice(12)] + 10f; // last patch only + + var qa = model.Backbone(a, mask, groups, types); + var qb = model.Backbone(b, mask, groups, types); + Assert.True(qa[0, TensorIndex.Slice(0, 3)].allclose(qb[0, TensorIndex.Slice(0, 3)])); + Assert.False(qa[0, 3].allclose(qb[0, 3])); + } + + /// A context that does not fill whole patches is left-padded, and padding must + /// be inert: whatever the pad holds, it is masked from attention and scaling alike. + [Fact] + public void AnyContextLengthIsAccepted() + { + var f = Forecaster(); + for (var context = 1; context <= 9; context++) + { + var series = Enumerable.Range(0, context + 3).Select(t => 10f + t).ToArray(); + var q = f.ForecastJoint(series, 1, context, 3); + Assert.Equal(3, q.GetLength(0)); + Assert.All(q.Cast(), x => Assert.True(double.IsFinite(x))); + } + } + + [Fact] + public void QuantilesAreMonotoneAndFinite() + { + var f = Forecaster(); + var q = f.ForecastJoint(Series(), 3, Context, Horizon, Known); + for (var h = 0; h < Horizon; h++) + for (var i = 0; i < q.GetLength(1); i++) + { + Assert.True(double.IsFinite(q[h, i])); + if (i > 0) Assert.True(q[h, i] >= q[h, i - 1]); + } + } + + [Fact] + public void MissingObservationsAreTolerated() + { + var f = Forecaster(); + var s = Series(); + s[3] = float.NaN; + s[2 * Total + Context + 1] = float.NaN; // a gap in the known future + Assert.All(f.ForecastJoint(s, 3, Context, Horizon, Known).Cast(), + x => Assert.True(double.IsFinite(x))); + } + + [Fact] + public void OhlcvaOfAnyLengthForecastsReturns() + { + var f = Forecaster(); + var bars = new float[13 * 6]; + for (var i = 0; i < 13; i++) + for (var c = 0; c < 6; c++) bars[i * 6 + c] = c < 4 ? 100f + i : 1000f; + var q = f.Forecast(bars, 5); + Assert.Equal(5, q.GetLength(0)); + Assert.Equal(Tiny.NumQuantiles, q.GetLength(1)); + } + + [Fact] + public void TheForecasterIsAnIQuantileForecaster() + { + IQuantileForecaster f = Forecaster(); + Assert.Equal(Tiny.Quantiles, f.Quantiles); + } + + [Fact] + public void ArgumentsAreValidated() + { + var f = Forecaster(); + var s = Series(); + Assert.Throws(() => f.ForecastJoint(s, 3, Context, Horizon, [true, false, true])); + Assert.Throws(() => f.ForecastJoint(s, 3, Context, Horizon, [false, true])); + Assert.Throws(() => f.ForecastJoint(s, 3, Context, 0)); + Assert.Throws(() => f.ForecastJoint(s, 3, Context, T0Forecaster.MaxHorizon + 1)); + Assert.Throws(() => f.ForecastJoint(s, 3, Context, Horizon + 1, Known)); + Assert.Throws(() => f.ForecastJoint(s, 3, 0, Horizon)); + Assert.Throws(() => f.ForecastJoint(s, 5, Context, Horizon)); + Assert.Throws(() => f.Forecast(new float[7], 1)); + } + + [Fact] + public void ConfigReadsBetaConventionsAndDefaultsToAlphas() + { + const string alpha = """{"embed_dim":512,"group_every_n":3,"mlp_hidden_dim":2048,"num_heads":8,"num_layers":24,"patch_size":32,"quantile_levels":[0.1,0.25,0.5,0.75,0.9],"scaler_use_arcsinh":true}"""; + const string beta = """{"embed_dim":1024,"group_every_n":3,"mlp_hidden_dim":2048,"num_heads":8,"num_layers":24,"patch_size":32,"quantile_levels":[0.1,0.5,0.9],"scaler_eps":0.01,"scaler_eps_mode":"std_clamp","scaler_use_arcsinh":true}"""; + + var a = T0Config.Parse(alpha); + Assert.Equal(0.1, a.ScalerEps); + Assert.False(a.ScalerStdClamp); + Assert.False(a.XposHalvesLayout); + Assert.True(a.ScalerZeroNoValue); + + var b = T0Config.Parse(beta); + Assert.Equal(0.01, b.ScalerEps); + Assert.True(b.ScalerStdClamp); + + Assert.Throws(() => T0Config.Parse(beta.Replace("std_clamp", "bogus"))); + } +} diff --git a/tests/Tsfm.Forecasting.T0.Tests/Tsfm.Forecasting.T0.Tests.csproj b/tests/Tsfm.Forecasting.T0.Tests/Tsfm.Forecasting.T0.Tests.csproj new file mode 100644 index 0000000..303b0e5 --- /dev/null +++ b/tests/Tsfm.Forecasting.T0.Tests/Tsfm.Forecasting.T0.Tests.csproj @@ -0,0 +1,19 @@ + + + + false + Tsfm.Forecasting.T0.Tests + + + + + + + + + + + + + + From c8e4fa2c6ff1fd23c61ecf3a3eb460fb87220b75 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 14:23:06 +0000 Subject: [PATCH 2/9] feat: batch forecasting across t0, TimesFM and Kronos; 0.2.0 IQuantileForecaster gains ForecastJointBatch and ForecastBatch over JointSeries / ReadOnlyMemory entries sharing a horizon, written [entry][step][quantile] into a caller-sized span. - t0 batches any entries in one pass, each its own attention group, shorter contexts left-padded and masked, cut into passes by MaxRowsPerPass. Checked against the reference in the parity harness. - TimesFM stacks same-shape entries on the model's batch axis. Single calls now run as a batch of one, bit-identical to 0.1.3. - KronosForecaster.InferBatch forecasts the final window of many series; Infer/InferPath share the generalised loop, bit-identical to 0.1.3. Also: clearer t0 weight installation (pinned, checksummed download commands; Load names them when files are missing), corrected wording on covariate-to-target information flow, and all packages move to 0.2.0. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01F1dTUhvq2UnFWdBmH7AA8G --- .github/workflows/ci.yml | 8 +- CHANGELOG.md | 32 +- Directory.Build.props | 2 +- README.md | 92 ++++- .../KronosForecaster.cs | 156 ++++++-- src/Tsfm.Forecasting.T0/T0Forecaster.cs | 333 +++++++++++++----- src/Tsfm.Forecasting.T0/T0Model.cs | 10 + .../TimesFmForecaster.cs | 317 ++++++++++++----- src/Tsfm.Forecasting/IQuantileForecaster.cs | 47 ++- .../BatchTests.cs | 78 ++++ tests/Tsfm.Forecasting.T0.Parity/Program.cs | 23 ++ tests/Tsfm.Forecasting.T0.Tests/BatchTests.cs | 116 ++++++ .../BatchTests.cs | 84 +++++ 13 files changed, 1072 insertions(+), 226 deletions(-) create mode 100644 tests/Tsfm.Forecasting.Kronos.IntegrationTests/BatchTests.cs create mode 100644 tests/Tsfm.Forecasting.T0.Tests/BatchTests.cs create mode 100644 tests/Tsfm.Forecasting.TimesFm.IntegrationTests/BatchTests.cs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1f5affd..aec810a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -43,10 +43,10 @@ jobs: - name: Test run: dotnet test tests/Tsfm.Forecasting.Kronos.Tests -c Release --nologo --no-build - # Vacuous here by design: the TimesFM checkpoint is ~1.2 GB and non-commercially - # licensed, so it is never fetched on a runner. Running them anyway proves they - # compile and start; the assertions only bite on a developer machine. - - name: TimesFM tests (self-skipping without weights) + # The checkpoint tests are vacuous here by design: the TimesFM checkpoint is ~1.2 GB + # and non-commercially licensed, so it is never fetched on a runner. The batch tests + # run on a random-init model and do bite. + - name: TimesFM tests (checkpoint tests self-skip) run: dotnet test tests/Tsfm.Forecasting.TimesFm.IntegrationTests -c Release --nologo # The invariant tests run on a tiny random-init model and need no weights; the diff --git a/CHANGELOG.md b/CHANGELOG.md index be56638..6c2783c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ # Changelog -## Unreleased +## 0.2.0 — unreleased ### Tsfm.Forecasting.T0 (Apache-2.0) — new @@ -22,15 +22,37 @@ extrapolation to levels off the trained grid, and the optional symmetric averaging, clipping and integer snapping. -### Tsfm.Forecasting +- Weights are not bundled. The README gives pinned download commands with checksums for + projects that only reference the package, and `Load` names them when files are missing. + +### Batching, across all three models + +- `IQuantileForecaster` (new, in `Tsfm.Forecasting`): the shared `ForecastJoint` / + `Forecast` / `Quantiles` shape, plus `ForecastJointBatch` and `ForecastBatch` over many + entries sharing a horizon, written `[entry][step][quantile]` into a caller-sized span + (`BatchOutputLength`). Entries are `JointSeries` / `ReadOnlyMemory`, so anchors + sliced from one history copy nothing. `TimesFmForecaster` and `T0Forecaster` implement it. +- t0 batches any entries in one pass — differing variates, contexts and known-future + patterns — each its own attention group, shorter contexts left-padded and masked. Entries + are ordered by context before being cut into passes of `MaxRowsPerPass` series. Verified + against the reference in the parity harness. +- TimesFM stacks entries of one shape on the model's batch axis, `MaxBatchPerPass` at a + time. The single calls now run as a batch of one; outputs are bit-identical to 0.1.3. +- `KronosForecaster.InferBatch` forecasts the final window of many series in one pass of + `batch × rollouts` rows; row `i` equals the last row `InferPath` gives for series `i` + alone, sampling included. `Infer` and `InferPath` share the refactored loop and are + bit-identical to 0.1.3. -- `IQuantileForecaster`: the shared `ForecastJoint` / `Forecast` / `Quantiles` shape. - `TimesFmForecaster` and `T0Forecaster` implement it, so one call site drives either. - Additive; TimesFM behaviour is unchanged. +### Packaging + +All six packages move to 0.2.0: the core gains the interface the model packages implement, +and a project reference packs an exact-version dependency. ### Samples and testing - `samples/GroceryCovariates` runs on TimesFM or t0 (`-- t0-beta`) through the interface. +- Batch tests for TimesFM and t0 run on small random-init models, so they bite in CI + without weights; Kronos's run with the integration tests. - `Tsfm.Forecasting.T0.Tests` checks structural invariants on a tiny random-init model, so they run in CI without weights: a target never reads its own future, known-future rows never read targets, time attention is causal, any context length is accepted. Behaviour diff --git a/Directory.Build.props b/Directory.Build.props index 10e3231..30e2302 100644 --- a/Directory.Build.props +++ b/Directory.Build.props @@ -21,7 +21,7 @@ - 0.1.3 + 0.2.0 Mark Hammond Copyright (c) 2026 Mark Hammond README.md diff --git a/README.md b/README.md index 73345b3..dd2f1d9 100644 --- a/README.md +++ b/README.md @@ -109,8 +109,8 @@ a channel left non-target would be handed its own future values. [t0](https://arxiv.org/abs/2609.24559) from The Forecasting Company is a decoder-style patch transformer alternating causal time attention with variate (group) attention, with native -known-future covariates. Both checkpoints are **Apache-2.0**. They are fetched rather than -packaged only because each exceeds nuget.org's 250 MB limit. +known-future covariates. Both checkpoints are **Apache-2.0**, but **the package carries no +weights**: install them yourself, as below. | Checkpoint | Params | Quantiles | Scaler eps | |---|---|---|---| @@ -120,9 +120,48 @@ packaged only because each exceeds nuget.org's 250 MB limit. ``` dotnet add package Tsfm.Forecasting.T0 dotnet add package TorchSharp-cpu -./scripts/fetch-t0-checkpoint.sh # both; or name one: t0-alpha | t0-beta ``` +### Installing the weights + +`T0Forecaster.Load(dir, device)` needs a directory holding two files from the Hugging Face +repository, unrenamed: + +``` +/config.json # architecture and scaler convention — required, not optional +/model.safetensors # 407 MB (alpha) or 1.0 GB (beta) +``` + +From a clone of this repository, the script fetches both checkpoints at the revisions the +port was verified against, and checks each payload's SHA-256: + +``` +./scripts/fetch-t0-checkpoint.sh # both, into checkpoints/t0-alpha and checkpoints/t0-beta +./scripts/fetch-t0-checkpoint.sh t0-beta # or one +``` + +From a project that only references the package, download them directly — pinning the +revision, since a refreshed checkpoint would silently change every forecast: + +``` +fetch() { # name, revision, sha256 + mkdir -p "models/$1" + for f in config.json model.safetensors; do + curl -fL -o "models/$1/$f" "https://huggingface.co/theforecastingcompany/$1/resolve/$2/$f" + done + echo "$3 models/$1/model.safetensors" | shasum -a 256 -c +} +fetch t0-alpha 741927847fe02367624749bb9290b2f09f128e7d 16c030d3fd70f06dc4238e9a8356e9b5a631d07f80f1bc76ba539991aed5897f +fetch t0-beta 941e522547ec63f0f2aad57089fa0f5b8498bdf8 a0fd8abd51275dd30afd21f4892c763ebd45e2cb1811eff26eed43f89a543f7d +``` + +`hf download theforecastingcompany/t0-beta config.json model.safetensors --revision +--local-dir models/t0-beta` does the same through the Hugging Face CLI. Deploy the directory +beside your application like any other data file; nothing is read from the environment or +a cache, and a missing file fails at `Load` with these instructions in the message. + +### Use + ```csharp using Tsfm.Forecasting.T0; @@ -136,7 +175,7 @@ double[,] r = t0.Forecast(ohlcva, horizon: 4); // forward r Both forecasters implement `IQuantileForecaster`, so one call site drives either; `Quantiles` names the columns, whose count differs per checkpoint. `samples/GroceryCovariates` -runs on either (`-- t0-beta`). +runs on either (`-- t0-beta`). Both also batch; see [Batching](#batching). What differs from TimesFM: @@ -145,14 +184,53 @@ What differs from TimesFM: evaluates at 8192. - **Up to 1024 steps per pass.** Every horizon patch is decoded at once. The reference continues past 1024 by an autoregressive rollout over quantile paths, which is not ported. -- **Known-future series cannot read the targets.** Group attention lets a known-future row - attend only to other known-future rows, so information flows from covariates into - targets and never back. They are scaled by global rather than causal statistics. +- **Information flows from covariates to targets, never back.** In variate attention a + target reads every series in its entry, but a known-future series reads only the + known-future series (itself included). A covariate's representation therefore never + depends on the target's history, and cannot smuggle it forward. Known-future series are + also scaled by global rather than causal statistics. - **Each checkpoint names its own scaler convention.** t0-beta's `config.json` carries `scaler_eps` and `scaler_eps_mode`; alpha's predates them, so their absence means alpha's values. Running beta under alpha's convention loads cleanly and silently degrades the forecast — which is why they are read rather than assumed. +## Batching + +Each forecaster has a batch form of its single call. Entries share a horizon and nothing +else — they never attend to one another — so a batched result equals the single call's, +exactly for TimesFM and Kronos and to float rounding for t0. What batching buys is fewer, +larger forward passes: one per anchor is the dominant cost of a backtest. + +| | Single | Batch | What shares a pass | +|---|---|---|---| +| t0 | `ForecastJoint`, `Forecast` | `ForecastJointBatch`, `ForecastBatch` | any entries, up to `MaxRowsPerPass` series; shorter contexts are left-padded and masked | +| TimesFM | `ForecastJoint`, `Forecast` | `ForecastJointBatch`, `ForecastBatch` | entries of one shape (variates, length, context), up to `MaxBatchPerPass` | +| Kronos | `Infer`, `InferPath` (windows along one series) | `InferBatch` (final window of many series) | `batch` windows × `rollouts` | + +The t0 and TimesFM batches are on `IQuantileForecaster`, take `JointSeries` entries (or +OHLCVA windows) as `ReadOnlyMemory`, and write `[entry][step][quantile]` into a span sized +by `BatchOutputLength`. Anchors along one history are slices of one array, so building a +backtest batch copies nothing: + +```csharp +var windows = anchors.Select(a => history.AsMemory((a - context + 1) * 6, context * 6)).ToArray(); +var returns = new double[forecaster.BatchOutputLength(windows.Length, horizon)]; +forecaster.ForecastBatch(windows, horizon, returns); +``` + +**Batching and rollouts.** The batch axis holds independent series; rollouts are samples +*within* one series. Kronos flattens both into one tensor batch of `series × rollouts` +rows, and a `RolloutAggregator` reduces only across the rollouts of one series, at one step +and channel — never across series. So the two compose without interaction, except in cost: +memory scales with the product, which is why Kronos's `batch` counts windows and leaves +rows to be `batch * rollouts`. Sampling is seeded per window from its anchor bar, so +results do not depend on how windows were grouped into passes; the corollary is that two +instruments anchored on the same bar draw the same uniforms (common random numbers). +TimesFM and t0 emit quantiles directly, with no rollouts and nothing to aggregate. t0's +upstream continuation past 1024 steps (not ported) would be the same pattern again: each +requested quantile becomes a path, paths are folded into the batch as extra groups, and +they are reduced per series to quantiles. + ## Choose your own backend Neither model package propagates a native runtime, because the right one is platform- diff --git a/src/Tsfm.Forecasting.Kronos/KronosForecaster.cs b/src/Tsfm.Forecasting.Kronos/KronosForecaster.cs index 9b4093c..15f88d2 100644 --- a/src/Tsfm.Forecasting.Kronos/KronosForecaster.cs +++ b/src/Tsfm.Forecasting.Kronos/KronosForecaster.cs @@ -150,29 +150,87 @@ public bool InferPath(ReadOnlySpan ohlcva, ReadOnlySpan barTimeMs, public static int PathLength(int barCount, int contextBars, int horizon) => OutputCount(barCount, contextBars) * horizon * Channels; +#pragma warning disable CS1573 + /// + /// + /// One forecast per series, anchored on each series' final bar — instruments at the + /// latest bar, say — batched across series rather than along one. Row i equals + /// the last row would produce for series[i] alone. + /// + /// Rollouts and series share the forward pass: series of + /// each make batch * rollouts rows. Each output reduces + /// only its own series' rollouts. + /// + /// Sampling is seeded by the anchor bar's timestamp alone, so series anchored on + /// the same bar draw the same uniforms — common random numbers. Comparisons across + /// instruments are then less noisy than independent draws would make them, but their + /// sampling errors are correlated, not independent. + /// + /// Each with at least bars. + /// One row per series. + /// One row per series. + /// One row per series, or empty. + /// Receives series x horizon x 6, or may be empty. + /// Reduces each channel's rollouts at each step; null is the + /// mean. Only applied when is requested. + /// Series per forward pass. + public void InferBatch(ReadOnlySpan series, + Span lean, Span upCount, Span dispersion, Span path, + int contextBars, int horizon, int rollouts, bool greedy, + float temperature, float topP, RolloutAggregator? aggregate = null, int batch = 8) + { + for (var i = 0; i < series.Length; i++) + { + var (o, t) = (series[i].Ohlcva.Length, series[i].BarTimeMs.Length); + if (o != t * Channels) + throw new ArgumentException($"series[{i}]: ohlcva must be L x 6 for its {t} bars", nameof(series)); + if (t < contextBars) + throw new ArgumentException( + $"series[{i}]: {t} bars cannot fill a context of {contextBars}", nameof(series)); + } + Drive(new Windows(series, contextBars), series.Length, lean, upCount, dispersion, path, + path.IsEmpty ? null : aggregate ?? RolloutAggregators.Mean, + contextBars, horizon, rollouts, greedy, temperature, topP, batch); + } +#pragma warning restore CS1573 + private bool InferCore(ReadOnlySpan ohlcva, ReadOnlySpan barTimeMs, Span lean, Span upCount, Span dispersion, Span path, RolloutAggregator? aggregate, int contextBars, int horizon, int rollouts, bool greedy, float temperature, float topP, int batch = 8) { - var k = contextBars; - var h = horizon; var length = barTimeMs.Length; if (ohlcva.Length != length * Channels) throw new ArgumentException("ohlcva must be L x 6", nameof(ohlcva)); + var count = OutputCount(length, contextBars); + if (contextBars <= MaxContext && count <= 0) return false; + Drive(new Windows(ohlcva, barTimeMs, contextBars), count, lean, upCount, dispersion, path, aggregate, + contextBars, horizon, rollouts, greedy, temperature, topP, batch); + return true; + } + + /// Evaluate windows, per + /// forward pass, writing output row i for window i. + private void Drive(in Windows source, int count, + Span lean, Span upCount, Span dispersion, Span path, + RolloutAggregator? aggregate, + int contextBars, int horizon, int rollouts, bool greedy, + float temperature, float topP, int batch) + { + var k = contextBars; + var h = horizon; + // Upstream truncates silently, which reads as a working call that quietly ignored // the bars beyond the limit. Refuse instead, and say what the limit is. if (k > MaxContext) throw new ArgumentOutOfRangeException(nameof(contextBars), $"{CheckpointName} attends over at most {MaxContext} bars; got {k}."); - var count = OutputCount(length, k); - if (count <= 0) return false; if (lean.Length != count || upCount.Length != count) throw new ArgumentException( - $"lean and upCount must each hold exactly {count} rows (= L - K + 1); " + + $"lean and upCount must each hold exactly {count} rows; " + $"got {lean.Length}/{upCount.Length}. Use {nameof(OutputCount)}."); // Dispersion is diagnostic. An empty span skips the standard deviation entirely. @@ -189,8 +247,8 @@ private bool InferCore(ReadOnlySpan ohlcva, ReadOnlySpan barTimeMs, if (!dispersion.IsEmpty) dispersion.Fill(float.NaN); if (!path.IsEmpty) path.Fill(float.NaN); upCount.Clear(); + if (count == 0) return; - var stamps = ArrayPool.Shared.Rent(length * 5); // Size scratch for a full batch and reuse it across the slice; per-batch // allocation dominates a materialisation-sized run. var maxRows = batch * rollouts; @@ -206,24 +264,19 @@ private bool InferCore(ReadOnlySpan ohlcva, ReadOnlySpan barTimeMs, try { - BuildStamps(barTimeMs, stamps.AsSpan(0, length * 5)); - var stepMs = length > 1 ? barTimeMs[1] - barTimeMs[0] : 0L; - using var _ = no_grad(); for (var start = 0; start < count; start += batch) { var take = Math.Min(batch, count - start); - EvaluateBatch(ohlcva, barTimeMs, stamps.AsSpan(0, length * 5), stepMs, start, take, + EvaluateBatch(source, start, take, k, h, rollouts, greedy, temperature, topP, lean, upCount, dispersion, path, aggregate, samples.AsSpan(0, rollouts), per.AsSpan(0, rollouts), x, stamp, mean, scale, anchor, uniforms, pick); } - return true; } finally { - ArrayPool.Shared.Return(stamps); ArrayPool.Shared.Return(x); ArrayPool.Shared.Return(stamp); ArrayPool.Shared.Return(mean); @@ -236,9 +289,49 @@ private bool InferCore(ReadOnlySpan ohlcva, ReadOnlySpan barTimeMs, } } + /// + /// Where window i lives: either the i-th sliding window along one series, + /// or the final window of the i-th of many. Either way it reads K bars from + /// of , anchored on the last of them. + /// + private readonly ref struct Windows + { + private readonly ReadOnlySpan _ohlcva; + private readonly ReadOnlySpan _times; + private readonly ReadOnlySpan _series; + private readonly int _k; + private readonly bool _many; + + public Windows(ReadOnlySpan ohlcva, ReadOnlySpan times, int k) + { + _ohlcva = ohlcva; + _times = times; + _k = k; + } + + public Windows(ReadOnlySpan series, int k) + { + _series = series; + _k = k; + _many = true; + } + + public ReadOnlySpan Ohlcva(int i) => _many ? _series[i].Ohlcva.Span : _ohlcva; + + public ReadOnlySpan Times(int i) => _many ? _series[i].BarTimeMs.Span : _times; + + public int Origin(int i) => _many ? _series[i].BarTimeMs.Length - _k : i; + + /// Bar interval, from the series' first spacing. + public long StepMs(int i) + { + var t = Times(i); + return t.Length > 1 ? t[1] - t[0] : 0L; + } + } + private void EvaluateBatch( - ReadOnlySpan ohlcva, ReadOnlySpan barTimeMs, ReadOnlySpan stamps, long stepMs, - int start, int take, int k, int h, int rollouts, bool greedy, float temperature, float topP, + in Windows source, int start, int take, int k, int h, int rollouts, bool greedy, float temperature, float topP, Span lean, Span upCount, Span dispersion, Span path, RolloutAggregator? aggregate, Span samples, Span per, float[] x, float[] stamp, float[] mean, float[] scale, float[] anchor, @@ -251,10 +344,20 @@ private void EvaluateBatch( // the model reads. for (var w = 0; w < take; w++) { - var origin = start + w; + var ohlcva = source.Ohlcva(start + w); + var times = source.Times(start + w); + var origin = source.Origin(start + w); + var stepMs = source.StepMs(start + w); NormaliseWindow(ohlcva, origin, k, mean.AsSpan(w * Channels, Channels), scale.AsSpan(w * Channels, Channels)); anchor[w] = ohlcva[(origin + k - 1) * Channels + 3]; // close of the anchor bar + // Calendar stamps for the window and its extrapolated horizon, written once and + // shared by every rollout. + var first = stamp.AsSpan(w * rollouts * (k + h) * 5, (k + h) * 5); + for (var t = 0; t < k; t++) WriteStamp(first, t, times[origin + t]); + for (var j = 0; j < h; j++) // forward stamps: extrapolated + WriteStamp(first, k + j, times[origin + k - 1] + stepMs * (j + 1)); + for (var r = 0; r < rollouts; r++) { var row = w * rollouts + r; @@ -264,9 +367,7 @@ private void EvaluateBatch( var v = (ohlcva[(origin + t) * Channels + c] - mean[w * Channels + c]) / (scale[w * Channels + c] + 1e-5f); x[(row * k + t) * Channels + c] = Math.Clamp(v, -5f, 5f); } - stamps.Slice(origin * 5, k * 5).CopyTo(stamp.AsSpan(row * (k + h) * 5, k * 5)); - for (var j = 0; j < h; j++) // forward stamps: extrapolated - WriteStamp(stamp, row * (k + h) + k + j, barTimeMs[origin + k - 1] + stepMs * (j + 1)); + if (r > 0) first.CopyTo(stamp.AsSpan(row * (k + h) * 5, (k + h) * 5)); } } @@ -281,8 +382,7 @@ private void EvaluateBatch( // Draw before the loop: sampling must be a pure function of the bar, not of how // windows were batched. - if (!greedy) DrawUniforms(barTimeMs, start, take, k, rollouts, h, uniforms); - + if (!greedy) DrawUniforms(source, start, take, k, rollouts, h, uniforms); for (var step = 0; step < h; step++) { var len = k + step; @@ -377,12 +477,12 @@ private static void NormaliseWindow( /// window's anchor bar. SplitMix64, so the stream cannot shift with a framework /// version. private static void DrawUniforms( - ReadOnlySpan barTimeMs, int start, int take, int k, int rollouts, int h, float[] u) + in Windows source, int start, int take, int k, int rollouts, int h, float[] u) { var perWindow = rollouts * h * 2; for (var w = 0; w < take; w++) { - var state = (ulong)SeedFor(barTimeMs[start + w + k - 1]); + var state = (ulong)SeedFor(source.Times(start + w)[source.Origin(start + w) + k - 1]); for (var i = 0; i < perWindow; i++) { state += 0x9E3779B97F4A7C15UL; @@ -432,11 +532,6 @@ private static Tensor Pick(Tensor logits, bool greedy, float temperature, float /// Calendar features: minute, hour, weekday, day-of-month, month. /// Weekday is Monday-zero, as pandas. is Sunday-zero; /// using it directly shifts every weekday embedding by one, silently. - private static void BuildStamps(ReadOnlySpan barTimeMs, Span stamps) - { - for (var i = 0; i < barTimeMs.Length; i++) WriteStamp(stamps, i, barTimeMs[i]); - } - public static void WriteStamp(Span dst, int row, long unixMs) { var t = DateTimeOffset.FromUnixTimeMilliseconds(unixMs).UtcDateTime; @@ -471,3 +566,10 @@ public void Dispose() _model.Dispose(); } } + +/// One series for : row-major +/// [L x 6] bars and their close times. Memory rather than spans so series can sit in +/// an array; slicing a larger buffer copies nothing. +/// Row-major [L x 6]: open, high, low, close, volume, amount. +/// Close time per bar, Unix milliseconds, [L]. +public readonly record struct KronosSeries(ReadOnlyMemory Ohlcva, ReadOnlyMemory BarTimeMs); diff --git a/src/Tsfm.Forecasting.T0/T0Forecaster.cs b/src/Tsfm.Forecasting.T0/T0Forecaster.cs index 686679e..42b5a9f 100644 --- a/src/Tsfm.Forecasting.T0/T0Forecaster.cs +++ b/src/Tsfm.Forecasting.T0/T0Forecaster.cs @@ -1,3 +1,5 @@ +using System.Runtime.InteropServices; + using static TorchSharp.torch; namespace Tsfm.Forecasting.T0; @@ -12,8 +14,9 @@ namespace Tsfm.Forecasting.T0; /// /// Every series not marked known is a forecast target, and targets are forecast /// jointly — so a channel cannot be handed its own future. Known-future series are -/// conditioned on but not forecast, and by construction cannot read the targets: their -/// group attention reads only other known-future rows. +/// conditioned on but not forecast, and by construction cannot read the targets: in +/// variate attention a target reads every series of its entry, but a known-future series +/// reads only known-future series, itself included. /// /// Inference is not thread safe. public sealed class T0Forecaster(T0Model model, T0Config config, Device device) : IQuantileForecaster @@ -63,52 +66,13 @@ public static T0Forecaster Load(string checkpointDir, Device device) ReadOnlySpan series, int variates, int contextBars, int horizon, ReadOnlySpan knownFuture = default, int targetVariate = 0) { - if (variates < 1) - throw new ArgumentOutOfRangeException(nameof(variates), "at least one series is required"); - if (series.Length % variates != 0) - throw new ArgumentException( - $"series length {series.Length} is not divisible by {variates} variates", nameof(series)); - var total = series.Length / variates; ValidateHorizon(horizon); - if (contextBars < 1 || contextBars > total) - throw new ArgumentOutOfRangeException(nameof(contextBars), - $"contextBars must lie in 1..{total}; got {contextBars}"); - if (targetVariate < 0 || targetVariate >= variates) - throw new ArgumentOutOfRangeException(nameof(targetVariate)); - if (!knownFuture.IsEmpty && knownFuture.Length != variates) - throw new ArgumentException( - $"knownFuture must hold {variates} flags or be empty", nameof(knownFuture)); - if (!knownFuture.IsEmpty && knownFuture[targetVariate]) - throw new ArgumentException( - "the target variate cannot be known into the future", nameof(knownFuture)); - var anyKnown = !knownFuture.IsEmpty && knownFuture.Contains(true); - if (anyKnown && horizon > total - contextBars) - throw new ArgumentOutOfRangeException(nameof(horizon), - $"horizon {horizon} exceeds the {total - contextBars} known steps supplied after the context"); - - // Targets first, then known-future rows: the reference's row order. Attention is - // indifferent to it, but summation order is not, and parity is checked bit-close. - var order = new int[variates]; - var known = new bool[variates]; - var rows = 0; - for (var v = 0; v < variates; v++) - if (knownFuture.IsEmpty || !knownFuture[v]) order[rows++] = v; - for (var v = 0; v < variates; v++) - if (!knownFuture.IsEmpty && knownFuture[v]) { known[rows] = true; order[rows++] = v; } - - var layout = new Layout(config.PatchSize, variates, contextBars, horizon); - var targetRow = Array.IndexOf(order, targetVariate); - for (var r = 0; r < variates; r++) - { - var src = series.Slice(order[r] * total, total); - layout.Row(r, src[..contextBars], known[r] ? src.Slice(contextBars, horizon) : default); - } + var total = ValidateJoint(series.Length, variates, contextBars, horizon, knownFuture, targetVariate, ""); + var layout = new Layout(config.PatchSize, contextBars, horizon, variates); + int[] targets = [AddJoint(layout, series, total, variates, contextBars, knownFuture, targetVariate)]; - var head = Run(layout, targetRow); var result = new double[horizon, config.NumQuantiles]; - for (var h = 0; h < horizon; h++) - for (var q = 0; q < config.NumQuantiles; q++) - result[h, q] = head[h * config.NumQuantiles + q]; + Run(layout, targets, [0], MemoryMarshal.CreateSpan(ref result[0, 0], result.Length)); return result; } @@ -123,26 +87,203 @@ public static T0Forecaster Load(string checkpointDir, Device device) /// [horizon, numQuantiles] of cumulative return relative to the anchor close. public double[,] Forecast(ReadOnlySpan ohlcva, int horizon) { - if (ohlcva.Length == 0 || ohlcva.Length % Channels != 0) - throw new ArgumentException($"ohlcva must hold whole bars of {Channels} channels", nameof(ohlcva)); ValidateHorizon(horizon); - var bars = ohlcva.Length / Channels; + ValidateBars(ohlcva.Length, ""); + var layout = new Layout(config.PatchSize, ohlcva.Length / Channels, horizon, Channels); + int[] targets = [AddBars(layout, ohlcva)]; + + var result = new double[horizon, config.NumQuantiles]; + var flat = MemoryMarshal.CreateSpan(ref result[0, 0], result.Length); + Run(layout, targets, [0], flat); + ToReturns(flat, ohlcva); + return result; + } + + /// Rows — series, summed over entries — attended in one forward pass. A batch + /// is split into passes of at most this many, since group attention costs grow with the + /// square of the rows in a pass. An entry wider than this still runs, alone. + public int MaxRowsPerPass { get; set; } = 256; + + /// + public int BatchOutputLength(int entries, int horizon) => entries * horizon * config.NumQuantiles; - var layout = new Layout(config.PatchSize, Channels, bars, horizon); - var column = new float[bars]; + /// + /// over many independent entries in as few passes as + /// allows. Each entry is its own group, so entries never + /// attend to one another; contexts may differ in length, each being left-padded to the + /// longest in its pass and the padding masked. + /// + /// Entries are ordered by context length before being split into passes, so a + /// short entry is not padded out to a long one's width. Results are written in the + /// caller's order regardless. + /// + /// Entries, each with the constraints of . + /// Steps ahead, shared, 1... + /// Receives [entry][step][quantile]; sized by + /// . + public void ForecastJointBatch(ReadOnlySpan batch, int horizon, Span quantiles) + { + ValidateHorizon(horizon); + ValidateDestination(quantiles.Length, batch.Length, horizon, nameof(quantiles)); + var totals = new int[batch.Length]; + var contexts = new int[batch.Length]; + var widths = new int[batch.Length]; + for (var i = 0; i < batch.Length; i++) + { + var e = batch[i]; + totals[i] = ValidateJoint(e.Series.Length, e.Variates, e.ContextBars, horizon, + e.KnownFuture.Span, e.TargetVariate, $"batch[{i}]: "); + (contexts[i], widths[i]) = (e.ContextBars, e.Variates); + } + + foreach (var pass in Passes(contexts, widths)) + { + var layout = new Layout(config.PatchSize, pass.Max(i => contexts[i]), horizon, pass.Sum(i => widths[i])); + var targets = new int[pass.Length]; + for (var j = 0; j < pass.Length; j++) + { + var e = batch[pass[j]]; + targets[j] = AddJoint(layout, e.Series.Span, totals[pass[j]], e.Variates, e.ContextBars, + e.KnownFuture.Span, e.TargetVariate); + } + Run(layout, targets, pass, quantiles); + } + } + + /// + /// over many OHLCVA windows — instruments, or anchors along one + /// history — batched as is. Windows may differ in + /// length; slicing one history into overlapping windows copies nothing. + /// + /// One row-major [L x 6] window per entry. + /// Steps ahead, shared, 1... + /// Receives [entry][step][quantile] of cumulative return + /// relative to each window's anchor close; sized by . + public void ForecastBatch(ReadOnlySpan> ohlcva, int horizon, Span returns) + { + ValidateHorizon(horizon); + ValidateDestination(returns.Length, ohlcva.Length, horizon, nameof(returns)); + var contexts = new int[ohlcva.Length]; + for (var i = 0; i < ohlcva.Length; i++) + { + ValidateBars(ohlcva[i].Length, $"ohlcva[{i}]: "); + contexts[i] = ohlcva[i].Length / Channels; + } + + var widths = Enumerable.Repeat(Channels, ohlcva.Length).ToArray(); + foreach (var pass in Passes(contexts, widths)) + { + var layout = new Layout(config.PatchSize, pass.Max(i => contexts[i]), horizon, pass.Length * Channels); + var targets = new int[pass.Length]; + for (var j = 0; j < pass.Length; j++) targets[j] = AddBars(layout, ohlcva[pass[j]].Span); + Run(layout, targets, pass, returns); + } + + var block = horizon * config.NumQuantiles; + for (var i = 0; i < ohlcva.Length; i++) ToReturns(returns.Slice(i * block, block), ohlcva[i].Span); + } + + /// Entry indices grouped into passes: ordered by context so a pass pads to + /// similar widths, and cut wherever the next entry would exceed the row budget. + private List Passes(int[] contexts, int[] rows) + { + var order = Enumerable.Range(0, contexts.Length).OrderBy(i => contexts[i]).ToArray(); + var passes = new List(); + var (from, used) = (0, 0); + for (var j = 0; j < order.Length; j++) + { + if (j > from && used + rows[order[j]] > MaxRowsPerPass) + { + passes.Add(order[from..j]); + (from, used) = (j, 0); + } + used += rows[order[j]]; + } + if (from < order.Length) passes.Add(order[from..]); + return passes; + } + + /// Add one joint entry as its own group: targets first, then known-future rows, + /// the reference's row order — attention is indifferent to it, but summation order is + /// not, and parity is checked bit-close. Returns the target's row. + private static int AddJoint(Layout layout, ReadOnlySpan series, int total, int variates, + int contextBars, ReadOnlySpan knownFuture, int targetVariate) + { + var group = layout.NextGroup(); + var target = -1; + for (var pass = 0; pass < 2; pass++) + for (var v = 0; v < variates; v++) + { + var known = !knownFuture.IsEmpty && knownFuture[v]; + if (known != (pass == 1)) continue; + var src = series.Slice(v * total, total); + var row = layout.AddRow(group, src[..contextBars], known ? src.Slice(contextBars, layout.Horizon) : default); + if (v == targetVariate) target = row; + } + return target; + } + + /// Add one OHLCVA window as a group of six targets; returns the close's row. + private static int AddBars(Layout layout, ReadOnlySpan ohlcva) + { + var group = layout.NextGroup(); + var bars = ohlcva.Length / Channels; + Span column = bars <= 1024 ? stackalloc float[bars] : new float[bars]; + var close = -1; for (var c = 0; c < Channels; c++) { for (var t = 0; t < bars; t++) column[t] = ohlcva[t * Channels + c]; - layout.Row(c, column, default); + var row = layout.AddRow(group, column, default); + if (c == CloseChannel) close = row; } + return close; + } - var head = Run(layout, CloseChannel); - var anchorClose = ohlcva[(bars - 1) * Channels + CloseChannel]; - var result = new double[horizon, config.NumQuantiles]; - for (var h = 0; h < horizon; h++) - for (var q = 0; q < config.NumQuantiles; q++) - result[h, q] = head[h * config.NumQuantiles + q] / anchorClose - 1.0; - return result; + private static void ToReturns(Span prices, ReadOnlySpan ohlcva) + { + var anchorClose = ohlcva[^(Channels - CloseChannel)]; + for (var i = 0; i < prices.Length; i++) prices[i] = prices[i] / anchorClose - 1.0; + } + + private static int ValidateJoint(int length, int variates, int contextBars, int horizon, + ReadOnlySpan knownFuture, int targetVariate, string where) + { + if (variates < 1) + throw new ArgumentOutOfRangeException(nameof(variates), $"{where}at least one series is required"); + if (length % variates != 0) + throw new ArgumentException( + $"{where}series length {length} is not divisible by {variates} variates", "series"); + var total = length / variates; + if (contextBars < 1 || contextBars > total) + throw new ArgumentOutOfRangeException(nameof(contextBars), + $"{where}contextBars must lie in 1..{total}; got {contextBars}"); + if (targetVariate < 0 || targetVariate >= variates) + throw new ArgumentOutOfRangeException(nameof(targetVariate), $"{where}no variate {targetVariate}"); + if (!knownFuture.IsEmpty && knownFuture.Length != variates) + throw new ArgumentException( + $"{where}knownFuture must hold {variates} flags or be empty", nameof(knownFuture)); + if (!knownFuture.IsEmpty && knownFuture[targetVariate]) + throw new ArgumentException( + $"{where}the target variate cannot be known into the future", nameof(knownFuture)); + if (!knownFuture.IsEmpty && knownFuture.Contains(true) && horizon > total - contextBars) + throw new ArgumentOutOfRangeException(nameof(horizon), + $"{where}horizon {horizon} exceeds the {total - contextBars} known steps supplied after the context"); + return total; + } + + private static void ValidateBars(int length, string where) + { + if (length == 0 || length % Channels != 0) + throw new ArgumentException($"{where}ohlcva must hold whole bars of {Channels} channels", "ohlcva"); + } + + private void ValidateDestination(int length, int entries, int horizon, string name) + { + if (length != BatchOutputLength(entries, horizon)) + throw new ArgumentException( + $"{name} must hold exactly {BatchOutputLength(entries, horizon)} elements " + + $"(entries x horizon x {config.NumQuantiles} quantiles); got {length}. Use {nameof(BatchOutputLength)}.", + name); } private static void ValidateHorizon(int horizon) @@ -151,9 +292,9 @@ private static void ValidateHorizon(int horizon) throw new ArgumentOutOfRangeException(nameof(horizon), $"horizon must lie in 1..{MaxHorizon}"); } - /// One forward pass; the target row's first horizon steps as - /// [horizon * quantiles], in data units. - private float[] Run(Layout layout, int targetRow) + /// One forward pass. Writes target row j's first horizon steps, + /// in data units, to at slot slots[j]. + private void Run(Layout layout, int[] targetRows, ReadOnlySpan slots, Span destination) { using var scope = NewDisposeScope(); using var _ = no_grad(); @@ -170,59 +311,81 @@ private float[] Run(Layout layout, int targetRow) // The output at patch i forecasts patch i + 1, so the horizon starts at the last // context patch. - var p = config.PatchSize; - var forecast = data[targetRow] - .narrow(0, layout.ContextPatches - 1, layout.ForecastPatches) - .reshape(layout.ForecastPatches * p, config.NumQuantiles) - .narrow(0, 0, layout.Horizon); + var (p, q, h) = (config.PatchSize, config.NumQuantiles, layout.Horizon); + var forecast = data.index_select(0, tensor(targetRows.Select(r => (long)r).ToArray()).to(device)) + .narrow(1, layout.ContextPatches - 1, layout.ForecastPatches) + .reshape(targetRows.Length, layout.ForecastPatches * p, q) + .narrow(1, 0, h); forecast = nan_to_num(forecast, 0.0, 0.0, 0.0); // the reference zeroes non-finite output - return forecast.cpu().contiguous().data().ToArray(); + var host = forecast.cpu().contiguous().data(); + + var block = h * q; + for (var j = 0; j < targetRows.Length; j++) + { + var dst = destination.Slice(slots[j] * block, block); + for (var i = 0; i < block; i++) dst[i] = host[j * block + i]; + } } /// - /// The per-step input the model reads: the context left-padded onto the patch grid, then - /// the horizon rounded up to whole patches. Over the horizon a target is withheld and a - /// known-future row carries its values; padding is inert (value 0, group and role -1). + /// The per-step input the model reads, for one or many groups. Every row is right-aligned + /// on the forecast start, which sits on a patch boundary: its context is left-padded to + /// the longest in the pass, rounded up to whole patches, and the horizon rounded up after + /// it. So a row's patch grid is the same alone or batched. Over the horizon a target is + /// withheld and a known-future row carries its values; padding is inert (value 0, group + /// and role -1). /// private sealed class Layout { - public readonly int Rows, Width, PadLeft, Context, Horizon, ContextPatches, ForecastPatches; + public readonly int Rows, Width, ContextWidth, Horizon, ContextPatches, ForecastPatches; public readonly float[] Values; public readonly long[] Mask, Groups, Types; + private int _row; + private long _group; - public Layout(int patch, int rows, int context, int horizon) + public Layout(int patch, int maxContext, int horizon, int rows) { - (Rows, Context, Horizon) = (rows, context, horizon); - PadLeft = (patch - context % patch) % patch; + (Rows, Horizon) = (rows, horizon); + ContextPatches = (maxContext + patch - 1) / patch; ForecastPatches = (horizon + patch - 1) / patch; - ContextPatches = (PadLeft + context) / patch; - Width = PadLeft + context + ForecastPatches * patch; + ContextWidth = ContextPatches * patch; + Width = ContextWidth + ForecastPatches * patch; Values = new float[rows * Width]; Mask = new long[rows * Width]; Groups = new long[rows * Width]; Types = new long[rows * Width]; } - /// Row index. + public long NextGroup() => _group++; + + /// Rows sharing a group attend to one another. /// Observed steps. /// Known future values, or empty for a target. - public void Row(int row, ReadOnlySpan context, ReadOnlySpan future) + /// The row written. + public int AddRow(long group, ReadOnlySpan context, ReadOnlySpan future) { + var row = _row++; var isFuture = !future.IsEmpty; + var pad = ContextWidth - context.Length; var o = row * Width; for (var i = 0; i < Width; i++) { - var pad = i < PadLeft; - Groups[o + i] = pad ? -1 : 0; - Types[o + i] = pad ? -1 : isFuture ? VariateType.Future : VariateType.Target; - if (pad) { Mask[o + i] = MaskType.Pad; continue; } + if (i < pad) + { + (Mask[o + i], Groups[o + i], Types[o + i]) = (MaskType.Pad, -1, -1); + continue; + } + Groups[o + i] = group; + Types[o + i] = isFuture ? VariateType.Future : VariateType.Target; - var t = i - PadLeft; - if (t < Context) Write(o + i, context[t]); + var t = i - pad; + var f = t - context.Length; + if (f < 0) Write(o + i, context[t]); else if (!isFuture) Mask[o + i] = MaskType.Withheld; - else if (t - Context < Horizon) Write(o + i, future[t - Context]); + else if (f < Horizon) Write(o + i, future[f]); else Mask[o + i] = MaskType.Pad; // the horizon's last patch, beyond what was asked } + return row; } private void Write(int at, float value) diff --git a/src/Tsfm.Forecasting.T0/T0Model.cs b/src/Tsfm.Forecasting.T0/T0Model.cs index 8b2a3be..803adef 100644 --- a/src/Tsfm.Forecasting.T0/T0Model.cs +++ b/src/Tsfm.Forecasting.T0/T0Model.cs @@ -122,6 +122,16 @@ public static Tensor Monotone(Tensor x) /// matched: an unbound tensor keeps its random init and still forecasts plausibly. public static T0Model Load(string checkpointDir, Device device) { + // Weights are not packaged, so a missing directory is the likeliest first failure; + // say how to get them rather than surfacing a bare FileNotFoundException. + foreach (var file in (string[])["config.json", "model.safetensors"]) + if (!File.Exists(Path.Combine(checkpointDir, file))) + throw new FileNotFoundException( + $"{Path.Combine(checkpointDir, file)} not found. t0 weights are not bundled: fetch them with " + + "scripts/fetch-t0-checkpoint.sh, or download config.json and model.safetensors from " + + "https://huggingface.co/theforecastingcompany/t0-alpha (or t0-beta) into that directory.", + Path.Combine(checkpointDir, file)); + var cfg = T0Config.Load(Path.Combine(checkpointDir, "config.json")); using var stream = File.OpenRead(Path.Combine(checkpointDir, "model.safetensors")); return Load(cfg, stream, device, checkpointDir); diff --git a/src/Tsfm.Forecasting.TimesFm/TimesFmForecaster.cs b/src/Tsfm.Forecasting.TimesFm/TimesFmForecaster.cs index 7a8187a..39dd93e 100644 --- a/src/Tsfm.Forecasting.TimesFm/TimesFmForecaster.cs +++ b/src/Tsfm.Forecasting.TimesFm/TimesFmForecaster.cs @@ -1,3 +1,5 @@ +using System.Runtime.InteropServices; + using static TorchSharp.torch; namespace Tsfm.Forecasting.TimesFm; @@ -58,77 +60,13 @@ public static TimesFmForecaster Load(string checkpointDir, Device device) ReadOnlySpan series, int variates, int contextBars, int horizon, ReadOnlySpan knownFuture = default, int targetVariate = 0) { - var p = config.InputPatchLen; - if (variates < 1 || variates > MaxVariates) - throw new ArgumentOutOfRangeException(nameof(variates), - $"variates must lie in 1..{MaxVariates}; got {variates}"); - if (series.Length % variates != 0) - throw new ArgumentException( - $"series length {series.Length} is not divisible by {variates} variates", nameof(series)); - var total = series.Length / variates; - - // The known future is gathered by rolling whole patches forward, so the series is - // patched and the context must end on a patch boundary, leaving at least one patch - // after it to roll into. - if (total % p != 0) - throw new ArgumentException( - $"each series must hold a multiple of the patch length {p}; got {total}", nameof(series)); - if (contextBars % p != 0 || contextBars < p || contextBars >= total) - throw new ArgumentException( - $"contextBars ({contextBars}) must be a positive multiple of the patch length {p} " + - $"and leave at least one patch of the {total} supplied", nameof(contextBars)); - if (horizon > total - contextBars) - throw new ArgumentOutOfRangeException(nameof(horizon), - $"horizon {horizon} exceeds the {total - contextBars} steps supplied after the context"); - if (horizon < 1 || horizon > config.OutputPatchLen) - throw new ArgumentOutOfRangeException(nameof(horizon), - $"horizon must lie in 1..{config.OutputPatchLen}"); - if (targetVariate < 0 || targetVariate >= variates) - throw new ArgumentOutOfRangeException(nameof(targetVariate)); - if (!knownFuture.IsEmpty && knownFuture.Length != variates) - throw new ArgumentException( - $"knownFuture must hold {variates} flags or be empty", nameof(knownFuture)); - if (!knownFuture.IsEmpty && knownFuture[targetVariate]) - throw new ArgumentException( - "the target variate cannot be known into the future", nameof(knownFuture)); - - var n = total / p; - var anchorPatch = contextBars / p - 1; // last patch made entirely of observed steps - using var scope = NewDisposeScope(); - using var _ = no_grad(); - - var flat = zeros(variates * total, ScalarType.Float32); - var dst = flat.data(); - for (var i = 0; i < variates * total; i++) dst[i] = series[i]; - var values = flat.view(variates, n, p).unsqueeze(0).contiguous().to(device); - - // Mask the future region of every variate whose future is not supplied, and mark - // as target exactly those variates. Only a target has its future slots withheld, - // so a variate left unmarked would be handed values it should not see. - var masks = zeros([1, variates, n, p], ScalarType.Bool, device); - var isTarget = ones([1, variates, n], ScalarType.Bool, device); - for (var v = 0; v < variates; v++) - { - var known = !knownFuture.IsEmpty && knownFuture[v]; - if (known) { isTarget[0, v] = false; continue; } - for (var j = anchorPatch + 1; j < n; j++) masks[0, v, j] = true; - } - - var logits = model.forward(values, masks, isTarget) - .view(1, variates, n, config.OutputPatchLen, config.NumQuantiles); - - var (_, mu, sigma) = TimesFmPreprocess.RunningStats( - nan_to_num(values, 0.0).clamp(-config.ValueClip, config.ValueClip), masks); - var m = mu[0, targetVariate, anchorPatch].item(); - var sd = sigma[0, targetVariate, anchorPatch].item(); - - var head = logits[0, targetVariate, anchorPatch].to(ScalarType.Float32).cpu(); - var acc = head.data(); + var total = ValidateJoint(series.Length, variates, contextBars, horizon, knownFuture, targetVariate, ""); + var known = new bool[variates]; + if (!knownFuture.IsEmpty) knownFuture.CopyTo(known); var result = new double[horizon, config.NumQuantiles]; - for (var h = 0; h < horizon; h++) - for (var q = 0; q < config.NumQuantiles; q++) - result[h, q] = acc[h * config.NumQuantiles + q] * sd + m; + JointCore(1, variates, total, contextBars, horizon, series.ToArray(), known, [targetVariate], [0], + MemoryMarshal.CreateSpan(ref result[0, 0], result.Length)); return result; } @@ -147,49 +85,240 @@ public static TimesFmForecaster Load(string checkpointDir, Device device) /// come from one forward pass, so a longer horizon costs nothing extra. /// [horizon, numQuantiles] of cumulative return relative to the anchor close. public double[,] Forecast(ReadOnlySpan ohlcva, int horizon) + { + ValidateBars(ohlcva.Length, horizon, ""); + var result = new double[horizon, config.NumQuantiles]; + BarsCore(1, ohlcva.Length / Channels, horizon, ohlcva.ToArray(), [0], + MemoryMarshal.CreateSpan(ref result[0, 0], result.Length)); + return result; + } + + /// Entries stacked into one forward pass, at most. Memory grows linearly with + /// it; variate attention never crosses entries. + public int MaxBatchPerPass { get; set; } = 64; + + /// + public int BatchOutputLength(int entries, int horizon) => entries * horizon * config.NumQuantiles; + + /// + /// over many entries. TimesFM has no padding of variates or + /// context, so only entries of one shape — variates, total length and context — share a + /// pass; a batch of mixed shapes runs as one pass per shape. Backtest anchors along one + /// history are a single shape, and batch fully. + /// + /// Entries, each with the constraints of . + /// Steps ahead, shared, 1..OutputPatchLen. + /// Receives [entry][step][quantile]; sized by + /// . + public void ForecastJointBatch(ReadOnlySpan batch, int horizon, Span quantiles) + { + ValidateDestination(quantiles.Length, batch.Length, horizon, nameof(quantiles)); + var shapes = new (int Variates, int Total, int Context)[batch.Length]; + for (var i = 0; i < batch.Length; i++) + { + var e = batch[i]; + var total = ValidateJoint(e.Series.Length, e.Variates, e.ContextBars, horizon, + e.KnownFuture.Span, e.TargetVariate, $"batch[{i}]: "); + shapes[i] = (e.Variates, total, e.ContextBars); + } + + foreach (var pass in Passes(shapes)) + { + var (v, total, context) = shapes[pass[0]]; + var values = new float[pass.Length * v * total]; + var known = new bool[pass.Length * v]; + var targets = new int[pass.Length]; + for (var j = 0; j < pass.Length; j++) + { + var e = batch[pass[j]]; + e.Series.Span.CopyTo(values.AsSpan(j * v * total, v * total)); + if (!e.KnownFuture.IsEmpty) e.KnownFuture.Span.CopyTo(known.AsSpan(j * v, v)); + targets[j] = e.TargetVariate; + } + JointCore(pass.Length, v, total, context, horizon, values, known, targets, pass, quantiles); + } + } + + /// over many OHLCVA windows, one pass per bar count. + /// One row-major [L x 6] window per entry, L a multiple of + /// the patch length. + /// Steps ahead, shared, 1..OutputPatchLen. + /// Receives [entry][step][quantile] of cumulative return; + /// sized by . + public void ForecastBatch(ReadOnlySpan> ohlcva, int horizon, Span returns) + { + ValidateDestination(returns.Length, ohlcva.Length, horizon, nameof(returns)); + var shapes = new (int Bars, int, int)[ohlcva.Length]; + for (var i = 0; i < ohlcva.Length; i++) + { + ValidateBars(ohlcva[i].Length, horizon, $"ohlcva[{i}]: "); + shapes[i] = (ohlcva[i].Length / Channels, 0, 0); + } + + foreach (var pass in Passes(shapes)) + { + var bars = shapes[pass[0]].Bars; + var values = new float[pass.Length * bars * Channels]; + for (var j = 0; j < pass.Length; j++) + ohlcva[pass[j]].Span.CopyTo(values.AsSpan(j * bars * Channels, bars * Channels)); + BarsCore(pass.Length, bars, horizon, values, pass, returns); + } + } + + /// Entry indices grouped by shape, then cut at . + private IEnumerable Passes(TShape[] shapes) where TShape : notnull + => Enumerable.Range(0, shapes.Length) + .GroupBy(i => shapes[i]) + .SelectMany(g => g.Chunk(Math.Max(1, MaxBatchPerPass))); + + /// + /// One pass over entries of one shape, stacked on the batch + /// axis. Writes entry j to at slot slots[j]. + /// + private void JointCore(int count, int variates, int total, int contextBars, int horizon, + float[] series, bool[] known, int[] targets, ReadOnlySpan slots, Span destination) { var p = config.InputPatchLen; - var bars = ohlcva.Length / Channels; - if (bars % p != 0) - throw new ArgumentException($"bar count {bars} must be a multiple of the patch length {p}"); - if (horizon < 1 || horizon > config.OutputPatchLen) - throw new ArgumentOutOfRangeException(nameof(horizon), - $"horizon must lie in 1..{config.OutputPatchLen}"); + var n = total / p; + var anchorPatch = contextBars / p - 1; // last patch made entirely of observed steps + using var scope = NewDisposeScope(); + using var _ = no_grad(); + + var values = tensor(series, [count, variates, n, p]).to(device); + + // Mask the future region of every variate whose future is not supplied, and mark + // as target exactly those variates. Only a target has its future slots withheld, + // so a variate left unmarked would be handed values it should not see. + var masks = zeros([count, variates, n, p], ScalarType.Bool, device); + var isTarget = ones([count, variates, n], ScalarType.Bool, device); + for (var b = 0; b < count; b++) + for (var v = 0; v < variates; v++) + { + if (known[b * variates + v]) { isTarget[b, v] = false; continue; } + if (anchorPatch + 1 < n) masks[b, v, TensorIndex.Slice(anchorPatch + 1)] = true; + } + + var logits = model.forward(values, masks, isTarget) + .view(count, variates, n, config.OutputPatchLen, config.NumQuantiles); + var (_, mu, sigma) = TimesFmPreprocess.RunningStats( + nan_to_num(values, 0.0).clamp(-config.ValueClip, config.ValueClip), masks); + Emit(count, variates, anchorPatch, targets, logits, mu, sigma, horizon, slots, destination, null); + } + + /// One pass over OHLCVA windows of one length. + private void BarsCore(int count, int bars, int horizon, float[] ohlcva, ReadOnlySpan slots, + Span destination) + { + var p = config.InputPatchLen; var n = bars / p; using var scope = NewDisposeScope(); using var _ = no_grad(); - // (bars, 6) -> (1, 6, n, p): variate-major, then patched along time. - var flat = zeros(bars * Channels, ScalarType.Float32); - var dst = flat.data(); - for (var i = 0; i < bars * Channels; i++) dst[i] = ohlcva[i]; - var values = flat.view(bars, Channels).t().reshape(1, Channels, n, p).contiguous().to(device); + // (count, bars, 6) -> (count, 6, n, p): variate-major, then patched along time. + var values = tensor(ohlcva, [count, bars, Channels]) + .permute(0, 2, 1).reshape(count, Channels, n, p).contiguous().to(device); - var masks = zeros([1, Channels, n, p], ScalarType.Bool, device); - var isTarget = ones([1, Channels, n], ScalarType.Bool, device); + var masks = zeros([count, Channels, n, p], ScalarType.Bool, device); + var isTarget = ones([count, Channels, n], ScalarType.Bool, device); var logits = model.forward(values, masks, isTarget) - .view(1, Channels, n, config.OutputPatchLen, config.NumQuantiles); + .view(count, Channels, n, config.OutputPatchLen, config.NumQuantiles); // The final patch carries the forecast; its running statistics undo the // normalisation the model produced its output in. var (_, mu, sigma) = TimesFmPreprocess.RunningStats( nan_to_num(values, 0.0).clamp(-config.ValueClip, config.ValueClip), masks); - var m = mu[0, CloseChannel, n - 1].item(); - var s = sigma[0, CloseChannel, n - 1].item(); - var head = logits[0, CloseChannel, n - 1].to(ScalarType.Float32).cpu(); - var acc = head.data(); + var anchors = new float[count]; + for (var b = 0; b < count; b++) anchors[b] = ohlcva[(b * bars + bars - 1) * Channels + CloseChannel]; + Emit(count, Channels, n - 1, Enumerable.Repeat(CloseChannel, count).ToArray(), logits, mu, sigma, + horizon, slots, destination, anchors); + } - var anchorClose = ohlcva[(bars - 1) * Channels + CloseChannel]; - var result = new double[horizon, config.NumQuantiles]; - for (var h = 0; h < horizon; h++) - for (var q = 0; q < config.NumQuantiles; q++) + /// Reverse RevIN on each entry's target at the anchor patch and write it out, + /// as returns on when given. + private void Emit(int count, int variates, int anchorPatch, int[] targets, Tensor logits, Tensor mu, + Tensor sigma, int horizon, ReadOnlySpan slots, Span destination, float[]? anchorClose) + { + var rows = tensor(targets.Select((t, b) => (long)(b * variates + t)).ToArray()).to(device); + var q = config.NumQuantiles; + var head = logits.select(2, anchorPatch).reshape(count * variates, config.OutputPatchLen, q) + .index_select(0, rows).to(ScalarType.Float32).cpu().contiguous().data(); + var m = mu.select(2, anchorPatch).reshape(count * variates).index_select(0, rows).cpu().data(); + var sd = sigma.select(2, anchorPatch).reshape(count * variates).index_select(0, rows).cpu().data(); + + var block = horizon * q; + var stride = config.OutputPatchLen * q; + for (var b = 0; b < count; b++) { - var price = acc[h * config.NumQuantiles + q] * s + m; // reverse revin - result[h, q] = price / anchorClose - 1.0; + var dst = destination.Slice(slots[b] * block, block); + for (var h = 0; h < horizon; h++) + for (var k = 0; k < q; k++) + { + var value = head[b * stride + h * q + k] * sd[b] + m[b]; // reverse revin + dst[h * q + k] = anchorClose is null ? value : value / anchorClose[b] - 1.0; + } } - return result; + } + + private int ValidateJoint(int length, int variates, int contextBars, int horizon, + ReadOnlySpan knownFuture, int targetVariate, string where) + { + var p = config.InputPatchLen; + if (variates < 1 || variates > MaxVariates) + throw new ArgumentOutOfRangeException(nameof(variates), + $"{where}variates must lie in 1..{MaxVariates}; got {variates}"); + if (length % variates != 0) + throw new ArgumentException( + $"{where}series length {length} is not divisible by {variates} variates", "series"); + var total = length / variates; + + // The known future is gathered by rolling whole patches forward, so the series is + // patched and the context must end on a patch boundary, leaving at least one patch + // after it to roll into. + if (total % p != 0) + throw new ArgumentException( + $"{where}each series must hold a multiple of the patch length {p}; got {total}", "series"); + if (contextBars % p != 0 || contextBars < p || contextBars >= total) + throw new ArgumentException( + $"{where}contextBars ({contextBars}) must be a positive multiple of the patch length {p} " + + $"and leave at least one patch of the {total} supplied", nameof(contextBars)); + if (horizon > total - contextBars) + throw new ArgumentOutOfRangeException(nameof(horizon), + $"{where}horizon {horizon} exceeds the {total - contextBars} steps supplied after the context"); + if (horizon < 1 || horizon > config.OutputPatchLen) + throw new ArgumentOutOfRangeException(nameof(horizon), + $"{where}horizon must lie in 1..{config.OutputPatchLen}"); + if (targetVariate < 0 || targetVariate >= variates) + throw new ArgumentOutOfRangeException(nameof(targetVariate), $"{where}no variate {targetVariate}"); + if (!knownFuture.IsEmpty && knownFuture.Length != variates) + throw new ArgumentException( + $"{where}knownFuture must hold {variates} flags or be empty", nameof(knownFuture)); + if (!knownFuture.IsEmpty && knownFuture[targetVariate]) + throw new ArgumentException( + $"{where}the target variate cannot be known into the future", nameof(knownFuture)); + return total; + } + + private void ValidateBars(int length, int horizon, string where) + { + var p = config.InputPatchLen; + var bars = length / Channels; + if (length == 0 || length % Channels != 0 || bars % p != 0) + throw new ArgumentException( + $"{where}ohlcva must hold whole bars, a multiple of the patch length {p}; got {length} values", "ohlcva"); + if (horizon < 1 || horizon > config.OutputPatchLen) + throw new ArgumentOutOfRangeException(nameof(horizon), + $"{where}horizon must lie in 1..{config.OutputPatchLen}"); + } + + private void ValidateDestination(int length, int entries, int horizon, string name) + { + if (length != BatchOutputLength(entries, horizon)) + throw new ArgumentException( + $"{name} must hold exactly {BatchOutputLength(entries, horizon)} elements " + + $"(entries x horizon x {config.NumQuantiles} quantiles); got {length}. Use {nameof(BatchOutputLength)}.", + name); } } diff --git a/src/Tsfm.Forecasting/IQuantileForecaster.cs b/src/Tsfm.Forecasting/IQuantileForecaster.cs index ac04a3d..2a4dac9 100644 --- a/src/Tsfm.Forecasting/IQuantileForecaster.cs +++ b/src/Tsfm.Forecasting/IQuantileForecaster.cs @@ -1,5 +1,21 @@ namespace Tsfm.Forecasting; +/// +/// One entry of a batched joint forecast: the arguments of +/// other than the horizon, which a batch +/// shares. Memory rather than spans so entries can sit in an array; slicing one large +/// buffer into many entries copies nothing. +/// +/// Variate-major [Variates x total], as for +/// . +/// Series in this entry. +/// Observed steps. +/// Per variate, or empty for none. +/// Which series the entry's quantiles describe. +public readonly record struct JointSeries( + ReadOnlyMemory Series, int Variates, int ContextBars, + ReadOnlyMemory KnownFuture = default, int TargetVariate = 0); + /// /// A model that emits a quantile forecast for one series of a jointly attended set, any of /// which may be known into the future. Implemented by the TimesFM and t0 adapters, so one @@ -8,6 +24,11 @@ namespace Tsfm.Forecasting; /// Shapes are shared; constraints are not. Each implementation documents its own /// limits on context, horizon and variate count — TimesFM, for instance, needs whole /// patches where t0 pads. +/// +/// The batch methods are the single calls over many independent entries, sharing a +/// horizon. Entries never attend to one another, so a batched result equals the single +/// call's up to float rounding; what a batch buys is fewer, larger forward passes. Results +/// land in a caller-supplied span, [entry][step][quantile] row-major. /// /// Implementations are not thread safe. public interface IQuantileForecaster @@ -17,9 +38,9 @@ public interface IQuantileForecaster IReadOnlyList Quantiles { get; } /// Joint forecast over several series. - /// Variate-major [variates x (contextBars + horizon)]: series - /// v at step t is series[v * (contextBars + horizon) + t]. The future - /// region of a variate not marked known is never read. + /// Variate-major [variates x total]: series v at + /// step t is series[v * total + t], where total = series.Length / + /// variates. The future region of a variate not marked known is never read. /// Series supplied. /// Observed steps. /// Steps ahead to return. @@ -37,4 +58,24 @@ public interface IQuantileForecaster /// Steps ahead to return. /// [horizon, Quantiles.Count] of cumulative return relative to the anchor close. double[,] Forecast(ReadOnlySpan ohlcva, int horizon); + + /// over many entries. + /// Independent entries; each may differ in variates, context and + /// known-future pattern, as far as the implementation allows. + /// Steps ahead, shared by every entry. + /// Receives [entry][step][quantile]; sized by + /// . + void ForecastJointBatch(ReadOnlySpan batch, int horizon, Span quantiles); + + /// over many OHLCVA windows — instruments, or anchors + /// along one history. + /// One row-major [L x 6] window per entry. + /// Steps ahead, shared by every entry. + /// Receives [entry][step][quantile] of cumulative return; + /// sized by . + void ForecastBatch(ReadOnlySpan> ohlcva, int horizon, Span returns); + + /// Elements a batch of writes. A wrong size is + /// refused rather than silently misaligning every entry. + int BatchOutputLength(int entries, int horizon) => entries * horizon * Quantiles.Count; } diff --git a/tests/Tsfm.Forecasting.Kronos.IntegrationTests/BatchTests.cs b/tests/Tsfm.Forecasting.Kronos.IntegrationTests/BatchTests.cs new file mode 100644 index 0000000..e2a8495 --- /dev/null +++ b/tests/Tsfm.Forecasting.Kronos.IntegrationTests/BatchTests.cs @@ -0,0 +1,78 @@ +using Tsfm.Forecasting.Kronos.Weights; + +using Xunit; + +using static TorchSharp.torch; + +namespace Tsfm.Forecasting.Kronos.IntegrationTests; + +/// +/// against the single-series path it batches. +/// Sampling is seeded by each window's anchor bar, so a sampled batch must reproduce the +/// single-series rows exactly, not merely in distribution. +/// +public class BatchTests : IDisposable +{ + private const int Context = 64, Horizon = 3, Rollouts = 5; + + private readonly KronosForecaster _forecaster = + KronosForecaster.Load(KronosMini.Instance, new Device("cpu")); + + public void Dispose() => _forecaster.Dispose(); + + [Fact] + public void EachRowEqualsTheLastRowOfItsSeriesAlone() + { + // Three instruments of different lengths and levels, on different bar grids. + KronosSeries[] series = [Synthetic(70, 100f, 0), Synthetic(64, 2500f, 7), Synthetic(90, 0.8f, 3)]; + var n = series.Length; + var (lean, up, disp) = (new float[n], new int[n], new float[n]); + var path = new float[n * Horizon * 6]; + + _forecaster.InferBatch(series, lean, up, disp, path, Context, Horizon, Rollouts, + greedy: false, 1f, 0.9f, batch: 2); + + for (var i = 0; i < n; i++) + { + var bars = series[i].BarTimeMs.Length; + var rows = KronosForecaster.OutputCount(bars, Context); + var (l, u, d) = (new float[rows], new int[rows], new float[rows]); + var p = new float[KronosForecaster.PathLength(bars, Context, Horizon)]; + _forecaster.InferPath(series[i].Ohlcva.Span, series[i].BarTimeMs.Span, l, u, d, p, + Context, Horizon, Rollouts, greedy: false, 1f, 0.9f); + + Assert.Equal(l[^1], lean[i]); + Assert.Equal(u[^1], up[i]); + Assert.Equal(d[^1], disp[i]); + Assert.Equal(p.AsSpan(p.Length - Horizon * 6).ToArray(), path.AsSpan(i * Horizon * 6, Horizon * 6).ToArray()); + } + } + + [Fact] + public void TooShortASeriesIsRefused() + { + KronosSeries[] series = [Synthetic(Context - 1, 100f, 0)]; + var ex = Assert.Throws(() => _forecaster.InferBatch(series, + new float[1], new int[1], default, default, Context, Horizon, Rollouts, greedy: true, 1f, 1f)); + Assert.Contains("series[0]", ex.Message); + } + + private static KronosSeries Synthetic(int bars, float level, int offsetBars) + { + var ohlcva = new float[bars * 6]; + var times = new long[bars]; + var t0 = new DateTimeOffset(2026, 3, 2, 0, 0, 0, TimeSpan.Zero).ToUnixTimeMilliseconds(); + for (var i = 0; i < bars; i++) + { + var mid = level * (1f + 0.01f * MathF.Sin((i + offsetBars) / 9f)); + ohlcva[i * 6 + 0] = mid * 0.999f; + ohlcva[i * 6 + 1] = mid * 1.002f; + ohlcva[i * 6 + 2] = mid * 0.997f; + ohlcva[i * 6 + 3] = mid * 1.001f; + ohlcva[i * 6 + 4] = 50f + i % 5; + ohlcva[i * 6 + 5] = ohlcva[i * 6 + 4] * mid; + times[i] = t0 + (i + offsetBars) * 300_000L; + } + return new KronosSeries(ohlcva, times); + } +} diff --git a/tests/Tsfm.Forecasting.T0.Parity/Program.cs b/tests/Tsfm.Forecasting.T0.Parity/Program.cs index 1e481dc..8af08ac 100644 --- a/tests/Tsfm.Forecasting.T0.Parity/Program.cs +++ b/tests/Tsfm.Forecasting.T0.Parity/Program.cs @@ -139,6 +139,29 @@ void Check(string stage, Tensor got, Tensor want, double tol, string note = "") var epCov = forecaster.ForecastJoint(epJoint, 3, ep.Length, epTruth.Length, [false, true, true]); Check("notebook ep + cov", Grid(epCov), E("ep_predict_cov"), 1e-4); + // 7. One batch of mixed lengths — 150, 64 and 512 steps, one to two targets — against + // the reference's single forecasts. Truncating each to a shared 24 steps is exact for + // target-only entries: targets are causal, so later withheld patches cannot reach them. + const int shared = 24; + JointSeries[] batch = + [ + new(joint.AsMemory(0, 2 * total), 2, t, TargetVariate: 1), + new(aligned, 1, 64), + new(m4, 1, m4.Length), + new(ep, 1, ep.Length), + ]; + string[] refs = ["predict_targets_only", "predict_aligned", "m4_predict", "ep_predict"]; + var outBatch = new double[forecaster.BatchOutputLength(batch.Length, shared)]; + forecaster.ForecastJointBatch(batch, shared, outBatch); + for (var i = 0; i < batch.Length; i++) + { + var want = E(refs[i]); + if (want.dim() == 3) want = want[1]; + var got = tensor(outBatch.AsSpan(i * shared * cfg.NumQuantiles, shared * cfg.NumQuantiles) + .ToArray().Select(d => (float)d).ToArray(), [shared, cfg.NumQuantiles]); + Check($"batch[{i}]", got, want.narrow(0, 0, shared), 1e-4, $"({refs[i]}, one pass)"); + } + var mid = Array.IndexOf(cfg.Quantiles, 0.5); Console.WriteLine($" electricity price |median-truth|: {MeanAbs(epPlain, epTruth, mid):F4} without covariates, " + $"{MeanAbs(epCov, epTruth, mid):F4} with"); diff --git a/tests/Tsfm.Forecasting.T0.Tests/BatchTests.cs b/tests/Tsfm.Forecasting.T0.Tests/BatchTests.cs new file mode 100644 index 0000000..d2caccd --- /dev/null +++ b/tests/Tsfm.Forecasting.T0.Tests/BatchTests.cs @@ -0,0 +1,116 @@ +using TorchSharp; +using Xunit; + +using static TorchSharp.torch; + +namespace Tsfm.Forecasting.T0.Tests; + +/// +/// Batched calls against the single calls they batch, on a tiny random-init model. Entries +/// are separate groups, so nothing may leak between them; differing contexts are padded +/// to a shared width, which moves rotary positions but not relative ones, so agreement is +/// to float rounding rather than exact. +/// +public class BatchTests +{ + private static readonly T0Config Tiny = new( + EmbedDim: 16, NumLayers: 3, NumHeads: 2, MlpHiddenDim: 32, PatchSize: 4, GroupEveryN: 3, + Quantiles: [0.1, 0.5, 0.9], ScalerUseArcsinh: true, ScalerEps: 0.1, ScalerStdClamp: false, + XposHalvesLayout: false, ScalerZeroNoValue: true); + + private static T0Forecaster Forecaster(int maxRowsPerPass) + { + manual_seed(7); + var model = new T0Model(Tiny); + model.eval(); + return new T0Forecaster(model, Tiny, new Device(DeviceType.CPU)) { MaxRowsPerPass = maxRowsPerPass }; + } + + private const int Horizon = 5; + + private static JointSeries Entry(int variates, int context, bool[] known, int target, int seed) + { + var r = new Random(seed); + var total = context + Horizon + 2; // spare steps past the horizon are ignored + var s = new float[variates * total]; + for (var v = 0; v < variates; v++) + for (var t = 0; t < total; t++) s[v * total + t] = (v + 1) * 10f + MathF.Sin(t / 3f + v) + (float)r.NextDouble(); + return new JointSeries(s, variates, context, known, target); + } + + private static readonly JointSeries[] Mixed = + [ + Entry(3, 13, [false, false, true], 1, 1), + Entry(1, 6, [], 0, 2), + Entry(4, 21, [false, true, true, false], 3, 3), + Entry(2, 8, [false, true], 0, 4), + Entry(1, 1, [], 0, 5), + ]; + + [Theory] + [InlineData(256)] // one pass + [InlineData(4)] // several passes, one entry wider than the budget + public void JointBatchMatchesSingles(int maxRowsPerPass) + { + var f = Forecaster(maxRowsPerPass); + var q = Tiny.NumQuantiles; + var batched = new double[f.BatchOutputLength(Mixed.Length, Horizon)]; + f.ForecastJointBatch(Mixed, Horizon, batched); + + for (var i = 0; i < Mixed.Length; i++) + { + var e = Mixed[i]; + var single = f.ForecastJoint(e.Series.Span, e.Variates, e.ContextBars, Horizon, e.KnownFuture.Span, e.TargetVariate); + for (var h = 0; h < Horizon; h++) + for (var k = 0; k < q; k++) + Assert.Equal(single[h, k], batched[(i * Horizon + h) * q + k], 4); + } + } + + [Fact] + public void AnEntryIsUnaffectedByItsCompanions() + { + var f = Forecaster(256); + var size = f.BatchOutputLength(2, Horizon); + var (a, b) = (new double[size], new double[size]); + f.ForecastJointBatch([Mixed[0], Mixed[2]], Horizon, a); + f.ForecastJointBatch([Mixed[0], Entry(4, 21, [false, true, true, false], 3, 99)], Horizon, b); + var block = Horizon * Tiny.NumQuantiles; + Assert.Equal(a.AsSpan(0, block).ToArray(), b.AsSpan(0, block).ToArray()); + } + + [Fact] + public void OhlcvaBatchMatchesSingles() + { + var f = Forecaster(256); + var bars = new float[40 * 6]; + for (var i = 0; i < 40; i++) + for (var c = 0; c < 6; c++) bars[i * 6 + c] = c < 4 ? 100f + MathF.Sin(i / 4f) + c * 0.1f : 500f + i; + + // Anchors along one history, as a backtest slices them: no copies. + ReadOnlyMemory[] windows = [bars.AsMemory(0, 20 * 6), bars.AsMemory(6, 25 * 6), bars]; + var batched = new double[f.BatchOutputLength(windows.Length, Horizon)]; + f.ForecastBatch(windows, Horizon, batched); + + var q = Tiny.NumQuantiles; + for (var i = 0; i < windows.Length; i++) + { + var single = f.Forecast(windows[i].Span, Horizon); + for (var h = 0; h < Horizon; h++) + for (var k = 0; k < q; k++) + Assert.Equal(single[h, k], batched[(i * Horizon + h) * q + k], 6); + } + } + + [Fact] + public void BadEntriesAndDestinationsAreRefused() + { + var f = Forecaster(256); + var size = f.BatchOutputLength(2, Horizon); + var ex = Assert.Throws(() => + f.ForecastJointBatch([Mixed[0], Mixed[1] with { KnownFuture = new[] { true } }], Horizon, new double[size])); + Assert.Contains("batch[1]", ex.Message); + Assert.Throws(() => f.ForecastJointBatch([Mixed[0], Mixed[1]], Horizon, new double[size - 1])); + Assert.Throws(() => f.ForecastBatch([new float[6], new float[7]], Horizon, new double[size])); + } +} diff --git a/tests/Tsfm.Forecasting.TimesFm.IntegrationTests/BatchTests.cs b/tests/Tsfm.Forecasting.TimesFm.IntegrationTests/BatchTests.cs new file mode 100644 index 0000000..6a3811c --- /dev/null +++ b/tests/Tsfm.Forecasting.TimesFm.IntegrationTests/BatchTests.cs @@ -0,0 +1,84 @@ +using Xunit; + +using static TorchSharp.torch; + +namespace Tsfm.Forecasting.TimesFm.IntegrationTests; + +/// +/// Batched calls against the single calls they batch. Run on a small random-init model, +/// so they need no checkpoint and bite in CI: what is under test is the stacking of +/// entries on the batch axis, not the weights. +/// +public class BatchTests +{ + private static readonly TimesFmConfig Tiny = new( + InputPatchLen: 8, OutputPatchLen: 16, Quantiles: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], + NumLayers: 2, ModelDims: 32, NumHeads: 4, UseVariateAttention: true, MaxVariates: 8, ValueClip: 1000); + + private static TimesFmForecaster Forecaster(int maxBatchPerPass) + { + manual_seed(11); + var model = new TimesFmModel(Tiny); + model.eval(); + return new TimesFmForecaster(model, Tiny, new Device("cpu")) { MaxBatchPerPass = maxBatchPerPass }; + } + + private static JointSeries Entry(int variates, int total, int context, bool[] known, int target, int seed) + { + var r = new Random(seed); + var s = new float[variates * total]; + for (var i = 0; i < s.Length; i++) s[i] = 20f * (i / total + 1) + (float)r.NextDouble(); + return new JointSeries(s, variates, context, known, target); + } + + [Theory] + [InlineData(64)] + [InlineData(2)] + public void JointBatchMatchesSinglesExactly(int maxBatchPerPass) + { + // Two shapes, interleaved; one pass per shape, or several when capped. + JointSeries[] batch = + [ + Entry(3, 32, 16, [false, false, true], 0, 1), + Entry(2, 24, 8, [], 1, 2), + Entry(3, 32, 16, [false, true, true], 0, 3), + Entry(3, 32, 16, [], 2, 4), + Entry(2, 24, 8, [false, true], 0, 5), + ]; + const int horizon = 6; + var f = Forecaster(maxBatchPerPass); + var q = Tiny.Quantiles.Length; + var batched = new double[f.BatchOutputLength(batch.Length, horizon)]; + f.ForecastJointBatch(batch, horizon, batched); + + for (var i = 0; i < batch.Length; i++) + { + var e = batch[i]; + var single = f.ForecastJoint(e.Series.Span, e.Variates, e.ContextBars, horizon, e.KnownFuture.Span, e.TargetVariate); + for (var h = 0; h < horizon; h++) + for (var k = 0; k < q; k++) + Assert.Equal(single[h, k], batched[(i * horizon + h) * q + k]); + } + } + + [Fact] + public void OhlcvaBatchMatchesSinglesExactly() + { + var bars = new float[32 * 6]; + for (var i = 0; i < bars.Length; i++) bars[i] = 100f + MathF.Sin(i / 7f); + ReadOnlyMemory[] windows = [bars.AsMemory(0, 16 * 6), bars.AsMemory(16 * 6), bars, bars.AsMemory(8 * 6, 16 * 6)]; + const int horizon = 5; + + var f = Forecaster(64); + var q = Tiny.Quantiles.Length; + var batched = new double[f.BatchOutputLength(windows.Length, horizon)]; + f.ForecastBatch(windows, horizon, batched); + for (var i = 0; i < windows.Length; i++) + { + var single = f.Forecast(windows[i].Span, horizon); + for (var h = 0; h < horizon; h++) + for (var k = 0; k < q; k++) + Assert.Equal(single[h, k], batched[(i * horizon + h) * q + k]); + } + } +} From 64587d8638c7c22bb82454cd6836be6477d72773 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 14:37:56 +0000 Subject: [PATCH 3/9] feat(kronos): make the sampling seed a host policy KronosForecaster.Seed decides each window's sampling stream from its anchor bar and a series key (KronosSeries.SeriesKey, or seriesKey: on Infer/InferPath). SamplingSeeds.AnchorTime, the default, keeps the current behaviour: series anchored on one bar share draws. PerSeries gives each key an independent stream; Salted derives a fresh, reproducible run from any policy. Batch invariance holds under all of them, and the default is bit-identical to before. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01F1dTUhvq2UnFWdBmH7AA8G --- CHANGELOG.md | 6 ++ README.md | 8 +- .../KronosForecaster.cs | 89 +++++++++++++++---- .../BatchTests.cs | 44 +++++++++ 4 files changed, 130 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6c2783c..14aca43 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -42,6 +42,12 @@ `batch × rollouts` rows; row `i` equals the last row `InferPath` gives for series `i` alone, sampling included. `Infer` and `InferPath` share the refactored loop and are bit-identical to 0.1.3. +- `KronosForecaster.Seed` makes the sampling seed a host policy over the anchor bar and a + series key (`KronosSeries.SeriesKey`, or `seriesKey:` on `Infer`/`InferPath`). + `SamplingSeeds.AnchorTime`, the default, is the 0.1.x behaviour: series anchored on one bar + share draws (common random numbers). `PerSeries` gives each key its own stream, and + `Salted` derives a fresh reproducible run from any policy. Batch invariance holds under + all of them; the default is bit-identical to 0.1.3. ### Packaging diff --git a/README.md b/README.md index dd2f1d9..02b2c8f 100644 --- a/README.md +++ b/README.md @@ -224,8 +224,12 @@ rows, and a `RolloutAggregator` reduces only across the rollouts of one series, and channel — never across series. So the two compose without interaction, except in cost: memory scales with the product, which is why Kronos's `batch` counts windows and leaves rows to be `batch * rollouts`. Sampling is seeded per window from its anchor bar, so -results do not depend on how windows were grouped into passes; the corollary is that two -instruments anchored on the same bar draw the same uniforms (common random numbers). +results do not depend on how windows were grouped into passes. By default two instruments +anchored on the same bar therefore draw the same uniforms — common random numbers, which +make cross-sectional comparisons less noisy but leave their errors correlated. That is a +host policy, `KronosForecaster.Seed`: `SamplingSeeds.PerSeries` gives each +`KronosSeries.SeriesKey` its own stream, and `SamplingSeeds.Salted` derives a fresh, +reproducible run from either. TimesFM and t0 emit quantiles directly, with no rollouts and nothing to aggregate. t0's upstream continuation past 1024 steps (not ported) would be the same pattern again: each requested quantile becomes a path, paths are folded into the batch as extra groups, and diff --git a/src/Tsfm.Forecasting.Kronos/KronosForecaster.cs b/src/Tsfm.Forecasting.Kronos/KronosForecaster.cs index 15f88d2..36b1519 100644 --- a/src/Tsfm.Forecasting.Kronos/KronosForecaster.cs +++ b/src/Tsfm.Forecasting.Kronos/KronosForecaster.cs @@ -22,6 +22,14 @@ public sealed class KronosForecaster : IDisposable private readonly KronosModel _model; private readonly Device _device; + /// + /// Seeds each window's sampling stream from its anchor bar and series key. The default, + /// , ignores the key: every series anchored on a + /// bar draws the same uniforms. See for the alternatives. + /// Whatever the policy, a window's draws never depend on how windows were batched. + /// + public SamplingSeed Seed { get; set; } = SamplingSeeds.AnchorTime; + /// Rows this slice produces, L − K + 1. Size buffers from this; a /// wrong count misaligns every row. /// @@ -98,15 +106,18 @@ public static KronosForecaster Load(ICheckpoint checkpoint, Device device) /// Logit temperature before sampling. /// Nucleus cutoff. /// Windows per forward pass. + /// Identifies the series to . The default policy + /// ignores it; needs a distinct, stable key per + /// instrument — a symbol hash, say. /// L − K + 1 rows, the last anchored on the final bar. Forward /// timestamps are extrapolated from the bar interval, so no trailing bars are consumed. /// Not thread safe. public bool Infer(ReadOnlySpan ohlcva, ReadOnlySpan barTimeMs, Span lean, Span upCount, Span dispersion, int contextBars, int horizon, int rollouts, bool greedy, - float temperature, float topP, int batch = 8) + float temperature, float topP, int batch = 8, long seriesKey = 0) => InferCore(ohlcva, barTimeMs, lean, upCount, dispersion, default, null, - contextBars, horizon, rollouts, greedy, temperature, topP, batch); + contextBars, horizon, rollouts, greedy, temperature, topP, batch, seriesKey); // CS1573 does not expand , though the IDE does, so the shared parameters // would have to be duplicated from Infer to satisfy a check the tooltips already pass. @@ -140,10 +151,10 @@ public bool Infer(ReadOnlySpan ohlcva, ReadOnlySpan barTimeMs, public bool InferPath(ReadOnlySpan ohlcva, ReadOnlySpan barTimeMs, Span lean, Span upCount, Span dispersion, Span path, int contextBars, int horizon, int rollouts, bool greedy, - float temperature, float topP, RolloutAggregator? aggregate = null, int batch = 8) + float temperature, float topP, RolloutAggregator? aggregate = null, int batch = 8, long seriesKey = 0) => InferCore(ohlcva, barTimeMs, lean, upCount, dispersion, path, aggregate ?? RolloutAggregators.Mean, - contextBars, horizon, rollouts, greedy, temperature, topP, batch); + contextBars, horizon, rollouts, greedy, temperature, topP, batch, seriesKey); #pragma warning restore CS1573 /// Elements writes: (L - K + 1) * horizon * Channels. @@ -161,10 +172,11 @@ public static int PathLength(int barCount, int contextBars, int horizon) /// each make batch * rollouts rows. Each output reduces /// only its own series' rollouts. /// - /// Sampling is seeded by the anchor bar's timestamp alone, so series anchored on - /// the same bar draw the same uniforms — common random numbers. Comparisons across - /// instruments are then less noisy than independent draws would make them, but their - /// sampling errors are correlated, not independent. + /// Sampling is seeded by . Under the default, the anchor bar's + /// timestamp alone, series anchored on the same bar draw the same uniforms — common + /// random numbers: comparisons across instruments are less noisy, but their sampling + /// errors are correlated. with distinct + /// s makes them independent. /// /// Each with at least bars. /// One row per series. @@ -198,7 +210,7 @@ private bool InferCore(ReadOnlySpan ohlcva, ReadOnlySpan barTimeMs, Span lean, Span upCount, Span dispersion, Span path, RolloutAggregator? aggregate, int contextBars, int horizon, int rollouts, bool greedy, - float temperature, float topP, int batch = 8) + float temperature, float topP, int batch, long seriesKey) { var length = barTimeMs.Length; if (ohlcva.Length != length * Channels) @@ -206,7 +218,7 @@ private bool InferCore(ReadOnlySpan ohlcva, ReadOnlySpan barTimeMs, var count = OutputCount(length, contextBars); if (contextBars <= MaxContext && count <= 0) return false; - Drive(new Windows(ohlcva, barTimeMs, contextBars), count, lean, upCount, dispersion, path, aggregate, + Drive(new Windows(ohlcva, barTimeMs, contextBars, seriesKey), count, lean, upCount, dispersion, path, aggregate, contextBars, horizon, rollouts, greedy, temperature, topP, batch); return true; } @@ -301,12 +313,14 @@ private readonly ref struct Windows private readonly ReadOnlySpan _series; private readonly int _k; private readonly bool _many; + private readonly long _key; - public Windows(ReadOnlySpan ohlcva, ReadOnlySpan times, int k) + public Windows(ReadOnlySpan ohlcva, ReadOnlySpan times, int k, long key) { _ohlcva = ohlcva; _times = times; _k = k; + _key = key; } public Windows(ReadOnlySpan series, int k) @@ -323,6 +337,8 @@ public Windows(ReadOnlySpan series, int k) public int Origin(int i) => _many ? _series[i].BarTimeMs.Length - _k : i; /// Bar interval, from the series' first spacing. + public long SeriesKey(int i) => _many ? _series[i].SeriesKey : _key; + public long StepMs(int i) { var t = Times(i); @@ -382,7 +398,7 @@ private void EvaluateBatch( // Draw before the loop: sampling must be a pure function of the bar, not of how // windows were batched. - if (!greedy) DrawUniforms(source, start, take, k, rollouts, h, uniforms); + if (!greedy) DrawUniforms(source, Seed, start, take, k, rollouts, h, uniforms); for (var step = 0; step < h; step++) { var len = k + step; @@ -477,12 +493,13 @@ private static void NormaliseWindow( /// window's anchor bar. SplitMix64, so the stream cannot shift with a framework /// version. private static void DrawUniforms( - in Windows source, int start, int take, int k, int rollouts, int h, float[] u) + in Windows source, SamplingSeed seed, int start, int take, int k, int rollouts, int h, float[] u) { var perWindow = rollouts * h * 2; for (var w = 0; w < take; w++) { - var state = (ulong)SeedFor(source.Times(start + w)[source.Origin(start + w) + k - 1]); + var anchorMs = source.Times(start + w)[source.Origin(start + w) + k - 1]; + var state = (ulong)seed(anchorMs, source.SeriesKey(start + w)); for (var i = 0; i < perWindow; i++) { state += 0x9E3779B97F4A7C15UL; @@ -572,4 +589,46 @@ public void Dispose() /// an array; slicing a larger buffer copies nothing. /// Row-major [L x 6]: open, high, low, close, volume, amount. /// Close time per bar, Unix milliseconds, [L]. -public readonly record struct KronosSeries(ReadOnlyMemory Ohlcva, ReadOnlyMemory BarTimeMs); +/// Identifies the series to ; +/// ignored by the default policy. Keep it stable across calls — a symbol hash, not a batch +/// index — or a series' draws change with the company it is batched in. +public readonly record struct KronosSeries(ReadOnlyMemory Ohlcva, ReadOnlyMemory BarTimeMs, long SeriesKey = 0); + +/// A sampling seed for one window, from its anchor bar's close time and its +/// series key. Must be a pure function: draws are reproduced from it. +public delegate long SamplingSeed(long anchorCloseTimeMs, long seriesKey); + +/// +/// Seeding policies for . +/// +/// The choice is between common and independent random numbers across series. Common +/// ones (the default) make a cross-sectional comparison — ranking instruments by P(up), say +/// — less noisy, because every instrument's sampling error moves together; but for the same +/// reason an average over instruments does not shrink that error. Independent ones behave as +/// separate experiments. +/// +public static class SamplingSeeds +{ + /// The anchor bar alone, ignoring the series: series anchored on one bar share + /// draws. The behaviour of 0.1.x and the default. + public static readonly SamplingSeed AnchorTime = (anchorMs, _) => KronosForecaster.SeedFor(anchorMs); + + /// The anchor bar mixed with the series key: independent streams for distinct + /// keys, and the same stream for the same series and bar on every call. + public static readonly SamplingSeed PerSeries = (anchorMs, key) => Mix(KronosForecaster.SeedFor(anchorMs), key); + + /// mixed with a run-level salt: a fresh, still + /// reproducible, set of draws per salt — for independent replications of a whole run. + public static SamplingSeed Salted(SamplingSeed inner, long salt) => (anchorMs, key) => Mix(inner(anchorMs, key), salt); + + /// SplitMix64's finaliser over the pair, kept non-negative like + /// . + private static long Mix(long seed, long key) + { + var z = (ulong)seed ^ ((ulong)key * 0x9E3779B97F4A7C15UL); + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9UL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBUL; + z ^= z >> 31; + return (long)(z & 0x7FFFFFFFFFFFFFFFUL); + } +} diff --git a/tests/Tsfm.Forecasting.Kronos.IntegrationTests/BatchTests.cs b/tests/Tsfm.Forecasting.Kronos.IntegrationTests/BatchTests.cs index e2a8495..08bb5a2 100644 --- a/tests/Tsfm.Forecasting.Kronos.IntegrationTests/BatchTests.cs +++ b/tests/Tsfm.Forecasting.Kronos.IntegrationTests/BatchTests.cs @@ -48,6 +48,50 @@ public void EachRowEqualsTheLastRowOfItsSeriesAlone() } } + /// Two instruments with identical bars on the same grid: the default policy + /// draws the same uniforms for both, PerSeries draws independently, and either way a + /// batched row still equals its single-series call with the same key. + [Fact] + public void SeedPolicyDecidesWhetherSeriesShareDraws() + { + var a = Synthetic(70, 100f, 0) with { SeriesKey = 11 }; + var b = a with { SeriesKey = 22 }; + float[] Leans(SamplingSeed seed) + { + _forecaster.Seed = seed; + var lean = new float[2]; + var path = new float[2 * Horizon * 6]; + _forecaster.InferBatch([a, b], lean, new int[2], default, path, Context, Horizon, Rollouts, + greedy: false, 1f, 1f); + return [path[3], path[Horizon * 6 + 3]]; // first step's close, each series + } + + try + { + var common = Leans(SamplingSeeds.AnchorTime); + Assert.Equal(common[0], common[1]); + + var independent = Leans(SamplingSeeds.PerSeries); + Assert.NotEqual(independent[0], independent[1]); + + // Batch invariance holds under the policy: B alone, keyed 22, reproduces its row. + var bars = b.BarTimeMs.Length; + var rows = KronosForecaster.OutputCount(bars, Context); + var p = new float[KronosForecaster.PathLength(bars, Context, Horizon)]; + _forecaster.InferPath(b.Ohlcva.Span, b.BarTimeMs.Span, new float[rows], new int[rows], default, p, + Context, Horizon, Rollouts, greedy: false, 1f, 1f, seriesKey: 22); + Assert.Equal(independent[1], p[p.Length - Horizon * 6 + 3]); + + var salted = Leans(SamplingSeeds.Salted(SamplingSeeds.AnchorTime, 1)); + Assert.Equal(salted[0], salted[1]); + Assert.NotEqual(common[0], salted[0]); + } + finally + { + _forecaster.Seed = SamplingSeeds.AnchorTime; + } + } + [Fact] public void TooShortASeriesIsRefused() { From 60759a579a10ebbc341ddca31ce6b45653a1b968 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 14:45:24 +0000 Subject: [PATCH 4/9] feat(timestral): add Mistral AI's Timestral family, with 5M and 28M embedded Port mistralai/timestral 1.0.0 as Tsfm.Forecasting.Timestral (Apache-2.0): causal per-patch scaler, gated time/variate transformer with partial RoPE, Decoupled Horizon Forecasting decoder tokens and the monotone 21-quantile head. TimestralForecaster implements IQuantileForecaster, batches included, and applies the reference forecaster's defaults (symmetric averaging, non-negative clipping, advancing refine frame, quantile cap); the mirrored input rides in the same forward pass. Verified stage by stage against the reference on all four sizes, with and without the defaults, including the covariate path and the shared notebook examples. Timestral-5M and -28M ship in Tsfm.Forecasting.Timestral.Weights.*, embedded at pinned revisions whose SHA-256 the build checks. The weights are under Mistral's revenue-restricted Modified MIT licence, carried verbatim as the package licence file with acceptance required; 123M and 209M are fetched by script. Core gains IModelWeights for tokenizer-free checkpoints. Also updates the README byline as requested. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01F1dTUhvq2UnFWdBmH7AA8G --- .github/workflows/ci.yml | 16 +- .github/workflows/release.yml | 4 +- CHANGELOG.md | 35 +- LICENSE-Timestral-Weights | 30 + NOTICE | 37 ++ README.md | 94 +++- Tsfm.Forecasting.slnx | 6 + reference/dump_timestral_parity.py | 172 ++++++ .../GroceryCovariates.csproj | 2 + samples/GroceryCovariates/Program.cs | 20 +- scripts/fetch-timestral-checkpoint.sh | 50 ++ scripts/verify-packages.sh | 15 +- .../Timestral28M.cs | 37 ++ ...m.Forecasting.Timestral.Weights.28M.csproj | 43 ++ .../Timestral5M.cs | 37 ++ ...fm.Forecasting.Timestral.Weights.5M.csproj | 43 ++ .../TimestralForecaster.cs | 532 ++++++++++++++++++ .../TimestralModel.cs | 184 ++++++ .../TimestralModules.cs | 305 ++++++++++ .../TimestralPreprocess.cs | 195 +++++++ .../Tsfm.Forecasting.Timestral.csproj | 24 + src/Tsfm.Forecasting/IModelWeights.cs | 78 +++ .../EmbeddedWeightsTests.cs | 77 +++ ...ecasting.Timestral.IntegrationTests.csproj | 20 + .../Program.cs | 183 ++++++ .../Tsfm.Forecasting.Timestral.Parity.csproj | 16 + .../BatchTests.cs | 84 +++ .../InvariantTests.cs | 154 +++++ .../Tsfm.Forecasting.Timestral.Tests.csproj | 19 + 29 files changed, 2497 insertions(+), 15 deletions(-) create mode 100644 LICENSE-Timestral-Weights create mode 100644 reference/dump_timestral_parity.py create mode 100755 scripts/fetch-timestral-checkpoint.sh create mode 100644 src/Tsfm.Forecasting.Timestral.Weights.28M/Timestral28M.cs create mode 100644 src/Tsfm.Forecasting.Timestral.Weights.28M/Tsfm.Forecasting.Timestral.Weights.28M.csproj create mode 100644 src/Tsfm.Forecasting.Timestral.Weights.5M/Timestral5M.cs create mode 100644 src/Tsfm.Forecasting.Timestral.Weights.5M/Tsfm.Forecasting.Timestral.Weights.5M.csproj create mode 100644 src/Tsfm.Forecasting.Timestral/TimestralForecaster.cs create mode 100644 src/Tsfm.Forecasting.Timestral/TimestralModel.cs create mode 100644 src/Tsfm.Forecasting.Timestral/TimestralModules.cs create mode 100644 src/Tsfm.Forecasting.Timestral/TimestralPreprocess.cs create mode 100644 src/Tsfm.Forecasting.Timestral/Tsfm.Forecasting.Timestral.csproj create mode 100644 src/Tsfm.Forecasting/IModelWeights.cs create mode 100644 tests/Tsfm.Forecasting.Timestral.IntegrationTests/EmbeddedWeightsTests.cs create mode 100644 tests/Tsfm.Forecasting.Timestral.IntegrationTests/Tsfm.Forecasting.Timestral.IntegrationTests.csproj create mode 100644 tests/Tsfm.Forecasting.Timestral.Parity/Program.cs create mode 100644 tests/Tsfm.Forecasting.Timestral.Parity/Tsfm.Forecasting.Timestral.Parity.csproj create mode 100644 tests/Tsfm.Forecasting.Timestral.Tests/BatchTests.cs create mode 100644 tests/Tsfm.Forecasting.Timestral.Tests/InvariantTests.cs create mode 100644 tests/Tsfm.Forecasting.Timestral.Tests/Tsfm.Forecasting.Timestral.Tests.csproj diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index aec810a..5d6118b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -54,6 +54,10 @@ jobs: - name: t0 tests run: dotnet test tests/Tsfm.Forecasting.T0.Tests -c Release --nologo + # Random-init model, no weights needed. + - name: Timestral tests + run: dotnet test tests/Tsfm.Forecasting.Timestral.Tests -c Release --nologo + # Packing embeds the checkpoints, so this job needs them. Platform-neutral. pack: runs-on: ubuntu-latest @@ -71,15 +75,17 @@ jobs: key: nuget-pack-${{ hashFiles('**/*.csproj', 'Directory.Build.props') }} restore-keys: nuget-pack- - # Key on the fetch script: revisions are pinned there and change deliberately. + # Key on the fetch scripts: revisions are pinned there and change deliberately. - name: Cache checkpoints uses: actions/cache@v4 with: path: checkpoints - key: checkpoints-${{ hashFiles('scripts/fetch-checkpoints.sh') }} + key: checkpoints-${{ hashFiles('scripts/fetch-checkpoints.sh', 'scripts/fetch-timestral-checkpoint.sh') }} - name: Fetch checkpoints - run: ./scripts/fetch-checkpoints.sh + run: | + ./scripts/fetch-checkpoints.sh + ./scripts/fetch-timestral-checkpoint.sh # 5M and 28M, the embedded pair (~132 MB) - name: Pack run: dotnet pack -c Release --nologo -o artifacts @@ -87,7 +93,9 @@ jobs: # Weights-backed, so they belong here rather than in the fast matrix, which fetches # no checkpoints. The example is built to prove it still compiles against the API. - name: Integration tests - run: dotnet test tests/Tsfm.Forecasting.Kronos.IntegrationTests -c Release --nologo + run: | + dotnet test tests/Tsfm.Forecasting.Kronos.IntegrationTests -c Release --nologo + dotnet test tests/Tsfm.Forecasting.Timestral.IntegrationTests -c Release --nologo - name: Build samples run: | diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 25f1d49..1285132 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -27,7 +27,9 @@ jobs: global-json-file: global.json - name: Fetch checkpoints - run: ./scripts/fetch-checkpoints.sh + run: | + ./scripts/fetch-checkpoints.sh + ./scripts/fetch-timestral-checkpoint.sh - name: Verify the tag matches the declared version if: startsWith(github.ref, 'refs/tags/v') diff --git a/CHANGELOG.md b/CHANGELOG.md index 14aca43..c0c5c04 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,7 +25,35 @@ - Weights are not bundled. The README gives pinned download commands with checksums for projects that only reference the package, and `Load` names them when files are missing. -### Batching, across all three models +### Tsfm.Forecasting.Timestral (Apache-2.0) — new + +- Inference for Mistral AI's Timestral family — 5M, 28M, 123M, 209M — ported from + `mistralai/timestral` 1.0.0: causal per-patch scaler, gated time/variate transformer with + partial RoPE, Decoupled Horizon Forecasting decoder tokens, monotone 21-quantile head. +- `TimestralForecaster` takes the shared shape (`ForecastJoint`, `Forecast`, both batches) + and applies the reference forecaster's defaults — symmetric averaging, non-negative + clipping, advancing refine frame, quantile cap — so published numbers are reproduced. + The mirrored input of symmetric averaging rides in the same forward pass. +- Verified stage by stage against the reference on all four sizes, with and without the + defaults, including the covariate path and the shared notebook examples. Faster than the + reference on CPU at every size. +- Context up to 32,768 steps, refused beyond; up to 1,024 steps forecast per pass. + +### Tsfm.Forecasting.Timestral.Weights.5M and .28M (Mistral Modified MIT) — new + +- Timestral-5M (21 MB) and -28M (111 MB) embedded at pinned revisions; the build checks + each payload's SHA-256 and refuses one that differs. +- **Revenue-restricted, not open source**: the licence withholds all rights from companies + above USD 20M monthly revenue. It ships verbatim as the package licence file, is + readable at run time (`Timestral5M.License`), and must be accepted on install. +- 123M and 209M are fetched by `scripts/fetch-timestral-checkpoint.sh`, not packaged. + +### Tsfm.Forecasting + +- `IModelWeights`, `DirectoryWeights`, `EmbeddedWeights`: a tokenizer-free checkpoint + abstraction beside `ICheckpoint`, used by the Timestral weights packages. + +### Batching, across all the models - `IQuantileForecaster` (new, in `Tsfm.Forecasting`): the shared `ForecastJoint` / `Forecast` / `Quantiles` shape, plus `ForecastJointBatch` and `ForecastBatch` over many @@ -51,12 +79,13 @@ ### Packaging -All six packages move to 0.2.0: the core gains the interface the model packages implement, +All packages move to 0.2.0: the core gains the interface the model packages implement, and a project reference packs an exact-version dependency. ### Samples and testing -- `samples/GroceryCovariates` runs on TimesFM or t0 (`-- t0-beta`) through the interface. +- `samples/GroceryCovariates` runs on TimesFM, t0 or Timestral (`-- t0-beta`, + `-- Timestral-28M`) through the interface. - Batch tests for TimesFM and t0 run on small random-init models, so they bite in CI without weights; Kronos's run with the integration tests. - `Tsfm.Forecasting.T0.Tests` checks structural invariants on a tiny random-init model, so diff --git a/LICENSE-Timestral-Weights b/LICENSE-Timestral-Weights new file mode 100644 index 0000000..b9396bf --- /dev/null +++ b/LICENSE-Timestral-Weights @@ -0,0 +1,30 @@ +Timestral weights — Mistral AI Modified MIT License +=================================================== + +This package redistributes, unmodified, the Timestral model weights published by Mistral AI +on the Hugging Face Hub. The weights are NOT under this repository's MIT or Apache-2.0 +licences. They are under the licence below, reproduced verbatim as its condition 1 requires. + +Note condition 2: no rights under it are granted to anyone whose company (or employer) had +global consolidated monthly revenue above USD 20 million in the preceding month. That +applies to every consumer of this package, not only to its publisher. + +The small loader class compiled into this package is this repository's own code, MIT +licensed; see LICENSE. + +Source revisions: + Timestral-5M https://huggingface.co/mistralai/Timestral-5M @ bdd20cfc15a3fea0e4cb5ccef37f8d3ca7aa335f + Timestral-28M https://huggingface.co/mistralai/Timestral-28M @ 180698bbf5a50dda3516aa6419798437fb03a430 + +-------------------------------------------------------------------------------- + +Modified MIT License + +Attribution notice: 2026 - Mistral AI + +Permission is hereby granted, free of charge, to any person obtaining a copy of the weights of this model and associated documentation files (the “Model”), to deal in the Model without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Model, and to permit persons to whom the Model is furnished to do so, subject to the following conditions: + +1. The above attribution notice and this permission notice shall be included in all copies or substantial portions of the Model. +2. You are not authorized to exercise any rights under this license if the global consolidated monthly revenue of your company (or that of your employer) exceeds $20 million (or its equivalent in another currency) for the preceding month. This restriction in (b) applies to the Model and any derivatives, modifications, or combined works based on it, whether provided by Mistral AI or by a third party. You may contact Mistral AI (sales@mistral.ai) to request a commercial license, which Mistral AI may grant you at its sole discretion, or choose to use the Model on Mistral AI's hosted services available at https://mistral.ai/. + +THE MODEL IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL MISTRAL AI BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE MODEL OR THE USE OR OTHER DEALINGS IN THE MODEL. diff --git a/NOTICE b/NOTICE index 6db2ace..ef5e785 100644 --- a/NOTICE +++ b/NOTICE @@ -77,3 +77,40 @@ Changes to the derived work, as required by Apache-2.0 section 4(b): No upstream source files and no weights are redistributed. The t0-alpha and t0-beta checkpoints are Apache-2.0, published by The Forecasting Company, and are fetched by scripts/fetch-t0-checkpoint.sh at pinned revisions. + +──────────────────────────────────────────────────────────────────────────────── +Tsfm.Forecasting.Timestral — Apache-2.0, NOT MIT +──────────────────────────────────────────────────────────────────────────────── + +That package contains software derived from Timestral +(https://github.com/mistralai/timestral), Copyright Mistral AI, licensed under the +Apache License, Version 2.0. + +Changes to the derived work, as required by Apache-2.0 section 4(b): + + - The inference path is translated from Python/PyTorch to C#/TorchSharp: the + causal per-patch scaler, instance scaling of known-future covariates, patch + representation, the time/variate transformer with gated attention and + partial rotary embedding, the monotone quantile head, the attention masks + and positions, and the forecaster's refine frame, quantile cap, symmetric + averaging and non-negative clipping. + - Decoder tokens are built for the horizon patches only; the reference also + builds, then discards, one for every earlier patch. + - Symmetric averaging runs the negated input in the same forward pass as the + original, as further groups, instead of in a second call. + - The advancing refine frame accumulates per-patch Welford states sequentially + rather than by a parallel prefix scan; both are float64. + - A context longer than the trained length is refused rather than truncated. + - Past-only covariates as a separate role, categorical target encoding and the + block-by-block continuation beyond one decoded block are not ported. + +──────────────────────────────────────────────────────────────────────────────── +Tsfm.Forecasting.Timestral.Weights.5M and .28M — Mistral AI Modified MIT +──────────────────────────────────────────────────────────────────────────────── + +These packages redistribute, unmodified, the Timestral-5M and Timestral-28M weights +published by Mistral AI. The weights are under Mistral AI's Modified MIT licence, +reproduced verbatim in LICENSE-Timestral-Weights inside each package. It is NOT an +open-source licence: its condition 2 grants no rights to anyone whose company (or +employer) had global consolidated monthly revenue above USD 20 million in the +preceding month. Attribution notice: 2026 - Mistral AI. diff --git a/README.md b/README.md index 02b2c8f..d55843c 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # Tsfm.Forecasting -Native .NET inference for time-series foundation models, on CPU or GPU via +Probabilistic forecasting with native .NET inference for time-series foundation models, on CPU or GPU via TorchSharp. Currently [Kronos](https://github.com/shiyu-coder/Kronos), [TimesFM 3.0](https://github.com/google-research/timesfm) and [t0](https://huggingface.co/theforecastingcompany/t0-beta) (t0-alpha, t0-beta). @@ -18,6 +18,9 @@ NuGet packages: | Tsfm.Forecasting.Kronos.Weights.Mini [![NuGet](https://img.shields.io/nuget/v/Tsfm.Forecasting.Kronos.Weights.Mini.svg)](https://www.nuget.org/packages/Tsfm.Forecasting.Kronos.Weights.Mini) | MIT | Kronos-mini (4.1M) + Tokenizer-2k | | Tsfm.Forecasting.TimesFm [![NuGet](https://img.shields.io/nuget/v/Tsfm.Forecasting.TimesFm.svg)](https://www.nuget.org/packages/Tsfm.Forecasting.TimesFm) | **Apache-2.0** | TimesFM 3.0 — **no weights** | | Tsfm.Forecasting.T0 [![NuGet](https://img.shields.io/nuget/v/Tsfm.Forecasting.T0.svg)](https://www.nuget.org/packages/Tsfm.Forecasting.T0) | **Apache-2.0** | t0-alpha, t0-beta — open weights, fetched separately | +| Tsfm.Forecasting.Timestral [![NuGet](https://img.shields.io/nuget/v/Tsfm.Forecasting.Timestral.svg)](https://www.nuget.org/packages/Tsfm.Forecasting.Timestral) | **Apache-2.0** | Timestral (Mistral AI) — code only | +| Tsfm.Forecasting.Timestral.Weights.5M [![NuGet](https://img.shields.io/nuget/v/Tsfm.Forecasting.Timestral.Weights.5M.svg)](https://www.nuget.org/packages/Tsfm.Forecasting.Timestral.Weights.5M) | **Mistral Modified MIT** — revenue-restricted | Timestral-5M (5.2M) | +| Tsfm.Forecasting.Timestral.Weights.28M [![NuGet](https://img.shields.io/nuget/v/Tsfm.Forecasting.Timestral.Weights.28M.svg)](https://www.nuget.org/packages/Tsfm.Forecasting.Timestral.Weights.28M) | **Mistral Modified MIT** — revenue-restricted | Timestral-28M (27.8M) | Refer to respective package for license and conditions of use (if any). @@ -194,6 +197,66 @@ What differs from TimesFM: values. Running beta under alpha's convention loads cleanly and silently degrades the forecast — which is why they are read rather than assumed. +## Timestral — 5M and 28M batteries included + +> **The weights are not open source.** Mistral AI publishes them under a Modified MIT +> licence whose condition 2 withholds *every* right from anyone whose company — or +> employer — had consolidated monthly revenue above USD 20 million in the preceding month. +> That binds each consumer, not only whoever publishes a package. The weights packages +> carry the licence verbatim (`LICENSE-Timestral-Weights`, and `Timestral5M.License` at +> run time) and require its acceptance on install. The code here is Apache-2.0 and grants +> nothing in the weights. + +[Timestral](https://github.com/mistralai/timestral) is a causal patch transformer with time +and variate attention, trained with Decoupled Horizon Forecasting: forecast-query tokens +read the encoded context, and a monotone head emits 21 quantiles per step. Context up to +32,768 steps; up to 1,024 forecast in one pass. + +``` +dotnet add package Tsfm.Forecasting.Timestral.Weights.5M # or .28M; brings the code package +dotnet add package TorchSharp-cpu +``` + +```csharp +using Tsfm.Forecasting.Timestral; +using Tsfm.Forecasting.Timestral.Weights; + +var ts = TimestralForecaster.Load(Timestral28M.Instance, new Device("cpu")); +double[,] q = ts.ForecastJoint(series, variates: 3, contextBars: 2048, horizon: 96, + knownFuture: [false, true, true]); // [step, 21 quantiles] +``` + +| Checkpoint | Params | Weights | Delivered as | CPU, 512 → 24 steps (.NET / Python) | +|---|---|---|---|---| +| Timestral-5M | 5.2M | 21 MB | `Tsfm.Forecasting.Timestral.Weights.5M` | 11 ms / 19 ms | +| Timestral-28M | 27.8M | 111 MB | `Tsfm.Forecasting.Timestral.Weights.28M` | 37 ms / 48 ms | +| Timestral-123M | 122.6M | 490 MB | `./scripts/fetch-timestral-checkpoint.sh Timestral-123M` | 128 ms / 141 ms | +| Timestral-209M | 209.3M | 837 MB | `./scripts/fetch-timestral-checkpoint.sh Timestral-209M` | 138 ms / 191 ms | + +Timings are 4 CPU cores with the defaults below. The port is faster than the reference +mainly because it folds the mirrored input of symmetric averaging into the same forward +pass, as extra groups, rather than making a second call. The larger two load from a +directory: `TimestralForecaster.Load("checkpoints/Timestral-209M", device)`, or any +`IModelWeights`. The fetch script pins revisions and checks SHA-256; the weights packages +check the same hashes at build time and refuse a payload that differs. + +The forecaster applies what the reference forecaster applies by default, so published +numbers are reproduced: + +- **Symmetric averaging** (`SymmetricAveraging`, on): forecast the negated input too, with + every series negated, and average the reflection in. +- **Non-negative clipping** (`ClipNonNegative`, on): a series whose observed context is + never negative is clipped at zero. Harmless on prices; never triggers on returns. +- **Advancing refine frame**: horizon patch `j` is denormalised with the statistics of the + context plus the block's own medians for patches before it, the frame it was trained in. +- **Quantile cap**: forecasts are held within 10⁴ context scales of the observed range. + +Known-future series are scaled by whole-row statistics, and a series observed only at 0 +and 1 is treated as an indicator. A context beyond 32,768 steps is refused rather than +silently truncated, as Kronos's is. Not ported: past-only covariates as a separate role +(a series not marked known is a joint target), categorical target encoding, and the +block-by-block continuation past 1,024 steps. + ## Batching Each forecaster has a batch form of its single call. Entries share a horizon and nothing @@ -205,9 +268,10 @@ larger forward passes: one per anchor is the dominant cost of a backtest. |---|---|---|---| | t0 | `ForecastJoint`, `Forecast` | `ForecastJointBatch`, `ForecastBatch` | any entries, up to `MaxRowsPerPass` series; shorter contexts are left-padded and masked | | TimesFM | `ForecastJoint`, `Forecast` | `ForecastJointBatch`, `ForecastBatch` | entries of one shape (variates, length, context), up to `MaxBatchPerPass` | +| Timestral | `ForecastJoint`, `Forecast` | `ForecastJointBatch`, `ForecastBatch` | as t0, up to `MaxRowsPerPass` rows; the mirrored half of symmetric averaging counts | | Kronos | `Infer`, `InferPath` (windows along one series) | `InferBatch` (final window of many series) | `batch` windows × `rollouts` | -The t0 and TimesFM batches are on `IQuantileForecaster`, take `JointSeries` entries (or +The t0, Timestral and TimesFM batches are on `IQuantileForecaster`, take `JointSeries` entries (or OHLCVA windows) as `ReadOnlyMemory`, and write `[entry][step][quantile]` into a span sized by `BatchOutputLength`. Anchors along one history are slices of one array, so building a backtest batch copies nothing: @@ -230,7 +294,9 @@ make cross-sectional comparisons less noisy but leave their errors correlated. T host policy, `KronosForecaster.Seed`: `SamplingSeeds.PerSeries` gives each `KronosSeries.SeriesKey` its own stream, and `SamplingSeeds.Salted` derives a fresh, reproducible run from either. -TimesFM and t0 emit quantiles directly, with no rollouts and nothing to aggregate. t0's +TimesFM, t0 and Timestral emit quantiles directly, with no rollouts and nothing to +aggregate. Timestral's symmetric averaging is the nearest thing: the input and its negation +travel as two groups of one pass, and are combined per series afterwards. t0's upstream continuation past 1024 steps (not ported) would be the same pattern again: each requested quantile becomes a path, paths are folded into the batch as extra groups, and they are reduced per series to quantiles. @@ -301,6 +367,18 @@ QK norm precedes the rotation, where TimesFM's follows it; and a padding patch h to attend to, so attention must return zero there, not NaN, or every row reading it is poisoned. +**Timestral.** All four sizes agree with `mistralai/timestral` 1.0.0 stage by stage: +scaling, masks and positions bit-exact, every other stage within 1.5e-6 relative, and +forecasts within 9e-7 — with the forecaster's defaults and with them off, and including the covariate path, +the refine frame and the electricity-price and M4 examples shared with t0. The reference +builds a decoder token for every patch of the window and discards those before the +horizon; the port builds only the horizon's, which is exact because the discarded ones +cannot reach the rest. Details that would each have forecast plausibly if missed: the +fused QKV projection is laid out `(qkv, head_dim, heads)`, not head-major; RoPE turns only +the first half of each head, in the split-half layout; SwiGLU's SiLU falls on the second +half, where t0's is on the first; and statistics are float64 Welford, which MPS cannot do, +so they run on the CPU. + ```bash ./scripts/fetch-checkpoints.sh # Kronos dotnet test tests/Tsfm.Forecasting.Kronos.Tests -c Release @@ -316,6 +394,11 @@ dotnet run --project tests/Tsfm.Forecasting.TimesFm.Parity -c Release ./.venv/bin/python reference/dump_t0_parity.py checkpoints/t0-alpha ./.venv/bin/python reference/dump_t0_parity.py checkpoints/t0-beta dotnet run --project tests/Tsfm.Forecasting.T0.Parity -c Release -- t0-alpha t0-beta + +./scripts/fetch-timestral-checkpoint.sh all # Timestral, revenue-restricted weights +./.venv/bin/pip install "git+https://github.com/mistralai/timestral.git@a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff" +for m in 5M 28M 123M 209M; do ./.venv/bin/python reference/dump_timestral_parity.py checkpoints/Timestral-$m; done +dotnet run --project tests/Tsfm.Forecasting.Timestral.Parity -c Release -- Timestral-5M Timestral-28M Timestral-123M Timestral-209M ``` ## Attribution @@ -336,5 +419,10 @@ Copyright 2026 The Forecasting Company, Apache-2.0, which in turn builds on Data [Chronos-2](https://github.com/amazon-science/chronos-forecasting), both Apache-2.0. No weights or upstream sources are redistributed; checkpoints are fetched at pinned revisions. +**Timestral.** Derived from [mistralai/timestral](https://github.com/mistralai/timestral), +Copyright Mistral AI, Apache-2.0. The 5M and 28M weights are redistributed unmodified +under Mistral AI's Modified MIT licence, which is reproduced in full in the weights packages; +the 123M and 209M weights are not redistributed. + This project is an independent port and is not affiliated with or endorsed by any of them. See NOTICE for the full attribution and the statement of changes. diff --git a/Tsfm.Forecasting.slnx b/Tsfm.Forecasting.slnx index b8645c4..bfc739b 100644 --- a/Tsfm.Forecasting.slnx +++ b/Tsfm.Forecasting.slnx @@ -5,14 +5,20 @@ + + + + + + diff --git a/reference/dump_timestral_parity.py b/reference/dump_timestral_parity.py new file mode 100644 index 0000000..ff4fc89 --- /dev/null +++ b/reference/dump_timestral_parity.py @@ -0,0 +1,172 @@ +"""Dump reference tensors so the .NET Timestral port can be bisected stage by stage. + +Runs mistralai/timestral (pinned below) against a local checkpoint and records each stage +of one covariate-path block — packed rows, scaling, patch embedding, masks, positions, +selected layers, head, denormalisation, refine frame — then end-to-end `forecast` outputs +with the reference forecaster's defaults (symmetric averaging and non-negative clipping +on) and with both off. + +The reference decodes a decoder token for every patch of the window and discards those +before the horizon; the port builds only the horizon's. They cannot influence the +horizon's, so decoder stages are recorded for the horizon tokens only. + + pip install git+https://github.com/mistralai/timestral.git@a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff pandas pyarrow + python reference/dump_timestral_parity.py checkpoints/Timestral-5M +""" + +import math +import pathlib +import sys + +import numpy as np +import torch +from safetensors.torch import save_file + +from timestral import TimestralForecaster +from timestral.layers import TwoStreamMask +from timestral.masks import ( + build_two_stream_positions, + build_two_stream_time_mask, + build_two_stream_variate_mask, + compute_enc_key_validity, + group_first_observed_patch, +) +from timestral.model import ROLE_KNOWN_COV, ROLE_TARGET +from timestral.scaler import denormalize_quantiles + +ROOT = pathlib.Path(__file__).resolve().parent.parent +T, H = 150, 40 # not patch multiples: left padding, and a two-patch horizon for the refine frame +DATA = ROOT / "reference" / "data" + + +def synthetic() -> tuple[torch.Tensor, torch.Tensor]: + """Two targets of distinct scale, a numeric known covariate and a 0/1 indicator, with + one NaN in a target and one in the numeric covariate.""" + g = torch.Generator().manual_seed(20261019) + t = torch.arange(T + H, dtype=torch.float32) + ctx = torch.stack([ + 100.0 + 5.0 * torch.sin(t[:T] / 7.0) + 0.02 * t[:T], + -3.0 + 0.25 * torch.cos(t[:T] / 5.0), + ]) + 0.05 * torch.randn(2, T, generator=g) + fut = torch.stack([20.0 + 2.0 * torch.sin(t / 7.0 + 0.3), (t % 24 < 6).float()]) + ctx[0, 40] = float("nan") + fut[0, 90] = float("nan") + return ctx, fut + + +def stages(fc: TimestralForecaster, ctx: torch.Tensor, fut: torch.Tensor) -> dict[str, torch.Tensor]: + model, c = fc.model, fc.config + ps = c.patch_size + inputs, mask = fc._as_context(ctx[None], None) + known = fc._as_covariates(fut[None], 1, T + H, "known_covariates") + rows, mask, role_ids, kf, kfm, _ = fc._assemble_rows(inputs, mask, H, None, known) + b, v, t = rows.shape + context, mask = rows.reshape(b * v, t), mask.reshape(b * v, t) + group_ids = torch.zeros(b * v, dtype=torch.long) + h = math.ceil(H / ps) + full = torch.cat([context, kf[:, : h * ps]], -1) + full_obs = torch.cat([mask, kfm[:, : h * ps]], -1) + out = {"buf_values": full, "buf_observed": full_obs.float(), "buf_groups": group_ids.float(), + "buf_roles": role_ids.float(), "ctx_patches": torch.tensor([t // ps], dtype=torch.float32)} + + n = full.shape[1] // ps + ll = n - h + known_rows = role_ids == ROLE_KNOWN_COV + scaled, loc, scale = model.scale(full, full_obs) + scaled, loc, scale = model._known_cov_scaling(full, full_obs, scaled, loc, scale, known_rows) + out |= {"scaled": scaled, "loc": loc, "scale": scale} + x_enc = model.representation(scaled, full_obs) + out["patch_embedding"] = x_enc + + obs = full_obs.unflatten(-1, (n, ps)).any(-1) + key_valid = compute_enc_key_validity(obs, c.attend_missing_patches) + s_rows = torch.full((b * v,), ll, dtype=torch.long) + time_mask = build_two_stream_time_mask(s_rows, n, c.block_size, key_valid, bidirectional_rows=known_rows) + variate_mask = build_two_stream_variate_mask( + group_ids, key_valid, n, target_rows=role_ids == ROLE_TARGET, bidirectional_rows=known_rows) + uniq, inv = torch.unique(group_ids, return_inverse=True) + first = group_first_observed_patch(obs, inv, uniq.numel())[inv] + positions = build_two_stream_positions(n, first) + + # Keep [Enc ; Dec(horizon)] — the slots the port builds. + keep = torch.cat([torch.arange(n), n + torch.arange(ll, n)]) + out["time_mask"] = time_mask[:, :, keep][:, :, :, keep].float() + out["variate_mask"] = variate_mask[keep].float() + out["positions"] = positions[:, keep].float() + + x = torch.cat([x_enc, model.representation.forecast_tokens(b * v, n)], 1) + two = TwoStreamMask(mask=time_mask, enc_len=n) + tf = model.transformer + cos, sin = tf.rope(x, positions) + layers = tf.layers + for i, layer in enumerate(layers): + x = layer(x, two, variate_mask, cos, sin) + if i in (0, 3, len(layers) - 1): + out[f"layer{i}_output"] = x[:, keep] + x = tf.out_norm(x) + out["transformer_output"] = x[:, keep] + preds = model.head(x[:, n + ll :]) + out["head_output"] = preds + anchor_loc = loc.unflatten(-1, (n, ps))[..., 0][:, ll - 1 : ll] + anchor_scale = scale.unflatten(-1, (n, ps))[..., 0][:, ll - 1 : ll] + raw = denormalize_quantiles(preds, anchor_loc.unsqueeze(-1), anchor_scale.unsqueeze(-1), c.use_asinh) + out["block_output"] = raw + out["refined"] = fc._refine_block(raw, context, mask, role_ids == ROLE_TARGET) + return out + + +def notebook_cases() -> dict[str, torch.Tensor]: + import pandas as pd + + m4 = pd.read_csv(DATA / "m4_hourly_train.csv") + sid = m4["item_id"].iloc[0] + full = m4.loc[m4["item_id"] == sid, "target"].to_numpy(dtype=np.float32) + train = pd.read_parquet(DATA / "electricity_price_train.parquet") + test = pd.read_parquet(DATA / "electricity_price_test.parquet") + cols = [c for c in test.columns if c not in ("id", "timestamp", "target")] + cov = np.concatenate([train[cols].to_numpy(dtype=np.float32)[-512:].T, test[cols].to_numpy(dtype=np.float32).T], 1) + return { + "m4_ctx": torch.as_tensor(full[:512]), + "m4_truth": torch.as_tensor(full[512:560]), + "ep_ctx": torch.as_tensor(train["target"].to_numpy(dtype=np.float32)[-512:]), + "ep_cov": torch.as_tensor(cov), + "ep_truth": torch.as_tensor(test["target"].to_numpy(dtype=np.float32)), + } + + +def main() -> None: + ckpt = pathlib.Path(sys.argv[1]).resolve() + fc = TimestralForecaster.from_pretrained(str(ckpt)) + torch.set_grad_enabled(False) + + ctx, fut = synthetic() + out = {"in_ctx": ctx, "in_fut": fut} + out |= stages(fc, ctx, fut) + + def predict(c, h, known=None, raw=False): + f = fc.forecast(c[None], h, known_covariates=None if known is None else known[None], + symmetric_averaging=not raw, clip_nonnegative=not raw) + return f.quantiles[0] # (V, h, Q) + + out["predict_joint"] = predict(ctx, H, fut) + out["predict_joint_raw"] = predict(ctx, H, fut, raw=True) + out["predict_targets_only"] = predict(ctx, H) + out["predict_aligned"] = predict(ctx[0, 64:128][None], 32)[0] + + nb = notebook_cases() + out |= nb + out["m4_predict"] = predict(nb["m4_ctx"][None], 48)[0] + out["ep_predict"] = predict(nb["ep_ctx"][None], 24)[0] + out["ep_predict_cov"] = predict(nb["ep_ctx"][None], 24, nb["ep_cov"])[0] + mid = fc.quantile_levels.index(0.5) + for name in ("ep_predict", "ep_predict_cov"): + err = (out[name][:, mid] - nb["ep_truth"]).abs().mean().item() + print(f" {name:16s} |median-truth| = {err:.4f}") + + dest = ROOT / "reference" / f"{ckpt.name}-parity.safetensors" + save_file({k: v.contiguous().float() for k, v in out.items()}, str(dest)) + print(f" wrote {dest} ({len(out)} tensors)") + + +if __name__ == "__main__": + main() diff --git a/samples/GroceryCovariates/GroceryCovariates.csproj b/samples/GroceryCovariates/GroceryCovariates.csproj index 2623283..57f0f1f 100644 --- a/samples/GroceryCovariates/GroceryCovariates.csproj +++ b/samples/GroceryCovariates/GroceryCovariates.csproj @@ -13,6 +13,8 @@ + + diff --git a/samples/GroceryCovariates/Program.cs b/samples/GroceryCovariates/Program.cs index dcb1a84..335d584 100644 --- a/samples/GroceryCovariates/Program.cs +++ b/samples/GroceryCovariates/Program.cs @@ -12,11 +12,14 @@ // // dotnet run --project samples/GroceryCovariates # TimesFM 3.0 // dotnet run --project samples/GroceryCovariates -- t0-beta # or t0-alpha +// dotnet run --project samples/GroceryCovariates -- Timestral-5M # embedded; or -28M, -123M, -209M using TorchSharp; using Tsfm.Forecasting; using Tsfm.Forecasting.T0; +using Tsfm.Forecasting.Timestral; +using Tsfm.Forecasting.Timestral.Weights; using Tsfm.Forecasting.TimesFm; using static TorchSharp.torch; @@ -44,7 +47,20 @@ IQuantileForecaster forecaster; int context, total, pad; -if (model.StartsWith("t0")) +if (model.StartsWith("Timestral")) +{ + // Like t0, Timestral pads a short context itself. 5M and 28M are embedded; the larger + // two load from a fetched directory. + forecaster = model switch + { + "Timestral-5M" => TimestralForecaster.Load(Timestral5M.Instance, device), + "Timestral-28M" => TimestralForecaster.Load(Timestral28M.Instance, device), + _ => TimestralForecaster.Load(checkpoint, device), + }; + (context, total, pad) = (observed, observed + horizon, 0); + Console.WriteLine($"{model} on {device.type}: {observed} days observed, {horizon} forecast\n"); +} +else if (model.StartsWith("t0")) { // t0 left-pads a short context onto its patch grid itself, and masks the padding. forecaster = T0Forecaster.Load(checkpoint, device); @@ -94,5 +110,5 @@ Console.WriteLine(); } -// The column nearest a level: the grids differ (TimesFM 9, t0-alpha 5, t0-beta 21 levels). +// The column nearest a level: the grids differ (TimesFM 9, t0-alpha 5, t0-beta and Timestral 21). int Level(double level) => Enumerable.Range(0, quantiles.Count).MinBy(i => Math.Abs(quantiles[i] - level)); diff --git a/scripts/fetch-timestral-checkpoint.sh b/scripts/fetch-timestral-checkpoint.sh new file mode 100755 index 0000000..eb68c90 --- /dev/null +++ b/scripts/fetch-timestral-checkpoint.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +# Fetch the Timestral checkpoints: 5M (21 MB), 28M (111 MB), 123M (490 MB), 209M (837 MB). +# +# THE WEIGHTS ARE NOT OPEN SOURCE IN THE OSI SENSE. Mistral AI publishes them under a +# Modified MIT licence that withholds every right from anyone whose company (or employer) +# had consolidated monthly revenue above USD 20 million in the preceding month. The +# Apache-2.0 code in this repository grants nothing in the weights. Read the LICENSE file +# fetched beside each checkpoint before using it. +# +# 5M and 28M are also embedded in the Tsfm.Forecasting.Timestral.Weights.* packages; their +# builds read from here. Revisions are pinned and payloads checked against their SHA-256. +# +# ./scripts/fetch-timestral-checkpoint.sh # 5M and 28M, the embedded pair +# ./scripts/fetch-timestral-checkpoint.sh all # all four +# ./scripts/fetch-timestral-checkpoint.sh Timestral-209M # one +set -euo pipefail + +ROOT="${TIMESTRAL_CHECKPOINT_ROOT:-$(cd "$(dirname "$0")/.." && pwd)/checkpoints}" +# macOS ships shasum, not sha256sum. +if command -v sha256sum > /dev/null; then SHA256="sha256sum"; else SHA256="shasum -a 256"; fi + +fetch () { + local name=$1 rev=$2 sha=$3 + local dir="$ROOT/$name" + local base="https://huggingface.co/mistralai/$name/resolve/$rev" + mkdir -p "$dir" + for f in config.json LICENSE README.md model.safetensors; do + if [ -s "$dir/$f" ]; then echo " have $name/$f"; continue; fi + echo " fetch $name/$f" + curl -fsSL --retry 3 -o "$dir/$f.part" "$base/$f" + mv "$dir/$f.part" "$dir/$f" + done + echo "$sha $dir/model.safetensors" | $SHA256 -c --quiet - \ + || { echo " $name/model.safetensors does not match the pinned revision $rev"; exit 1; } + echo " ok $name @ ${rev:0:12}" +} + +want="${1:-embedded}" +pick () { [[ $want == all || $want == "$1" || ($want == embedded && $2 == yes) ]]; } + +echo "Timestral weights: Modified MIT, revenue-restricted. See LICENSE beside each checkpoint." +pick Timestral-5M yes && fetch Timestral-5M \ + bdd20cfc15a3fea0e4cb5ccef37f8d3ca7aa335f 7ace108c4987e7f2a99a75a445240019d1d31c2117139d1bd05ffdfd25aca738 +pick Timestral-28M yes && fetch Timestral-28M \ + 180698bbf5a50dda3516aa6419798437fb03a430 dfc83c6220937ce5fc6687b35a4df94cd25e28a2cfe6fb92a23dad6301b2015d +pick Timestral-123M no && fetch Timestral-123M \ + 5a147287addda85f6acd8da3fb8858fe9af7e8f7 599fafa73cadbbc800b5a346ca900fa741ee3e5fd22e3b4a4430e151e74f4bc2 +pick Timestral-209M no && fetch Timestral-209M \ + 9874d252ab0aae16f7f7c999f98f98d6ad7d8365 1456effb9927e3daa8c7234f10497e3424233869245a9c72b8ecebaf7ad2b7fa +exit 0 diff --git a/scripts/verify-packages.sh b/scripts/verify-packages.sh index 00e68ba..b21c58a 100755 --- a/scripts/verify-packages.sh +++ b/scripts/verify-packages.sh @@ -14,7 +14,18 @@ for pkg in "$DIR"/*.nupkg; do echo "── $name (${size} MB)" contents=$(unzip -l "$pkg") - for required in LICENSE README.md; do + required_files="LICENSE README.md" + # Revenue-restricted weights must carry their own licence, and demand acceptance of it. + case "$name" in + Tsfm.Forecasting.Timestral.Weights.*) + required_files="$required_files LICENSE-Timestral-Weights" + if unzip -p "$pkg" '*.nuspec' | grep -q 'true'; then + echo " ok requires licence acceptance" + else + echo " MISSING requireLicenseAcceptance"; fail=1 + fi ;; + esac + for required in $required_files; do if grep -q "$required" <<< "$contents"; then echo " ok carries $required" else @@ -30,7 +41,7 @@ for pkg in "$DIR"/*.nupkg; do # A weights package with no embedded payload builds and publishes cleanly, then fails # at load for every consumer — worth an explicit check. case "$name" in - Tsfm.Forecasting.Kronos.Weights.*) + Tsfm.Forecasting.Kronos.Weights.*|Tsfm.Forecasting.Timestral.Weights.*) if [ "$size" -lt 10 ]; then echo " EMPTY weights package is only ${size} MB; the checkpoint did not embed"; fail=1 else diff --git a/src/Tsfm.Forecasting.Timestral.Weights.28M/Timestral28M.cs b/src/Tsfm.Forecasting.Timestral.Weights.28M/Timestral28M.cs new file mode 100644 index 0000000..e6e8713 --- /dev/null +++ b/src/Tsfm.Forecasting.Timestral.Weights.28M/Timestral28M.cs @@ -0,0 +1,37 @@ +using System.Reflection; + +using Tsfm.Forecasting; + +namespace Tsfm.Forecasting.Timestral.Weights; + +/// +/// Timestral-28M (27.8M), embedded. Context up to 32768 steps; 1024 per decoded block. +/// +/// Pinned revision 180698bbf5a50dda3516aa6419798437fb03a430, embedded rather than fetched, so the weights a build +/// was compiled against cannot drift underneath it. The build refuses a payload whose +/// SHA-256 differs from the revision's. +/// +/// Not open source. Mistral AI's Modified MIT licence () +/// grants no rights to anyone whose company, or employer, had consolidated monthly revenue +/// above USD 20 million in the preceding month. +/// +public sealed class Timestral28M : EmbeddedWeights +{ + public static readonly Timestral28M Instance = new(); + + public override string Name => "mistralai/Timestral-28M@180698bbf5a50dda3516aa6419798437fb03a430"; + + /// The weights' licence, verbatim as Mistral AI publishes it. + public static string License + { + get + { + using var reader = new StreamReader(typeof(Timestral28M).Assembly.GetManifestResourceStream("timestral.license")!); + return reader.ReadToEnd(); + } + } + + protected override Assembly Host => typeof(Timestral28M).Assembly; + protected override string ModelResource => "timestral.model"; + protected override string ModelConfigResource => "timestral.config"; +} diff --git a/src/Tsfm.Forecasting.Timestral.Weights.28M/Tsfm.Forecasting.Timestral.Weights.28M.csproj b/src/Tsfm.Forecasting.Timestral.Weights.28M/Tsfm.Forecasting.Timestral.Weights.28M.csproj new file mode 100644 index 0000000..b85615e --- /dev/null +++ b/src/Tsfm.Forecasting.Timestral.Weights.28M/Tsfm.Forecasting.Timestral.Weights.28M.csproj @@ -0,0 +1,43 @@ + + + + Tsfm.Forecasting.Timestral.Weights.28M + Tsfm.Forecasting.Timestral.Weights + Tsfm.Forecasting.Timestral.Weights.28M + Timestral-28M (27.8M) weights for Tsfm.Forecasting.Timestral, embedded as assembly resources (~111 MB), redistributed unmodified from Mistral AI at the revision recorded in LICENSE-Timestral-Weights. NOT open source: Mistral's Modified MIT licence grants no rights to anyone whose company or employer exceeds USD 20M monthly revenue. + timestral;mistral;timeseries;forecasting;weights;model + true + + LICENSE-Timestral-Weights + true + + $(MSBuildThisFileDirectory)../../checkpoints + $(TimestralCheckpointRoot)/Timestral-28M + dfc83c6220937ce5fc6687b35a4df94cd25e28a2cfe6fb92a23dad6301b2015d + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/Tsfm.Forecasting.Timestral.Weights.5M/Timestral5M.cs b/src/Tsfm.Forecasting.Timestral.Weights.5M/Timestral5M.cs new file mode 100644 index 0000000..bcb6f46 --- /dev/null +++ b/src/Tsfm.Forecasting.Timestral.Weights.5M/Timestral5M.cs @@ -0,0 +1,37 @@ +using System.Reflection; + +using Tsfm.Forecasting; + +namespace Tsfm.Forecasting.Timestral.Weights; + +/// +/// Timestral-5M (5.2M), embedded. Context up to 32768 steps; 1024 per decoded block. +/// +/// Pinned revision bdd20cfc15a3fea0e4cb5ccef37f8d3ca7aa335f, embedded rather than fetched, so the weights a build +/// was compiled against cannot drift underneath it. The build refuses a payload whose +/// SHA-256 differs from the revision's. +/// +/// Not open source. Mistral AI's Modified MIT licence () +/// grants no rights to anyone whose company, or employer, had consolidated monthly revenue +/// above USD 20 million in the preceding month. +/// +public sealed class Timestral5M : EmbeddedWeights +{ + public static readonly Timestral5M Instance = new(); + + public override string Name => "mistralai/Timestral-5M@bdd20cfc15a3fea0e4cb5ccef37f8d3ca7aa335f"; + + /// The weights' licence, verbatim as Mistral AI publishes it. + public static string License + { + get + { + using var reader = new StreamReader(typeof(Timestral5M).Assembly.GetManifestResourceStream("timestral.license")!); + return reader.ReadToEnd(); + } + } + + protected override Assembly Host => typeof(Timestral5M).Assembly; + protected override string ModelResource => "timestral.model"; + protected override string ModelConfigResource => "timestral.config"; +} diff --git a/src/Tsfm.Forecasting.Timestral.Weights.5M/Tsfm.Forecasting.Timestral.Weights.5M.csproj b/src/Tsfm.Forecasting.Timestral.Weights.5M/Tsfm.Forecasting.Timestral.Weights.5M.csproj new file mode 100644 index 0000000..d170df6 --- /dev/null +++ b/src/Tsfm.Forecasting.Timestral.Weights.5M/Tsfm.Forecasting.Timestral.Weights.5M.csproj @@ -0,0 +1,43 @@ + + + + Tsfm.Forecasting.Timestral.Weights.5M + Tsfm.Forecasting.Timestral.Weights + Tsfm.Forecasting.Timestral.Weights.5M + Timestral-5M (5.2M) weights for Tsfm.Forecasting.Timestral, embedded as assembly resources (~21 MB), redistributed unmodified from Mistral AI at the revision recorded in LICENSE-Timestral-Weights. NOT open source: Mistral's Modified MIT licence grants no rights to anyone whose company or employer exceeds USD 20M monthly revenue. + timestral;mistral;timeseries;forecasting;weights;model + true + + LICENSE-Timestral-Weights + true + + $(MSBuildThisFileDirectory)../../checkpoints + $(TimestralCheckpointRoot)/Timestral-5M + 7ace108c4987e7f2a99a75a445240019d1d31c2117139d1bd05ffdfd25aca738 + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/Tsfm.Forecasting.Timestral/TimestralForecaster.cs b/src/Tsfm.Forecasting.Timestral/TimestralForecaster.cs new file mode 100644 index 0000000..f69d48a --- /dev/null +++ b/src/Tsfm.Forecasting.Timestral/TimestralForecaster.cs @@ -0,0 +1,532 @@ +using System.Runtime.InteropServices; + +using static TorchSharp.torch; + +namespace Tsfm.Forecasting.Timestral; + +/// +/// Quantile forecasts from a Timestral checkpoint, in the shape shared with +/// TimesFmForecaster and T0Forecaster. +/// +/// Around the network it applies what the reference forecaster applies by default, +/// so published numbers are reproduced: each horizon patch is moved into the advancing +/// refine frame it was trained in, forecasts are capped at 10⁴ context scales beyond the +/// observed range, the forecast of the negated input is reflected and averaged in +/// (), and series whose context is non-negative are +/// clipped at zero (). The mirrored input rides in the same +/// forward pass as the original, as further groups, rather than costing a second call. +/// +/// Every series not marked known is a target, forecast jointly. Known-future series +/// are conditioned on, not forecast, and never read a target. +/// +/// Inference is not thread safe. +public sealed class TimestralForecaster(TimestralModel model, Device device, string name = "") + : IQuantileForecaster +{ + private const int CloseChannel = 3; // open, high, low, close, volume, amount + private const int Channels = 6; + private const double QuantileCap = 1e4; + + private readonly TimestralConfig _c = model.Config; + + public TimestralConfig Config => _c; + + public TimestralModel Model => model; + + /// Checkpoint provenance, from the weights it was loaded from. + public string CheckpointName => name; + + /// + public IReadOnlyList Quantiles => _c.Quantiles; + + /// Longest horizon decoded, one block: 1024 steps for the published checkpoints. + /// The reference continues block by block, feeding medians back; that is not ported. + public int MaxHorizon => _c.MaxHorizon; + + /// Longest context, the trained length: 32768 for the published checkpoints. + /// The reference silently keeps the most recent steps; a longer context is refused here. + public int MaxContext => _c.ContextLength; + + /// Average with the reflected forecast of the negated input — every series, + /// covariates included. On by default, as in the reference. + public bool SymmetricAveraging { get; set; } = true; + + /// Clip at zero the forecast of a series whose observed context is all + /// non-negative. On by default, as in the reference. + public bool ClipNonNegative { get; set; } = true; + + /// Rows — series, summed over entries and doubled under + /// — in one forward pass. Variate masks grow with its + /// square and time masks with the square of the context, so long contexts want less. + public int MaxRowsPerPass { get; set; } = 64; + + public static TimestralForecaster Load(IModelWeights weights, Device device) + => new(TimestralModel.Load(weights, device), device, weights.Name); + + /// Load from a snapshot directory holding config.json and + /// model.safetensors. + public static TimestralForecaster Load(string checkpointDir, Device device) + => Load(new DirectoryWeights(checkpointDir, howToFetch: + "Fetch the weights with scripts/fetch-timestral-checkpoint.sh, or download config.json and " + + "model.safetensors from https://huggingface.co/mistralai/Timestral-209M (or another size); " + + "5M and 28M also ship in Tsfm.Forecasting.Timestral.Weights.*."), device); + + /// + /// Joint forecast over several series, any of which may be known into the future. + /// + /// Variate-major [variates x total]: series v at step + /// t is series[v * total + t]. Steps from + /// on are read only for known variates, and only the first . + /// NaN is a missing observation. + /// Series supplied. + /// Observed steps, 1... + /// Steps ahead, 1... + /// Per variate, or empty for none. A variate marked known is + /// NOT a forecast target. + /// Which series the returned quantiles describe. + /// [horizon, 21] in the units of the target series. + public double[,] ForecastJoint( + ReadOnlySpan series, int variates, int contextBars, int horizon, + ReadOnlySpan knownFuture = default, int targetVariate = 0) + { + ValidateHorizon(horizon); + var total = ValidateJoint(series.Length, variates, contextBars, horizon, knownFuture, targetVariate, ""); + var result = new double[horizon, _c.NumQuantiles]; + var layout = new Layout(_c.PatchSize, contextBars, horizon, variates, Mirrored); + var entry = AddJoint(layout, series, total, variates, contextBars, knownFuture, targetVariate); + Run(layout, [entry], [0], MemoryMarshal.CreateSpan(ref result[0, 0], result.Length)); + return result; + } + + /// Forward-return quantiles from OHLCVA bars, every channel a target. + /// Row-major [L x 6]: open, high, low, close, volume, amount. + /// Steps ahead, 1... + /// [horizon, 21] of cumulative return relative to the anchor close. + public double[,] Forecast(ReadOnlySpan ohlcva, int horizon) + { + ValidateHorizon(horizon); + ValidateBars(ohlcva.Length, ""); + var result = new double[horizon, _c.NumQuantiles]; + var flat = MemoryMarshal.CreateSpan(ref result[0, 0], result.Length); + var layout = new Layout(_c.PatchSize, ohlcva.Length / Channels, horizon, Channels, Mirrored); + Run(layout, [AddBars(layout, ohlcva)], [0], flat); + ToReturns(flat, ohlcva); + return result; + } + + /// + public int BatchOutputLength(int entries, int horizon) => entries * horizon * _c.NumQuantiles; + + /// + /// over many entries, each its own group, in passes of at + /// most rows. Contexts may differ: shorter ones are + /// left-padded with unobserved patches, which are neither keys nor positions, so an + /// entry's forecast does not depend on its companions. + /// + /// Entries, each with the constraints of . + /// Steps ahead, shared. + /// Receives [entry][step][quantile]; sized by + /// . + public void ForecastJointBatch(ReadOnlySpan batch, int horizon, Span quantiles) + { + ValidateHorizon(horizon); + ValidateDestination(quantiles.Length, batch.Length, horizon, nameof(quantiles)); + var totals = new int[batch.Length]; + var contexts = new int[batch.Length]; + var rows = new int[batch.Length]; + var variates = new int[batch.Length]; + for (var i = 0; i < batch.Length; i++) + { + var e = batch[i]; + totals[i] = ValidateJoint(e.Series.Length, e.Variates, e.ContextBars, horizon, e.KnownFuture.Span, + e.TargetVariate, $"batch[{i}]: "); + (contexts[i], variates[i], rows[i]) = (e.ContextBars, e.Variates, e.Variates * (Mirrored ? 2 : 1)); + } + + foreach (var pass in Passes(contexts, rows)) + { + var layout = new Layout(_c.PatchSize, pass.Max(i => contexts[i]), horizon, pass.Sum(i => variates[i]), + Mirrored); + var entries = new Entry[pass.Length]; + for (var j = 0; j < pass.Length; j++) + { + var e = batch[pass[j]]; + entries[j] = AddJoint(layout, e.Series.Span, totals[pass[j]], e.Variates, e.ContextBars, + e.KnownFuture.Span, e.TargetVariate); + } + Run(layout, entries, pass, quantiles); + } + } + + /// over many OHLCVA windows, batched as + /// is. + /// One row-major [L x 6] window per entry. + /// Steps ahead, shared. + /// Receives [entry][step][quantile] of cumulative return; + /// sized by . + public void ForecastBatch(ReadOnlySpan> ohlcva, int horizon, Span returns) + { + ValidateHorizon(horizon); + ValidateDestination(returns.Length, ohlcva.Length, horizon, nameof(returns)); + var contexts = new int[ohlcva.Length]; + for (var i = 0; i < ohlcva.Length; i++) + { + ValidateBars(ohlcva[i].Length, $"ohlcva[{i}]: "); + contexts[i] = ohlcva[i].Length / Channels; + } + + var rows = Enumerable.Repeat(Channels * (Mirrored ? 2 : 1), ohlcva.Length).ToArray(); + foreach (var pass in Passes(contexts, rows)) + { + var layout = new Layout(_c.PatchSize, pass.Max(i => contexts[i]), horizon, pass.Length * Channels, Mirrored); + var entries = new Entry[pass.Length]; + for (var j = 0; j < pass.Length; j++) entries[j] = AddBars(layout, ohlcva[pass[j]].Span); + Run(layout, entries, pass, returns); + } + + var block = horizon * _c.NumQuantiles; + for (var i = 0; i < ohlcva.Length; i++) ToReturns(returns.Slice(i * block, block), ohlcva[i].Span); + } + + private bool Mirrored => SymmetricAveraging; + + /// Entry indices ordered by context, cut wherever the next would exceed the row + /// budget; an entry wider than the budget runs alone. + private List Passes(int[] contexts, int[] rows) + { + var order = Enumerable.Range(0, contexts.Length).OrderBy(i => contexts[i]).ToArray(); + var passes = new List(); + var (from, used) = (0, 0); + for (var j = 0; j < order.Length; j++) + { + if (j > from && used + rows[order[j]] > MaxRowsPerPass) + { + passes.Add(order[from..j]); + (from, used) = (j, 0); + } + used += rows[order[j]]; + } + if (from < order.Length) passes.Add(order[from..]); + return passes; + } + + /// An entry's target row in the layout. + private readonly record struct Entry(int TargetRow); + + /// Targets first, then known-future rows: the reference's row order. + private static Entry AddJoint(Layout layout, ReadOnlySpan series, int total, int variates, int contextBars, + ReadOnlySpan knownFuture, int targetVariate) + { + var group = layout.NextGroup(); + var target = -1; + for (var pass = 0; pass < 2; pass++) + for (var v = 0; v < variates; v++) + { + var known = !knownFuture.IsEmpty && knownFuture[v]; + if (known != (pass == 1)) continue; + var src = series.Slice(v * total, total); + var row = layout.AddRow(group, src[..contextBars], known ? src.Slice(contextBars, layout.Horizon) : default); + if (v == targetVariate) target = row; + } + return new Entry(target); + } + + private static Entry AddBars(Layout layout, ReadOnlySpan ohlcva) + { + var group = layout.NextGroup(); + var bars = ohlcva.Length / Channels; + Span column = bars <= 1024 ? stackalloc float[bars] : new float[bars]; + var close = -1; + for (var c = 0; c < Channels; c++) + { + for (var t = 0; t < bars; t++) column[t] = ohlcva[t * Channels + c]; + var row = layout.AddRow(group, column, default); + if (c == CloseChannel) close = row; + } + return new Entry(close); + } + + private static void ToReturns(Span prices, ReadOnlySpan ohlcva) + { + var anchorClose = ohlcva[^(Channels - CloseChannel)]; + for (var i = 0; i < prices.Length; i++) prices[i] = prices[i] / anchorClose - 1.0; + } + + /// + /// One forward pass over every row of the layout — the originals and, under symmetric + /// averaging, their negations — then the reference's post-processing per entry. Writes + /// entry j at slot slots[j] of . + /// + private void Run(Layout layout, Entry[] entries, ReadOnlySpan slots, Span destination) + { + using var scope = NewDisposeScope(); + using var _ = no_grad(); + var (p, q, h) = (_c.PatchSize, _c.NumQuantiles, layout.Horizon); + var hp = layout.HorizonPatches; + + long[] dims = [layout.TotalRows, layout.Width]; + var values = tensor(layout.Values, dims).to(device); + var observed = tensor(layout.Observed, dims).to(device); + var groups = tensor(layout.Groups).to(device); + var roles = tensor(layout.Roles).to(device); + + var raw = model.PredictBlock(values, observed, groups, roles, layout.ContextPatches, hp); + + // Only target rows are post-processed and read, so only their context is needed. + var picked = entries.SelectMany(e => Mirrored ? new long[] { e.TargetRow, e.TargetRow + layout.Rows } + : new long[] { e.TargetRow }).ToArray(); + var index = tensor(picked).to(device); + var context = layout.ContextWidth; + var ctx = values.index_select(0, index).narrow(1, 0, context); + var ctxObs = observed.index_select(0, index).narrow(1, 0, context); + var preds = Refine(raw.index_select(0, index), ctx, ctxObs, hp); + preds = Cap(preds, ctx, ctxObs); + if (preds.isnan().any().item()) + throw new InvalidOperationException("forecast produced NaN; the input is degenerate"); + + // (picked, h·p, Q) -> host; the horizon rounds up to whole patches, so trim. + var host = preds.narrow(1, 0, h).to(ScalarType.Float32).cpu().contiguous().data().ToArray(); + var nonNegative = NonNegative(ctx, ctxObs).cpu().data().ToArray(); + var stride = h * q; + for (var j = 0; j < entries.Length; j++) + { + var dst = destination.Slice(slots[j] * stride, stride); + var a = (Mirrored ? 2 * j : j) * stride; + var m = a + stride; + for (var t = 0; t < h; t++) + for (var k = 0; k < q; k++) + { + // The reference averages in float32: 0.5 * (q + flip(-mirrored)). + var v = Mirrored ? 0.5f * (host[a + t * q + k] + -host[m + t * q + (q - 1 - k)]) : host[a + t * q + k]; + if (ClipNonNegative && nonNegative[Mirrored ? 2 * j : j]) v = MathF.Max(v, 0f); + dst[t * q + k] = v; + } + } + } + + /// Whether each row's observed context exists and is all non-negative. + private static Tensor NonNegative(Tensor ctx, Tensor ctxObs) + => ctxObs.any(-1) & (ctx.ge(0) | ~ctxObs).all(-1); + + /// + /// Move a decoded block from the frame of the last context patch into the advancing + /// refine frame it was trained in: horizon patch j is denormalized with the + /// statistics of the observed context plus the block's medians for patches 0..j-1. + /// The first patch keeps the anchor frame. Welford states in float64, on the host. + /// + private Tensor Refine(Tensor quantiles, Tensor context, Tensor observed, int horizonPatches) + { + if (horizonPatches <= 1) return quantiles; + using var scope = NewDisposeScope(); + var (p, rows) = (_c.PatchSize, (int)quantiles.shape[0]); + var (loc, scale) = TimestralPreprocess.CausalPatchStats(context, observed, p, _c.MinimumScale); + var loc0 = loc.select(1, loc.shape[1] - 1); + var scale0 = scale.select(1, scale.shape[1] - 1); + + var median = quantiles.select(-1, _c.MedianIndex).to(ScalarType.Float64).cpu().data().ToArray(); + var l0 = loc0.to(ScalarType.Float64).cpu().data().ToArray(); + var s0 = scale0.to(ScalarType.Float64).cpu().data().ToArray(); + var x = context.to(ScalarType.Float64).cpu().data().ToArray(); + var w = observed.cpu().data().ToArray(); + + var width = (int)context.shape[1]; + var locs = new float[rows * horizonPatches]; + var scales = new float[rows * horizonPatches]; + var clip = Math.Sqrt(float.MaxValue) / 1e3; + for (var r = 0; r < rows; r++) + { + // The context's state, one patch at a time. + double cnt = 0, mean = 0, m2 = 0; + for (var start = 0; start < width; start += p) + { + double c = 0, s = 0; + for (var t = start; t < start + p; t++) if (w[r * width + t]) { c++; s += x[r * width + t]; } + var pm = c == 0 ? 0 : s / Math.Max(c, 1.0); + double pm2 = 0; + for (var t = start; t < start + p; t++) + if (w[r * width + t]) pm2 += (x[r * width + t] - pm) * (x[r * width + t] - pm); + (cnt, mean, m2) = Combine(cnt, mean, m2, c, pm, pm2); + } + + var offset = mean; // medians are clamped relative to the context mean, so a large level survives + for (var j = 0; j < horizonPatches; j++) + { + var sd = Math.Max(Math.Sqrt(Math.Max(m2 / Math.Max(cnt - 1.0, 1.0), 0.0)), _c.MinimumScale); + locs[r * horizonPatches + j] = j == 0 ? (float)l0[r] : (float)mean; + scales[r * horizonPatches + j] = j == 0 ? (float)s0[r] : (float)sd; + + double bs = 0; + var folded = new double[p]; + for (var t = 0; t < p; t++) + { + var sinhMedian = (median[(r * horizonPatches + j) * p + t] - l0[r]) / s0[r]; + folded[t] = Math.Clamp(sinhMedian * sd + mean - offset, -clip, clip) + offset; + bs += folded[t]; + } + var bm = bs / p; + double bm2 = 0; + foreach (var f in folded) bm2 += (f - bm) * (f - bm); + (cnt, mean, m2) = Combine(cnt, mean, m2, p, bm, bm2); + } + } + + var lj = tensor(locs, [rows, horizonPatches]).to(quantiles.device).repeat_interleave(p, 1).unsqueeze(-1); + var sj = tensor(scales, [rows, horizonPatches]).to(quantiles.device).repeat_interleave(p, 1).unsqueeze(-1); + var norm = (quantiles - loc0.view(rows, 1, 1)) / scale0.view(rows, 1, 1); + return TimestralPreprocess.ClampNonFinite(norm * sj + lj).MoveToOuterDisposeScope(); + } + + private static (double, double, double) Combine(double wa, double ma, double m2a, double wb, double mb, double m2b) + { + var w = wa + wb; + if (w == 0) return (0, 0, 0); + var delta = mb - ma; + return (w, ma + delta * wb / w, m2a + m2b + delta * delta * wa * wb / w); + } + + /// Clamp to [min − k·s, max + k·s] of the observed context, s the + /// scale of its last patch and k 10⁴. + private Tensor Cap(Tensor quantiles, Tensor context, Tensor observed) + { + using var scope = NewDisposeScope(); + var hi = nan_to_num(context.masked_fill(~observed, float.NegativeInfinity).amax([-1L], keepdim: true), 0, 0, 0); + var lo = nan_to_num(context.masked_fill(~observed, float.PositiveInfinity).amin([-1L], keepdim: true), 0, 0, 0); + var (_, scale) = TimestralPreprocess.CausalPatchStats(context, observed, _c.PatchSize, _c.MinimumScale); + var band = scale.narrow(1, scale.shape[1] - 1, 1) * QuantileCap; + return quantiles.clamp((lo - band).unsqueeze(-1), (hi + band).unsqueeze(-1)).MoveToOuterDisposeScope(); + } + + private void ValidateHorizon(int horizon) + { + if (horizon < 1 || horizon > MaxHorizon) + throw new ArgumentOutOfRangeException(nameof(horizon), $"horizon must lie in 1..{MaxHorizon}"); + } + + private int ValidateJoint(int length, int variates, int contextBars, int horizon, + ReadOnlySpan knownFuture, int targetVariate, string where) + { + if (variates < 1) + throw new ArgumentOutOfRangeException(nameof(variates), $"{where}at least one series is required"); + if (length % variates != 0) + throw new ArgumentException($"{where}series length {length} is not divisible by {variates} variates", "series"); + var total = length / variates; + if (contextBars < 1 || contextBars > total) + throw new ArgumentOutOfRangeException(nameof(contextBars), + $"{where}contextBars must lie in 1..{total}; got {contextBars}"); + if (contextBars > MaxContext) + throw new ArgumentOutOfRangeException(nameof(contextBars), + $"{where}{CheckpointName} attends over at most {MaxContext} steps; got {contextBars}. Pass the most recent {MaxContext}."); + if (targetVariate < 0 || targetVariate >= variates) + throw new ArgumentOutOfRangeException(nameof(targetVariate), $"{where}no variate {targetVariate}"); + if (!knownFuture.IsEmpty && knownFuture.Length != variates) + throw new ArgumentException($"{where}knownFuture must hold {variates} flags or be empty", nameof(knownFuture)); + if (!knownFuture.IsEmpty && knownFuture[targetVariate]) + throw new ArgumentException($"{where}the target variate cannot be known into the future", nameof(knownFuture)); + if (!knownFuture.IsEmpty && knownFuture.Contains(true)) + { + if (!_c.CovariateEnabled) + throw new ArgumentException($"{where}{CheckpointName} was not trained with covariates", nameof(knownFuture)); + if (horizon > total - contextBars) + throw new ArgumentOutOfRangeException(nameof(horizon), + $"{where}horizon {horizon} exceeds the {total - contextBars} known steps supplied after the context"); + } + return total; + } + + private void ValidateBars(int length, string where) + { + if (length == 0 || length % Channels != 0) + throw new ArgumentException($"{where}ohlcva must hold whole bars of {Channels} channels", "ohlcva"); + if (length / Channels > MaxContext) + throw new ArgumentOutOfRangeException("ohlcva", $"{where}at most {MaxContext} bars; got {length / Channels}"); + } + + private void ValidateDestination(int length, int entries, int horizon, string name) + { + if (length != BatchOutputLength(entries, horizon)) + throw new ArgumentException( + $"{name} must hold exactly {BatchOutputLength(entries, horizon)} elements " + + $"(entries x horizon x {_c.NumQuantiles} quantiles); got {length}. Use {nameof(BatchOutputLength)}.", + name); + } + + /// + /// Packed rows for one pass: every row's context right-aligned on a patch boundary, the + /// longest rounded up to whole patches and the rest left-padded unobserved. A horizon + /// region follows only when some row is a known-future covariate, carrying its values. + /// Under symmetric averaging the negated copy of every row follows the originals, each + /// copied group its own series. + /// + private sealed class Layout + { + public readonly int Rows, TotalRows, ContextPatches, HorizonPatches, Horizon, ContextWidth; + private readonly int _patch, _capacity; + private readonly bool _mirrored; + private readonly List<(long Group, float[] Context, float[]? Future)> _rows = []; + private long _group; + + public Layout(int patch, int maxContext, int horizon, int rows, bool mirrored) + { + (_patch, Horizon, Rows, _capacity, _mirrored) = (patch, horizon, rows, rows, mirrored); + ContextPatches = (maxContext + patch - 1) / patch; + HorizonPatches = (horizon + patch - 1) / patch; + ContextWidth = ContextPatches * patch; + TotalRows = mirrored ? 2 * rows : rows; + } + + public long NextGroup() => _group++; + + public int AddRow(long group, ReadOnlySpan context, ReadOnlySpan future) + { + _rows.Add((group, context.ToArray(), future.IsEmpty ? null : future.ToArray())); + return _rows.Count - 1; + } + + private bool AnyKnown => _rows.Any(r => r.Future is not null); + + /// Patches the encoder spans: the context, and the horizon too when any row + /// carries known-future values. + public int Width => (ContextPatches + (AnyKnown ? HorizonPatches : 0)) * _patch; + + public float[] Values => Build().Values; + public bool[] Observed => Build().Observed; + public long[] Groups => Build().Groups; + public long[] Roles => Build().Roles; + + private (float[] Values, bool[] Observed, long[] Groups, long[] Roles)? _built; + + private (float[] Values, bool[] Observed, long[] Groups, long[] Roles) Build() + { + if (_built is { } b) return b; + if (_rows.Count != _capacity) throw new InvalidOperationException("layout is incomplete"); + var width = Width; + var values = new float[TotalRows * width]; + var observed = new bool[TotalRows * width]; + var groups = new long[TotalRows]; + var roles = new long[TotalRows]; + for (var copy = 0; copy < (_mirrored ? 2 : 1); copy++) + for (var i = 0; i < Rows; i++) + { + var (group, context, future) = _rows[i]; + var row = copy * Rows + i; + var sign = copy == 0 ? 1f : -1f; + groups[row] = group + copy * _group; + roles[row] = future is null ? Role.Target : Role.KnownCovariate; + var o = row * width; + var pad = ContextWidth - context.Length; + for (var t = 0; t < context.Length; t++) Write(o + pad + t, context[t] * sign); + if (future is null) continue; + for (var t = 0; t < future.Length; t++) Write(o + ContextWidth + t, future[t] * sign); + } + _built = (values, observed, groups, roles); + return _built.Value; + + void Write(int at, float v) + { + if (!float.IsFinite(v)) return; // NaN is missing; the reference reads isfinite + values[at] = v; + observed[at] = true; + } + } + } +} diff --git a/src/Tsfm.Forecasting.Timestral/TimestralModel.cs b/src/Tsfm.Forecasting.Timestral/TimestralModel.cs new file mode 100644 index 0000000..51d4a27 --- /dev/null +++ b/src/Tsfm.Forecasting.Timestral/TimestralModel.cs @@ -0,0 +1,184 @@ +using System.Text.Json; + +using TorchSharp; + +using static TorchSharp.torch; + +namespace Tsfm.Forecasting.Timestral; + +/// +/// Checkpoint hyper-parameters, read from the published config.json. Strict, as the +/// reference is: an unknown key may change the forward pass, so it is refused rather than +/// ignored. +/// +public sealed record TimestralConfig( + int PatchSize, long DModel, long NumHeads, long MlpHiddenDim, string[] BlockSchedule, double[] Quantiles, + int BlockSize, int ContextLength, double RopeTheta, double RopePartialFactor, double NormEps, + bool AttnBias, bool MlpBias, bool TimeAttnGate, bool VariateAttnGate, double MinimumScale, bool UseAsinh, + bool CovariateEnabled, bool SplitContextHorizonAttn, bool AttendMissingPatches, double QuantileWidthFactor) +{ + public long HeadDim => DModel / NumHeads; + public long RopeDim => (long)(HeadDim * RopePartialFactor); + public int NumQuantiles => Quantiles.Length; + public int MedianIndex => Array.IndexOf(Quantiles, 0.5); + + /// Most steps decoded in one forward pass. + public int MaxHorizon => BlockSize * PatchSize; + + public static TimestralConfig Parse(string json) + { + using var doc = JsonDocument.Parse(json); + var r = doc.RootElement; + string[] known = + [ + "patch_size", "d_model", "num_heads", "mlp_hidden_dim", "block_schedule", "quantile_knots", "block_size", + "context_length", "rope_theta", "rope_partial_factor", "norm_eps", "attn_bias", "mlp_bias", + "time_attn_gate", "variate_attn_gate", "minimum_scale", "use_asinh", "stabilize_with_global", + "covariate_enabled", "split_context_horizon_attn", "attend_missing_patches", "quantile_width_factor", + ]; + var unknown = r.EnumerateObject().Select(p => p.Name).Except(known).ToArray(); + if (unknown.Length > 0) + throw new InvalidDataException($"unknown config keys: {string.Join(", ", unknown)}"); + + double D(string k, double d) => r.TryGetProperty(k, out var e) ? e.GetDouble() : d; + bool B(string k, bool d) => r.TryGetProperty(k, out var e) ? e.GetBoolean() : d; + bool Gate(string k) + { + var g = r.TryGetProperty(k, out var e) ? e.GetString() : "sdpa_sigmoid"; + return g switch + { + "sdpa_sigmoid" => true, + "none" => false, + _ => throw new InvalidDataException($"unknown attention gate '{g}'"), + }; + } + + // The advancing refine frame the forecaster decodes in assumes purely causal + // statistics; the global stabiliser would change every patch's frame. + if (B("stabilize_with_global", false)) + throw new InvalidDataException("stabilize_with_global is not supported by the advancing refine frame"); + + var cfg = new TimestralConfig( + r.GetProperty("patch_size").GetInt32(), + r.GetProperty("d_model").GetInt64(), + r.GetProperty("num_heads").GetInt64(), + r.GetProperty("mlp_hidden_dim").GetInt64(), + [.. r.GetProperty("block_schedule").EnumerateArray().Select(e => e.GetString()!)], + [.. r.GetProperty("quantile_knots").EnumerateArray().Select(e => e.GetDouble())], + r.GetProperty("block_size").GetInt32(), + r.GetProperty("context_length").GetInt32(), + D("rope_theta", 10000.0), D("rope_partial_factor", 0.5), D("norm_eps", 1e-6), + B("attn_bias", false), B("mlp_bias", false), Gate("time_attn_gate"), Gate("variate_attn_gate"), + D("minimum_scale", 1e-6), B("use_asinh", true), B("covariate_enabled", true), + B("split_context_horizon_attn", true), B("attend_missing_patches", true), D("quantile_width_factor", 2.0)); + + if (cfg.MedianIndex < 0) throw new InvalidDataException("quantile_knots must contain 0.5"); + if (!cfg.AttendMissingPatches) + throw new InvalidDataException("attend_missing_patches=false is not ported"); + return cfg; + } +} + +/// +/// Timestral: a causal patch transformer trained with Decoupled Horizon Forecasting. Context +/// patches are encoded causally; the horizon is read by learned forecast-query tokens that +/// attend to the encoded context, and a monotone head emits 21 quantiles per step. +/// +/// Each row is a target, a past-only covariate or a known-future covariate. With +/// known-future rows the encoder also spans the horizon, so their future values are keys; +/// targets' horizon slots stay blank. +/// +public sealed class TimestralModel : nn.Module +{ + private readonly TimestralConfig _c; + + public PatchRepresentation representation; + public Transformer transformer; + public QuantileHead head; + + public TimestralModel(TimestralConfig c) : base(nameof(TimestralModel)) + { + _c = c; + representation = new PatchRepresentation(c.PatchSize, c.DModel); + transformer = new Transformer(c); + head = new QuantileHead(c); + RegisterComponents(); + } + + public TimestralConfig Config => _c; + + /// + /// Forecast one block of after packed rows. + /// + /// (rows, W * patch): the context, then — when any row is a + /// known-future covariate — the horizon, carrying those rows' future values. + /// Observation mask, same shape. + /// (rows) series id; rows of a series attend each other. + /// (rows) . + /// Context patches, L. W is L, or L + h with known rows. + /// Patches to forecast, h. + /// Quantiles (rows, h * patch, Q) in data units, in the frame of the last + /// context patch. + public Tensor PredictBlock(Tensor values, Tensor observed, Tensor groups, Tensor roles, + int contextPatches, int horizonPatches) + { + using var scope = NewDisposeScope(); + var (r, p) = (values.shape[0], _c.PatchSize); + var w = (int)(values.shape[1] / p); + + var (scaled, loc, scale) = TimestralPreprocess.Scale(values, observed, roles, _c); + var xEnc = representation.Embed(scaled, observed); + var xDec = representation.ForecastTokens(r, horizonPatches); + + var perPatch = observed.reshape(r, w, p).any(-1); + var keyValid = TimestralPreprocess.KeyValid(perPatch); + var timeMask = TimestralPreprocess.TimeMask(keyValid, roles.eq(Role.KnownCovariate), contextPatches, horizonPatches); + var variateMask = TimestralPreprocess.VariateMask(groups, keyValid, roles, horizonPatches); + var first = TimestralPreprocess.FirstObserved(perPatch, groups); + var positions = TimestralPreprocess.Positions(first, w, contextPatches, horizonPatches); + + var h = transformer.Forward(cat([xEnc, xDec], 1), timeMask, variateMask, positions, w); + var preds = head.Forward(h.narrow(1, w, horizonPatches)); + + var anchor = (contextPatches - 1) * p; + var anchorLoc = loc.narrow(1, anchor, 1).unsqueeze(-1); + var anchorScale = scale.narrow(1, anchor, 1).unsqueeze(-1); + return TimestralPreprocess.Denormalize(preds, anchorLoc, anchorScale, _c.UseAsinh).MoveToOuterDisposeScope(); + } + + /// Bind published weights by their PyTorch names. Every parameter must be + /// matched, and every checkpoint tensor used: an unbound one keeps its random init and + /// still forecasts plausibly. + public static TimestralModel Load(IModelWeights weights, Device device) + { + var cfg = TimestralConfig.Parse(weights.ModelConfigJson); + var net = new TimestralModel(cfg); + using var stream = weights.OpenModel(); + using var st = Safetensors.Load(stream, weights.Name); + + var named = net.named_parameters().ToDictionary(p => p.name, p => p.parameter); + var missing = named.Keys.Where(k => !st.Has(k)).ToArray(); + var extra = st.Tensors.Keys.Where(k => !named.ContainsKey(k)).ToArray(); + if (missing.Length > 0 || extra.Length > 0) + throw new InvalidDataException( + $"{weights.Name}: {missing.Length} parameters absent (e.g. {string.Join(", ", missing.Take(3))}), " + + $"{extra.Length} unexpected (e.g. {string.Join(", ", extra.Take(3))})"); + + using (no_grad()) + { + foreach (var (name, param) in named) + { + using var t = st.Get(name); + if (!t.shape.SequenceEqual(param.shape)) + throw new InvalidDataException( + $"{name}: checkpoint {string.Join('x', t.shape)} vs model {string.Join('x', param.shape)}"); + param.copy_(t); + } + } + + Console.Error.WriteLine($"bound {named.Count} parameters from {weights.Name}"); + net.to(device); + net.eval(); + return net; + } +} diff --git a/src/Tsfm.Forecasting.Timestral/TimestralModules.cs b/src/Tsfm.Forecasting.Timestral/TimestralModules.cs new file mode 100644 index 0000000..1468189 --- /dev/null +++ b/src/Tsfm.Forecasting.Timestral/TimestralModules.cs @@ -0,0 +1,305 @@ +using TorchSharp.Modules; + +using static TorchSharp.torch; +using static TorchSharp.torch.nn; + +namespace Tsfm.Forecasting.Timestral; + +/// RMS norm, weight named scale. +public sealed class RmsNorm(long dims, double eps) : Module(nameof(RmsNorm)) +{ + public Parameter scale = new(ones(dims)); + + public override Tensor forward(Tensor x) + { + using var s = NewDisposeScope(); + return (x * (x.pow(2).mean([-1L], keepdim: true) + eps).rsqrt() * scale).MoveToOuterDisposeScope(); + } +} + +/// output(silu(hidden(x))) + residual(x) — SiLU, where t0's uses ReLU. +public sealed class ResidualBlock : Module +{ + public Linear hidden_layer, output_layer, residual_layer; + + public ResidualBlock(long input, long hidden, long output, bool bias) : base(nameof(ResidualBlock)) + { + hidden_layer = Linear(input, hidden, hasBias: bias); + output_layer = Linear(hidden, output, hasBias: bias); + residual_layer = Linear(input, output, hasBias: bias); + RegisterComponents(); + } + + public override Tensor forward(Tensor x) + { + using var s = NewDisposeScope(); + return (output_layer.forward(functional.silu(hidden_layer.forward(x))) + residual_layer.forward(x)) + .MoveToOuterDisposeScope(); + } +} + +/// SwiGLU with one fused input projection, computing gate * silu(up): the +/// SiLU falls on the SECOND half, the reverse of t0's. +public sealed class FusedSwiGlu : Module +{ + public Linear w12, w3; + + public FusedSwiGlu(long embed, long hidden, bool bias) : base(nameof(FusedSwiGlu)) + { + w12 = Linear(embed, 2 * hidden, hasBias: bias); + w3 = Linear(hidden, embed, hasBias: bias); + RegisterComponents(); + } + + public override Tensor forward(Tensor x) + { + using var s = NewDisposeScope(); + var halves = w12.forward(x).chunk(2, -1); + return w3.forward(halves[0] * functional.silu(halves[1])).MoveToOuterDisposeScope(); + } +} + +/// +/// Partial rotary embedding in the split-half layout: dimension i pairs with +/// i + headDim/2, and only the first rotaryDim/2 pairs turn — the rest see +/// cos 1, sin 0. Positions are integers per row, anchored at the series' first observed +/// patch, so left padding does not shift them. +/// +public sealed class PartialRope(long headDim, long rotaryDim, double theta) +{ + /// Cos and sin, each (rows, seq, headDim), for integer positions + /// (rows, seq). + public (Tensor Cos, Tensor Sin) Tables(Tensor positions) + { + using var s = NewDisposeScope(); + var idx = arange(0, rotaryDim, 2, ScalarType.Float32, positions.device) / rotaryDim; + var rotated = 1.0 / full_like(idx, theta).pow(idx); + var inv = cat([rotated, zeros(headDim / 2 - rotated.shape[0], ScalarType.Float32, positions.device)]); + var freqs = positions.to(ScalarType.Float32).unsqueeze(-1) * inv; + var emb = cat([freqs, freqs], -1); + return (emb.cos().MoveToOuterDisposeScope(), emb.sin().MoveToOuterDisposeScope()); + } + + public static Tensor Apply(Tensor x, Tensor cos, Tensor sin) + { + using var s = NewDisposeScope(); + var half = x.shape[^1] / 2; + var rotated = cat([-x.narrow(-1, half, half), x.narrow(-1, 0, half)], -1); + return (x * cos + rotated * sin).MoveToOuterDisposeScope(); + } +} + +/// +/// Gated multi-head attention over time (rows attend along patches, with RoPE) or across +/// variates (patches attend along rows, unordered). The output is multiplied by +/// sigmoid(_gate(x)) before the output projection. Unbiased throughout. +/// +public sealed class Attention : Module +{ + private readonly long _heads, _headDim, _embed; + + public bool IsTime { get; } + + public Linear wQKV, wO; + public Linear? _gate; + + public Attention(long embed, long heads, bool isTime, bool bias, bool gated) : base(nameof(Attention)) + { + (_embed, _heads, _headDim, IsTime) = (embed, heads, embed / heads, isTime); + wQKV = Linear(embed, 3 * embed, hasBias: bias); + wO = Linear(embed, embed, hasBias: bias); + if (gated) _gate = Linear(embed, embed, hasBias: bias); + RegisterComponents(); + } + + /// (rows, patches, embed) + /// Time: (rows, 1, S, S). Variate: (S, 1, rows, rows). True attends. + /// RoPE table (rows, S, headDim); time only. + /// RoPE table; time only. + /// Encoder tokens leading the time sequence. Their queries read + /// encoder keys only; the rest read everything — the quadrant they skip is masked anyway. + public Tensor Forward(Tensor x, Tensor mask, Tensor? cos, Tensor? sin, long encLen) + { + using var s = NewDisposeScope(); + var (rows, patches) = (x.shape[0], x.shape[1]); + var inputs = IsTime ? x : x.permute(1, 0, 2); + var (b, n) = (inputs.shape[0], inputs.shape[1]); + + // The fused projection is laid out (qkv, head_dim, heads): head-minor, not head-major. + var qkv = wQKV.forward(inputs.contiguous()).reshape(b, n, 3, _headDim, _heads).permute(2, 0, 1, 4, 3); + var (q, k, v) = (qkv[0].contiguous(), qkv[1].contiguous(), qkv[2].contiguous()); + + Tensor o; + if (IsTime) + { + var (c, sn) = (cos!.unsqueeze(2), sin!.unsqueeze(2)); + (q, k) = (PartialRope.Apply(q, c, sn), PartialRope.Apply(k, c, sn)); + var dec = n - encLen; + var enc = Sdpa(q.narrow(1, 0, encLen), k.narrow(1, 0, encLen), v.narrow(1, 0, encLen), + mask.narrow(2, 0, encLen).narrow(3, 0, encLen)); + o = dec == 0 ? enc : cat([enc, Sdpa(q.narrow(1, encLen, dec), k, v, mask.narrow(2, encLen, dec))], 1); + } + else + { + o = Sdpa(q, k, v, mask); + } + + if (_gate is not null) o = o * _gate.forward(inputs).sigmoid(); + o = IsTime ? o.reshape(rows, patches, _embed) : o.reshape(patches, rows, _embed).permute(1, 0, 2); + return wO.forward(o).MoveToOuterDisposeScope(); + } + + /// SDPA on (batch, seq, heads, headDim); returns (batch, seq, heads * headDim). + private static Tensor Sdpa(Tensor q, Tensor k, Tensor v, Tensor mask) + { + var (b, n) = (q.shape[0], q.shape[1]); + var o = functional.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), + attn_mask: mask); + return o.transpose(1, 2).contiguous().view(b, n, -1); + } +} + +/// One layer: its attentions in schedule order, each a pre-norm residual, then +/// one pre-norm SwiGLU. +public sealed class TransformerBlock : Module +{ + public ModuleList attn_norms; + public ModuleList attns; + public RmsNorm mlp_norm; + public FusedSwiGlu mlp; + + public TransformerBlock(string kind, TimestralConfig c) : base(nameof(TransformerBlock)) + { + bool[] axes = kind switch + { + "time" => [true], + "variate" => [false], + "time_variate" => [true, false], + "variate_time" => [false, true], + _ => throw new InvalidDataException($"unknown layer kind '{kind}'"), + }; + attn_norms = new ModuleList([.. axes.Select(_ => new RmsNorm(c.DModel, c.NormEps))]); + attns = new ModuleList([.. axes.Select(t => new Attention(c.DModel, c.NumHeads, t, c.AttnBias, + t ? c.TimeAttnGate : c.VariateAttnGate))]); + mlp_norm = new RmsNorm(c.DModel, c.NormEps); + mlp = new FusedSwiGlu(c.DModel, c.MlpHiddenDim, c.MlpBias); + RegisterComponents(); + } + + public Tensor Forward(Tensor x, Tensor timeMask, Tensor variateMask, Tensor cos, Tensor sin, long encLen) + { + using var s = NewDisposeScope(); + for (var i = 0; i < attns.Count; i++) + { + var a = attns[i]; + x = x + a.Forward(attn_norms[i].forward(x), a.IsTime ? timeMask : variateMask, cos, sin, encLen); + } + return (x + mlp.forward(mlp_norm.forward(x))).MoveToOuterDisposeScope(); + } +} + +/// The layer stack and its final norm. +public sealed class Transformer : Module +{ + private readonly PartialRope _rope; + + public ModuleList layers; + public RmsNorm out_norm; + + public Transformer(TimestralConfig c) : base(nameof(Transformer)) + { + _rope = new PartialRope(c.HeadDim, c.RopeDim, c.RopeTheta); + layers = new ModuleList([.. c.BlockSchedule.Select(k => new TransformerBlock(k, c))]); + out_norm = new RmsNorm(c.DModel, c.NormEps); + RegisterComponents(); + } + + /// [Enc ; Dec] tokens (rows, S, embed). + /// (rows, 1, S, S). + /// (S, 1, rows, rows). + /// Integer time positions (rows, S). + /// Encoder tokens leading the sequence. + public Tensor Forward(Tensor x, Tensor timeMask, Tensor variateMask, Tensor positions, long encLen) + { + using var s = NewDisposeScope(); + var (cos, sin) = _rope.Tables(positions); + foreach (var layer in layers) x = layer.Forward(x, timeMask, variateMask, cos, sin, encLen); + return out_norm.forward(x).MoveToOuterDisposeScope(); + } +} + +/// Patch embedding from [values ‖ missing], and the learned forecast-query +/// token every horizon patch starts from. +public sealed class PatchRepresentation : Module +{ + private readonly int _patch; + + public ResidualBlock patch_proj; + public Parameter forecast_query; + + public PatchRepresentation(int patch, long dModel) : base(nameof(PatchRepresentation)) + { + _patch = patch; + patch_proj = new ResidualBlock(2 * patch, 4 * dModel, dModel, bias: true); + forecast_query = new Parameter(zeros(dModel)); + RegisterComponents(); + } + + /// (rows, T) scaled values and observation mask to (rows, T / patch, dModel). + public Tensor Embed(Tensor scaled, Tensor observed) + { + using var s = NewDisposeScope(); + var (r, n) = (scaled.shape[0], scaled.shape[1] / _patch); + var missing = (~observed).to(scaled.dtype).reshape(r, n, _patch); + return patch_proj.forward(cat([scaled.reshape(r, n, _patch), missing], -1)).MoveToOuterDisposeScope(); + } + + public Tensor ForecastTokens(long rows, long patches) => forecast_query.view(1, 1, -1).expand(rows, patches, -1); +} + +/// Decoder tokens to patch monotone quantile vectors each, in scaled space. +public sealed class QuantileHead : Module +{ + private readonly TimestralConfig _c; + + public ResidualBlock projection; + + public QuantileHead(TimestralConfig c) : base(nameof(QuantileHead)) + { + _c = c; + projection = new ResidualBlock(c.DModel, 4 * c.DModel, c.PatchSize * c.NumQuantiles, bias: false); + RegisterComponents(); + } + + /// (rows, N, dModel) to (rows, N * patch, Q). + public Tensor Forward(Tensor x) + { + using var s = NewDisposeScope(); + var raw = projection.forward(x); + var (r, n) = (raw.shape[0], raw.shape[1]); + raw = raw.reshape(r, n, _c.PatchSize, _c.NumQuantiles); + return Monotone(raw, _c.Quantiles, _c.QuantileWidthFactor) + .reshape(r, n * _c.PatchSize, _c.NumQuantiles).MoveToOuterDisposeScope(); + } + + /// + /// Non-crossing quantiles: the median is free, and the rest fan out from it by + /// cumsum(width * dp * softplus(raw)), dp the probability gap to the next + /// level inward, so unequal knot spacing is respected. + /// + public static Tensor Monotone(Tensor raw, double[] levels, double width) + { + using var s = NewDisposeScope(); + var m = Array.IndexOf(levels, 0.5); + var gaps = new float[levels.Length]; + for (var i = m + 1; i < levels.Length; i++) gaps[i] = (float)(levels[i] - levels[i - 1]); + for (var i = 0; i < m; i++) gaps[i] = (float)(levels[i + 1] - levels[i]); + + var inc = width * tensor(gaps).to(raw.device) * functional.softplus(raw); + var median = raw.narrow(-1, m, 1); + var up = inc.narrow(-1, m + 1, levels.Length - m - 1).cumsum(-1); + var down = inc.narrow(-1, 0, m).flip(-1).cumsum(-1).flip(-1); + var big = float.MaxValue; + return cat([median - down, median, median + up], -1).clamp(-big, big).MoveToOuterDisposeScope(); + } +} diff --git a/src/Tsfm.Forecasting.Timestral/TimestralPreprocess.cs b/src/Tsfm.Forecasting.Timestral/TimestralPreprocess.cs new file mode 100644 index 0000000..ebe9657 --- /dev/null +++ b/src/Tsfm.Forecasting.Timestral/TimestralPreprocess.cs @@ -0,0 +1,195 @@ +using static TorchSharp.torch; + +namespace Tsfm.Forecasting.Timestral; + +/// Role of a row. Values match the reference's role ids. +public static class Role +{ + public const long Target = 0, PastCovariate = 1, KnownCovariate = 2; +} + +/// +/// Scaling, attention masks and positions around the transformer. Every per-step array is +/// (rows, T), with T a whole number of patches; an observation mask is true +/// where a value was observed. +/// +public static class TimestralPreprocess +{ + /// + /// Causal per-patch location and scale, broadcast over each patch's steps: each patch + /// takes the running mean and deviation of every observed value up to its last step. + /// Computed in float64 on the CPU, as the reference does — MPS has no float64. + /// + public static (Tensor Loc, Tensor Scale) CausalPatchStats(Tensor values, Tensor observed, int patch, double minimumScale) + { + using var scope = NewDisposeScope(); + var device = values.device; + var x = values.to(ScalarType.Float64, CPU); + var w = observed.to(ScalarType.Float64, CPU); + + var den = w.cumsum(-1).clamp_min(1.0); + var means = (w * x).cumsum(-1) / den; + var shifted = cat([zeros_like(means.narrow(1, 0, 1)), means.narrow(1, 0, means.shape[1] - 1)], -1); + var m2 = ((x - shifted) * (x - means) * w).cumsum(-1); + var scales = (m2 / (den - 1.0).clamp_min(1.0)).sqrt().clamp_min(minimumScale); + + var (r, n) = (values.shape[0], values.shape[1] / patch); + Tensor PerPatch(Tensor t) => t.to(ScalarType.Float32).reshape(r, n, patch).select(-1, patch - 1) + .repeat_interleave(patch, -1).to(device); + return (PerPatch(means).MoveToOuterDisposeScope(), PerPatch(scales).MoveToOuterDisposeScope()); + } + + /// float32 machine epsilon, the reference's floor on a divisor. + private const double Float32Eps = 1.1920928955078125e-07; + + /// asinh((x - loc) / scale) where observed, 0 elsewhere. + public static Tensor Affine(Tensor x, Tensor loc, Tensor scale, Tensor observed, bool asinh) + { + using var s = NewDisposeScope(); + var scaled = ((x - loc) / scale.clamp_min(Float32Eps)).masked_fill(~observed, 0.0); + if (asinh) scaled = where(observed, scaled.asinh(), scaled); + return scaled.MoveToOuterDisposeScope(); + } + + /// + /// Non-causal per-row location and scale for known-future rows, over context and + /// horizon alike. A row observed only at 0 and 1 is an indicator: it keeps loc 0, + /// scale 1, and skips the asinh. + /// + public static (Tensor Loc, Tensor Scale, Tensor IsBinary) InstanceStats(Tensor x, Tensor observed, double minimumScale) + { + using var s = NewDisposeScope(); + var m = observed.to(x.dtype); + var n = m.sum(-1, keepdim: true).clamp_min(1.0); + var loc = (x * m).sum(-1, keepdim: true) / n; + var variance = ((x - loc) * m).pow(2).sum(-1, keepdim: true) / n; + var scale = variance.sqrt().clamp_min(minimumScale); + var zeroOne = (x.eq(0.0) | x.eq(1.0) | ~observed).all(-1, keepdim: true); + var binary = zeroOne & observed.any(-1, keepdim: true); + loc = where(binary, zeros_like(loc), loc); + scale = where(binary, ones_like(scale), scale); + return (loc.MoveToOuterDisposeScope(), scale.MoveToOuterDisposeScope(), binary.MoveToOuterDisposeScope()); + } + + /// + /// Scaled input for every row: causal for targets and past covariates, instance + /// statistics for known-future rows. + /// + public static (Tensor Scaled, Tensor Loc, Tensor Scale) Scale( + Tensor x, Tensor observed, Tensor roles, TimestralConfig c) + { + using var s = NewDisposeScope(); + var (loc, scale) = CausalPatchStats(x, observed, c.PatchSize, c.MinimumScale); + var scaled = Affine(x, loc, scale, observed, c.UseAsinh); + + var known = roles.eq(Role.KnownCovariate); + if (known.any().item()) + { + var (kl, ks, binary) = InstanceStats(x, observed, c.MinimumScale); + var ksc = Affine(x, kl, ks, observed, asinh: false); + if (c.UseAsinh) ksc = where(observed & ~binary, ksc.asinh(), ksc); + var sel = known.unsqueeze(1); + scaled = where(sel, ksc, scaled); + loc = where(sel, kl.expand_as(loc), loc); + scale = where(sel, ks.expand_as(scale), scale); + } + return (scaled.MoveToOuterDisposeScope(), loc.MoveToOuterDisposeScope(), scale.MoveToOuterDisposeScope()); + } + + /// Encoder key validity (rows, W): a patch from the series' first + /// observed one on, observed or not. Leading unobserved patches are padding. + public static Tensor KeyValid(Tensor observedPerPatch) + => observedPerPatch.to(ScalarType.Int32).cumsum(-1).gt(0); + + /// First patch any row of each row's group observes, per row; 0 if none. + public static Tensor FirstObserved(Tensor observedPerPatch, Tensor groups) + { + using var s = NewDisposeScope(); + var w = observedPerPatch.shape[1]; + var groupCount = groups.max().item() + 1; + var count = zeros([groupCount, w], ScalarType.Float32, groups.device) + .index_add_(0, groups, observedPerPatch.to(ScalarType.Float32), 1.0); + return count.gt(0).to(ScalarType.Int32).argmax(-1).index_select(0, groups).MoveToOuterDisposeScope(); + } + + /// + /// Time mask over [Enc(W) ; Dec(h)], (rows, 1, W+h, W+h), true attends. + /// + /// Enc to Enc: causal over valid keys, the diagonal always open; full for + /// known-future rows. + /// Enc to Dec: never. + /// Dec to Enc: every valid context key, j < contextPatches — never the + /// horizon's encoder tokens, which only known-future rows populate. + /// Dec to Dec: causal. + /// + /// + public static Tensor TimeMask(Tensor keyValid, Tensor knownRows, int contextPatches, int horizonPatches) + { + using var s = NewDisposeScope(); + var (r, w) = (keyValid.shape[0], keyValid.shape[1]); + var h = horizonPatches; + var dev = keyValid.device; + var diagonal = TorchSharp.torch.eye(w, dtype: ScalarType.Bool, device: dev).unsqueeze(0); + var keys = keyValid.unsqueeze(1); + var causal = ones([w, w], ScalarType.Bool, dev).tril().unsqueeze(0); + var encEnc = where(knownRows.view(r, 1, 1), keys.expand(r, w, w), causal & keys) | diagonal; + var encDec = zeros([r, w, h], ScalarType.Bool, dev); + var inContext = arange(w, ScalarType.Int64, dev).lt(contextPatches).view(1, 1, w); + var decEnc = (inContext & keys).expand(r, h, w); + var decDec = ones([h, h], ScalarType.Bool, dev).tril().unsqueeze(0).expand(r, h, h); + var top = cat([encEnc, encDec], 2); + var bottom = cat([decEnc, decDec], 2); + return cat([top, bottom], 1).unsqueeze(1).MoveToOuterDisposeScope(); + } + + /// + /// Variate mask (W+h, 1, rows, rows): a row reads rows of its own group, on + /// encoder slots only valid keys. Two firewalls: a covariate never reads a target, and a + /// known-future row reads only known-future rows — else a covariate's view of the + /// future could relay a target's own future back into it. The diagonal is always open. + /// + public static Tensor VariateMask(Tensor groups, Tensor keyValid, Tensor roles, int horizonPatches) + { + using var s = NewDisposeScope(); + var r = keyValid.shape[0]; + var dev = keyValid.device; + var sameGroup = groups.unsqueeze(1).eq(groups.unsqueeze(0)); + var slotValid = cat([keyValid, ones([r, horizonPatches], ScalarType.Bool, dev)], 1).t(); // (S, R) + var target = roles.eq(Role.Target); + var known = roles.eq(Role.KnownCovariate); + var forbid = (~target.unsqueeze(1) & target.unsqueeze(0)) | (known.unsqueeze(1) & ~known.unsqueeze(0)); + var mask = sameGroup.unsqueeze(0) & slotValid.unsqueeze(1) & ~forbid.unsqueeze(0); + var diagonal = TorchSharp.torch.eye(r, dtype: ScalarType.Bool, device: dev).unsqueeze(0); + return (mask | diagonal).unsqueeze(1).MoveToOuterDisposeScope(); + } + + /// Positions (rows, W+h): encoder patch p and the decoder token for + /// horizon patch contextPatches + i sit at their patch index less the group's first + /// observed patch. + public static Tensor Positions(Tensor firstObserved, int encPatches, int contextPatches, int horizonPatches) + { + var dev = firstObserved.device; + var p = cat([arange(encPatches, ScalarType.Int64, dev), arange(contextPatches, contextPatches + horizonPatches, + ScalarType.Int64, dev)]); + return p.unsqueeze(0) - firstObserved.unsqueeze(1); + } + + /// sinh(x) * scale + loc, then ±inf (a sinh overflow) to the nearest + /// finite value along the quantile axis, or 0 when the whole axis is infinite. + public static Tensor Denormalize(Tensor preds, Tensor loc, Tensor scale, bool asinh) + { + using var s = NewDisposeScope(); + var x = asinh ? preds.sinh() : preds; + return ClampNonFinite(x * scale + loc).MoveToOuterDisposeScope(); + } + + public static Tensor ClampNonFinite(Tensor v) + { + using var s = NewDisposeScope(); + var inf = float.PositiveInfinity; + var maxFinite = where(v.isfinite(), v, full_like(v, -inf)).amax([-1L], keepdim: true); + var minFinite = where(v.isfinite(), v, full_like(v, inf)).amin([-1L], keepdim: true); + var c = where(v.eq(inf), maxFinite.expand_as(v), where(v.eq(-inf), minFinite.expand_as(v), v)); + return where(c.isinf(), zeros_like(c), c).MoveToOuterDisposeScope(); + } +} diff --git a/src/Tsfm.Forecasting.Timestral/Tsfm.Forecasting.Timestral.csproj b/src/Tsfm.Forecasting.Timestral/Tsfm.Forecasting.Timestral.csproj new file mode 100644 index 0000000..294875f --- /dev/null +++ b/src/Tsfm.Forecasting.Timestral/Tsfm.Forecasting.Timestral.csproj @@ -0,0 +1,24 @@ + + + + true + Tsfm.Forecasting.Timestral + Tsfm.Forecasting.Timestral + Tsfm.Forecasting.Timestral + + Apache-2.0 + Native .NET inference for Mistral AI's Timestral family (5M, 28M, 123M, 209M): a patch transformer with time and variate attention, past and known-future covariates and 21 quantiles. Code only. Weights are under Mistral's Modified MIT licence, which withholds all rights from companies above USD 20M monthly revenue: 5M and 28M ship in Tsfm.Forecasting.Timestral.Weights.*, the larger two are fetched. + timeseries;forecasting;foundation-model;timestral;mistral;covariates;quantile;torchsharp + + + + + + + + + + + diff --git a/src/Tsfm.Forecasting/IModelWeights.cs b/src/Tsfm.Forecasting/IModelWeights.cs new file mode 100644 index 0000000..01fde46 --- /dev/null +++ b/src/Tsfm.Forecasting/IModelWeights.cs @@ -0,0 +1,78 @@ +namespace Tsfm.Forecasting; + +/// +/// A single-network checkpoint: config.json beside model.safetensors. The +/// tokenizer-free counterpart of , for models that read values +/// directly. Streams, not paths: an embedded resource is located by the runtime rather +/// than by configuration. +/// +public interface IModelWeights +{ + /// Repository name and revision, for provenance. + string Name { get; } + + /// The safetensors payload. + Stream OpenModel(); + + /// The published config.json, verbatim. + string ModelConfigJson { get; } +} + +/// Weights from a snapshot directory holding config.json and +/// model.safetensors, as published. +public sealed class DirectoryWeights : IModelWeights +{ + private readonly string _dir; + + /// The snapshot directory. + /// Provenance label; defaults to the directory name. + /// Appended to the error when a file is missing, so the + /// message says where the weights come from. + public DirectoryWeights(string dir, string? name = null, string? howToFetch = null) + { + _dir = dir; + Name = name ?? Path.GetFileName(dir.TrimEnd(Path.DirectorySeparatorChar)); + foreach (var file in (string[])["config.json", "model.safetensors"]) + if (!File.Exists(Path.Combine(dir, file))) + throw new FileNotFoundException( + $"{Path.Combine(dir, file)} not found. {howToFetch}".TrimEnd(), Path.Combine(dir, file)); + } + + public string Name { get; } + + public override string ToString() => Name; + + public Stream OpenModel() => File.OpenRead(Path.Combine(_dir, "model.safetensors")); + + public string ModelConfigJson => File.ReadAllText(Path.Combine(_dir, "config.json")); +} + +/// Weights carried as embedded resources. Derived types supply the assembly and +/// resource names. +public abstract class EmbeddedWeights : IModelWeights +{ + protected abstract System.Reflection.Assembly Host { get; } + protected abstract string ModelResource { get; } + protected abstract string ModelConfigResource { get; } + + public abstract string Name { get; } + + public override string ToString() => Name; + + public Stream OpenModel() => Open(ModelResource); + + public string ModelConfigJson + { + get + { + using var reader = new StreamReader(Open(ModelConfigResource)); + return reader.ReadToEnd(); + } + } + + private Stream Open(string resource) + => Host.GetManifestResourceStream(resource) + ?? throw new InvalidOperationException( + $"resource '{resource}' missing from {Host.GetName().Name}; " + + "the weights assembly was built without its payload"); +} diff --git a/tests/Tsfm.Forecasting.Timestral.IntegrationTests/EmbeddedWeightsTests.cs b/tests/Tsfm.Forecasting.Timestral.IntegrationTests/EmbeddedWeightsTests.cs new file mode 100644 index 0000000..5c3d156 --- /dev/null +++ b/tests/Tsfm.Forecasting.Timestral.IntegrationTests/EmbeddedWeightsTests.cs @@ -0,0 +1,77 @@ +using Tsfm.Forecasting.Timestral.Weights; + +using Xunit; +using Xunit.Abstractions; + +using static TorchSharp.torch; + +namespace Tsfm.Forecasting.Timestral.IntegrationTests; + +/// +/// The embedded 5M and 28M checkpoints, end to end. Numerical agreement with the reference +/// is the parity harness's job; these check that the packages carry what they claim and +/// that the forecaster behaves on real weights. +/// +public class EmbeddedWeightsTests(ITestOutputHelper output) +{ + public static TheoryData Checkpoints => [Timestral5M.Instance, Timestral28M.Instance]; + + private const int Context = 96, Horizon = 24, Total = Context + Horizon; + + private static float[] Series(float driverFuture) + { + var s = new float[3 * Total]; + for (var t = 0; t < Total; t++) + { + s[Total + t] = t < Context ? 20f + 4f * MathF.Sin(t / 6f) : 20f + driverFuture; + s[2 * Total + t] = t % 7 < 5 ? 1f : 0f; + } + for (var t = 0; t < Context; t++) s[t] = 100f + 3f * s[Total + Math.Max(t - 1, 0)] + s[2 * Total + t]; + return s; + } + + [Theory] + [MemberData(nameof(Checkpoints))] + public void LoadsWithProvenanceAndTheTrainedGrid(IModelWeights weights) + { + var f = TimestralForecaster.Load(weights, new Device("cpu")); + Assert.Contains("@", f.CheckpointName); + Assert.Equal(21, f.Quantiles.Count); + Assert.Equal(1024, f.MaxHorizon); + Assert.Equal(32768, f.MaxContext); + } + + [Fact] + public void TheLicenceTravelsWithTheWeights() + { + Assert.Contains("Modified MIT License", Timestral5M.License); + Assert.Contains("$20 million", Timestral28M.License); + } + + [Theory] + [MemberData(nameof(Checkpoints))] + public void KnownFutureValuesLiftTheForecast(IModelWeights weights) + { + var f = TimestralForecaster.Load(weights, new Device("cpu")); + var a = f.ForecastJoint(Series(+8f), 3, Context, Horizon, [false, true, true]); + var b = f.ForecastJoint(Series(-8f), 3, Context, Horizon, [false, true, true]); + var mid = f.Config.MedianIndex; + var lift = Enumerable.Range(0, Horizon).Average(h => a[h, mid] - b[h, mid]); + output.WriteLine($"{weights.Name}: mean median lift {lift:F2}"); + Assert.True(lift > 0, "a target that tracks its driver should rise with the driver's known future"); + } + + [Theory] + [MemberData(nameof(Checkpoints))] + public void ForecastsAreOrderedAndFinite(IModelWeights weights) + { + var f = TimestralForecaster.Load(weights, new Device("cpu")); + var q = f.ForecastJoint(Series(0f), 3, Context, Horizon, [false, true, true]); + for (var h = 0; h < Horizon; h++) + for (var k = 0; k < q.GetLength(1); k++) + { + Assert.True(double.IsFinite(q[h, k])); + if (k > 0) Assert.True(q[h, k] >= q[h, k - 1]); + } + } +} diff --git a/tests/Tsfm.Forecasting.Timestral.IntegrationTests/Tsfm.Forecasting.Timestral.IntegrationTests.csproj b/tests/Tsfm.Forecasting.Timestral.IntegrationTests/Tsfm.Forecasting.Timestral.IntegrationTests.csproj new file mode 100644 index 0000000..945c935 --- /dev/null +++ b/tests/Tsfm.Forecasting.Timestral.IntegrationTests/Tsfm.Forecasting.Timestral.IntegrationTests.csproj @@ -0,0 +1,20 @@ + + + + false + Tsfm.Forecasting.Timestral.IntegrationTests + + + + + + + + + + + + + + + diff --git a/tests/Tsfm.Forecasting.Timestral.Parity/Program.cs b/tests/Tsfm.Forecasting.Timestral.Parity/Program.cs new file mode 100644 index 0000000..b71affc --- /dev/null +++ b/tests/Tsfm.Forecasting.Timestral.Parity/Program.cs @@ -0,0 +1,183 @@ +// Bisecting parity harness for Timestral. Compares the .NET port against tensors dumped +// from mistralai/timestral by reference/dump_timestral_parity.py, one stage at a time, then +// end to end — with the reference forecaster's defaults and without them. +// +// Each stage is fed the REFERENCE output of the stage before it, so a failure localises to +// one stage rather than being inherited from an earlier one. +// +// dotnet run --project tests/Tsfm.Forecasting.Timestral.Parity -c Release -- Timestral-5M Timestral-28M + +using TorchSharp; + +using Tsfm.Forecasting; +using Tsfm.Forecasting.Timestral; + +using static TorchSharp.torch; + +var root = Path.GetFullPath(Path.Combine(AppContext.BaseDirectory, "../../../../..")); +var names = args.Length > 0 ? args : ["Timestral-5M"]; + +Device device; +try { device = new Device(DeviceType.MPS); using (ones(1).to(device)) { } } +catch { device = new Device(DeviceType.CPU); } +Console.WriteLine($"device={device.type}"); + +var ok = true; +foreach (var name in names) ok &= Run(name); +Console.WriteLine(); +Console.WriteLine(ok ? "PARITY: PASS" : "PARITY: FAIL — the first failing stage above localises the fault"); +return ok ? 0 : 1; + +bool Run(string name) +{ + var fixturePath = Path.Combine(root, "reference", $"{name}-parity.safetensors"); + if (!File.Exists(fixturePath)) + { + Console.Error.WriteLine($"missing {fixturePath}; run reference/dump_timestral_parity.py checkpoints/{name}"); + return false; + } + + using var _ = no_grad(); + using var fx = Safetensors.Load(fixturePath); + Tensor E(string n) => fx.Get(n).to(device); + + var forecaster = TimestralForecaster.Load(Path.Combine(root, "checkpoints", name), device); + var (model, c) = (forecaster.Model, forecaster.Config); + var pass = true; + + Console.WriteLine($"\n── {name}: {c.BlockSchedule.Length} layers x {c.DModel}, {c.NumHeads} heads, rope {c.RopeDim}/{c.HeadDim}"); + Console.WriteLine($"{"stage",-24} {"max|Δ|",12} {"rel",12} note"); + + void Check(string stage, Tensor got, Tensor want, double tol, string note = "") + { + var absMax = (got.to(ScalarType.Float32).to(want.device) - want).abs().max().item(); + var scale = want.abs().max().item(); + var rel = absMax / Math.Max(scale, 1e-12); + var good = rel <= tol || absMax <= 1e-6; + pass &= good; + Console.WriteLine($"{stage,-24} {absMax,12:E3} {rel,12:E3} {(good ? "ok" : "FAIL")} {note}"); + } + + // 1. Scaling: causal for targets, instance statistics for known-future rows. + var values = E("buf_values"); + var observed = E("buf_observed").to(ScalarType.Bool); + var groups = E("buf_groups").to(ScalarType.Int64); + var roles = E("buf_roles").to(ScalarType.Int64); + var ll = (int)fx.Get("ctx_patches").item(); + var p = c.PatchSize; + var (r, w) = (values.shape[0], (int)(values.shape[1] / p)); + var h = w - ll; + var (scaled, loc, scale) = TimestralPreprocess.Scale(values, observed, roles, c); + Check("loc", loc, E("loc"), 1e-6); + Check("scale", scale, E("scale"), 1e-6); + Check("scaled", scaled, E("scaled"), 1e-5); + + // 2. Embedding, masks and positions. + Check("patch_embedding", model.representation.Embed(E("scaled"), observed), E("patch_embedding"), 1e-5); + var perPatch = observed.reshape(r, w, p).any(-1); + var keyValid = TimestralPreprocess.KeyValid(perPatch); + var timeMask = TimestralPreprocess.TimeMask(keyValid, roles.eq(Role.KnownCovariate), ll, h); + var variateMask = TimestralPreprocess.VariateMask(groups, keyValid, roles, h); + var positions = TimestralPreprocess.Positions(TimestralPreprocess.FirstObserved(perPatch, groups), w, ll, h); + Check("time_mask", timeMask.to(ScalarType.Float32), E("time_mask"), 0); + Check("variate_mask", variateMask.to(ScalarType.Float32), E("variate_mask"), 0); + Check("positions", positions.to(ScalarType.Float32), E("positions"), 0); + + // 3. Layers, each from the reference's previous output. + var tf = model.transformer; + var x0 = cat([E("patch_embedding"), model.representation.ForecastTokens(r, h)], 1); + var x = x0; + var rope = new PartialRope(c.HeadDim, c.RopeDim, c.RopeTheta).Tables(E("positions").to(ScalarType.Int64)); + for (var i = 0; i < tf.layers.Count; i++) + { + x = tf.layers[i].Forward(x, timeMask, variateMask, rope.Cos, rope.Sin, w); + var refName = $"layer{i}_output"; + if (!fx.Has(refName)) continue; + Check(refName, x, E(refName), 1e-4, c.BlockSchedule[i]); + x = E(refName); + } + Check("transformer_output", tf.Forward(x0, timeMask, variateMask, positions, w), E("transformer_output"), 2e-4, + "(whole stack)"); + + // 4. Head, then the whole block end to end. + Check("head_output", model.head.Forward(E("transformer_output").narrow(1, w, h)), E("head_output"), 1e-5); + Check("block_output", model.PredictBlock(values, observed, groups, roles, ll, h), E("block_output"), 1e-4); + + // 5. End to end through the public API. + var ctx = fx.Get("in_ctx"); // (2, T) + var fut = fx.Get("in_fut"); // (2, T + H) + var (t, horizon) = ((int)ctx.shape[1], (int)(fut.shape[1] - ctx.shape[1])); + var total = t + horizon; + var joint = new float[4 * total]; + for (var v = 0; v < 2; v++) + for (var i = 0; i < t; i++) joint[v * total + i] = ctx[v, i].item(); + for (var v = 0; v < 2; v++) + for (var i = 0; i < total; i++) joint[(2 + v) * total + i] = fut[v, i].item(); + + for (var target = 0; target < 2; target++) + { + Check($"joint[{target}]", Grid(forecaster.ForecastJoint(joint, 4, t, horizon, [false, false, true, true], target)), + E("predict_joint")[target], 1e-4, "(defaults)"); + forecaster.SymmetricAveraging = forecaster.ClipNonNegative = false; + Check($"joint_raw[{target}]", Grid(forecaster.ForecastJoint(joint, 4, t, horizon, [false, false, true, true], target)), + E("predict_joint_raw")[target], 1e-4, "(no averaging, no clip)"); + forecaster.SymmetricAveraging = forecaster.ClipNonNegative = true; + Check($"targets_only[{target}]", Grid(forecaster.ForecastJoint(joint.AsSpan(0, 2 * total), 2, t, horizon, + targetVariate: target)), E("predict_targets_only")[target], 1e-4); + } + var aligned = Floats(ctx[0].narrow(0, 64, 64)); + Check("aligned", Grid(forecaster.ForecastJoint(aligned, 1, 64, 32)), E("predict_aligned"), 1e-4, "(no padding)"); + + // 6. The quickstart notebooks' series, shared with the t0 harness. + var m4 = Floats(fx.Get("m4_ctx")); + Check("notebook m4 H1", Grid(forecaster.ForecastJoint(m4, 1, m4.Length, 48)), E("m4_predict"), 1e-4); + var ep = Floats(fx.Get("ep_ctx")); + var cov = fx.Get("ep_cov"); + var truth = Floats(fx.Get("ep_truth")); + var plain = forecaster.ForecastJoint(ep, 1, ep.Length, truth.Length); + Check("notebook ep", Grid(plain), E("ep_predict"), 1e-4); + var width = (int)cov.shape[1]; + var epJoint = new float[3 * width]; + ep.CopyTo(epJoint, 0); + Floats(cov[0]).CopyTo(epJoint, width); + Floats(cov[1]).CopyTo(epJoint, 2 * width); + var withCov = forecaster.ForecastJoint(epJoint, 3, ep.Length, truth.Length, [false, true, true]); + Check("notebook ep + cov", Grid(withCov), E("ep_predict_cov"), 1e-4); + + // 7. One batch of mixed lengths and kinds — including a known-future entry, which + // widens the whole pass's encoder over the horizon — against the single forecasts. + // All share the electricity horizon of 24; the refs were taken at 40, 32, 48 and 24, and + // only target-only entries truncate exactly (their later patches cannot reach earlier). + const int shared = 24; + JointSeries[] batch = + [ + new(joint.AsMemory(0, 2 * total), 2, t, TargetVariate: 1), + new(aligned, 1, 64), + new(m4, 1, m4.Length), + new(epJoint, 3, ep.Length, new[] { false, true, true }), + ]; + var outBatch = new double[forecaster.BatchOutputLength(batch.Length, shared)]; + forecaster.ForecastJointBatch(batch, shared, outBatch); + for (var i = 0; i < batch.Length; i++) + { + var e = batch[i]; + var single = Grid(forecaster.ForecastJoint(e.Series.Span, e.Variates, e.ContextBars, shared, + e.KnownFuture.Span, e.TargetVariate)); + var got = tensor(outBatch.AsSpan(i * shared * c.NumQuantiles, shared * c.NumQuantiles).ToArray() + .Select(d => (float)d).ToArray(), [shared, c.NumQuantiles]); + Check($"batch[{i}]", got, single.to(device), 1e-4, "(vs single call)"); + } + + var mid = c.MedianIndex; + Console.WriteLine($" electricity price |median-truth|: {MeanAbs(plain, truth, mid):F4} without covariates, " + + $"{MeanAbs(withCov, truth, mid):F4} with"); + return pass; +} + +static float[] Floats(Tensor x) => x.contiguous().data().ToArray(); + +static Tensor Grid(double[,] q) => tensor(q.Cast().Select(d => (float)d).ToArray(), + [q.GetLength(0), q.GetLength(1)]); + +static double MeanAbs(double[,] q, float[] truth, int col) + => Enumerable.Range(0, truth.Length).Average(h => Math.Abs(q[h, col] - truth[h])); diff --git a/tests/Tsfm.Forecasting.Timestral.Parity/Tsfm.Forecasting.Timestral.Parity.csproj b/tests/Tsfm.Forecasting.Timestral.Parity/Tsfm.Forecasting.Timestral.Parity.csproj new file mode 100644 index 0000000..7b5db49 --- /dev/null +++ b/tests/Tsfm.Forecasting.Timestral.Parity/Tsfm.Forecasting.Timestral.Parity.csproj @@ -0,0 +1,16 @@ + + + + Exe + Tsfm.Forecasting.Timestral.Parity + + + + + + + + + + + diff --git a/tests/Tsfm.Forecasting.Timestral.Tests/BatchTests.cs b/tests/Tsfm.Forecasting.Timestral.Tests/BatchTests.cs new file mode 100644 index 0000000..8f651b6 --- /dev/null +++ b/tests/Tsfm.Forecasting.Timestral.Tests/BatchTests.cs @@ -0,0 +1,84 @@ +using Xunit; + +namespace Tsfm.Forecasting.Timestral.Tests; + +/// Batched calls against the single calls they batch, on the tiny model. A pass +/// with a known-future entry widens every row's encoder over the horizon; entries without +/// one must be unaffected by that. +public class BatchTests +{ + private const int Horizon = 12; + + private static JointSeries Entry(int variates, int context, bool[] known, int target, int seed) + { + var r = new Random(seed); + var total = context + Horizon + 3; + var s = new float[variates * total]; + for (var v = 0; v < variates; v++) + for (var t = 0; t < total; t++) s[v * total + t] = (v + 1) * 10f + MathF.Sin(t / 3f + v) + (float)r.NextDouble(); + return new JointSeries(s, variates, context, known, target); + } + + private static readonly JointSeries[] Mixed = + [ + Entry(3, 37, [false, false, true], 1, 1), + Entry(1, 9, [], 0, 2), + Entry(2, 70, [false, true], 0, 3), + Entry(1, 1, [], 0, 4), + Entry(2, 16, [], 1, 5), + ]; + + [Theory] + [InlineData(64, true)] + [InlineData(64, false)] + [InlineData(3, true)] // several passes, some entries wider than the budget + public void JointBatchMatchesSingles(int maxRows, bool averaging) + { + var f = InvariantTests.Forecaster(); + f.MaxRowsPerPass = maxRows; + f.SymmetricAveraging = averaging; + var q = f.Config.NumQuantiles; + var batched = new double[f.BatchOutputLength(Mixed.Length, Horizon)]; + f.ForecastJointBatch(Mixed, Horizon, batched); + + for (var i = 0; i < Mixed.Length; i++) + { + var e = Mixed[i]; + var single = f.ForecastJoint(e.Series.Span, e.Variates, e.ContextBars, Horizon, e.KnownFuture.Span, e.TargetVariate); + for (var h = 0; h < Horizon; h++) + for (var k = 0; k < q; k++) + Assert.Equal(single[h, k], batched[(i * Horizon + h) * q + k], 3); + } + } + + [Fact] + public void OhlcvaBatchMatchesSingles() + { + var f = InvariantTests.Forecaster(); + var bars = new float[50 * 6]; + for (var i = 0; i < 50; i++) + for (var c = 0; c < 6; c++) bars[i * 6 + c] = c < 4 ? 100f + MathF.Sin(i / 4f) + c * 0.1f : 500f + i; + ReadOnlyMemory[] windows = [bars.AsMemory(0, 20 * 6), bars.AsMemory(6, 33 * 6), bars]; + var batched = new double[f.BatchOutputLength(windows.Length, Horizon)]; + f.ForecastBatch(windows, Horizon, batched); + var q = f.Config.NumQuantiles; + for (var i = 0; i < windows.Length; i++) + { + var single = f.Forecast(windows[i].Span, Horizon); + for (var h = 0; h < Horizon; h++) + for (var k = 0; k < q; k++) + Assert.Equal(single[h, k], batched[(i * Horizon + h) * q + k], 5); + } + } + + [Fact] + public void BadEntriesAndDestinationsAreRefused() + { + var f = InvariantTests.Forecaster(); + var size = f.BatchOutputLength(2, Horizon); + var ex = Assert.Throws(() => + f.ForecastJointBatch([Mixed[0], Mixed[1] with { KnownFuture = new[] { true } }], Horizon, new double[size])); + Assert.Contains("batch[1]", ex.Message); + Assert.Throws(() => f.ForecastJointBatch([Mixed[0], Mixed[1]], Horizon, new double[size + 1])); + } +} diff --git a/tests/Tsfm.Forecasting.Timestral.Tests/InvariantTests.cs b/tests/Tsfm.Forecasting.Timestral.Tests/InvariantTests.cs new file mode 100644 index 0000000..092194d --- /dev/null +++ b/tests/Tsfm.Forecasting.Timestral.Tests/InvariantTests.cs @@ -0,0 +1,154 @@ +using TorchSharp; +using Xunit; + +using static TorchSharp.torch; + +namespace Tsfm.Forecasting.Timestral.Tests; + +/// +/// Properties that hold for any weights, on a tiny random-init model — the reference's own +/// test configuration — so they run in CI without a checkpoint. +/// +public class InvariantTests +{ + internal static readonly TimestralConfig Tiny = new( + PatchSize: 8, DModel: 32, NumHeads: 4, MlpHiddenDim: 64, + BlockSchedule: ["time", "time_variate", "time", "time_variate"], + Quantiles: [0.05, 0.1, 0.25, 0.4, 0.5, 0.6, 0.75, 0.9, 0.95], + BlockSize: 4, ContextLength: 256, RopeTheta: 10000, RopePartialFactor: 0.5, NormEps: 1e-6, + AttnBias: false, MlpBias: false, TimeAttnGate: true, VariateAttnGate: true, MinimumScale: 1e-6, + UseAsinh: true, CovariateEnabled: true, SplitContextHorizonAttn: true, AttendMissingPatches: true, + QuantileWidthFactor: 3.0); + + internal static TimestralForecaster Forecaster(bool raw = false) + { + manual_seed(0); + var model = new TimestralModel(Tiny); + using (no_grad()) + foreach (var (_, p) in model.named_parameters()) p.normal_(0, 0.05); + model.eval(); + return new TimestralForecaster(model, new Device(DeviceType.CPU), "tiny") + { SymmetricAveraging = !raw, ClipNonNegative = !raw }; + } + + private const int Context = 37, Horizon = 20, Total = Context + Horizon; + private static readonly bool[] Known = [false, false, true]; + + internal static float[] Series(Func? tweak = null) + { + var s = new float[3 * Total]; + for (var t = 0; t < Total; t++) + { + s[t] = 50f + 3f * MathF.Sin(t / 2f); + s[Total + t] = -2f * MathF.Cos(t / 3f); + s[2 * Total + t] = 10f + t % 4; + } + if (tweak is not null) + for (var v = 0; v < 3; v++) + for (var t = 0; t < Total; t++) s[v * Total + t] += tweak(v, t); + return s; + } + + [Fact] + public void TargetFutureRegionIsNeverRead() + { + var f = Forecaster(); + Assert.Equal(f.ForecastJoint(Series(), 3, Context, Horizon, Known), + f.ForecastJoint(Series((v, t) => v < 2 && t >= Context ? 1e4f : 0f), 3, Context, Horizon, Known)); + } + + [Fact] + public void KnownFutureValuesReachTheTarget() + { + var f = Forecaster(); + Assert.NotEqual(f.ForecastJoint(Series(), 3, Context, Horizon, Known), + f.ForecastJoint(Series((v, t) => v == 2 && t >= Context ? 5f : 0f), 3, Context, Horizon, Known)); + } + + /// Covariates are exogenous: perturbing a target's context must not move a + /// known-future row's output. + [Fact] + public void KnownFutureRowsDoNotReadTargets() + { + var f = Forecaster(raw: true); + using var _ = no_grad(); + const int ctxPatches = 4, horizonPatches = 2, p = 8; + const int w = (ctxPatches + horizonPatches) * p; + Tensor Known(float level) + { + var values = zeros(2, w); + values[0, TensorIndex.Slice(0, ctxPatches * p)] = arange(ctxPatches * p, ScalarType.Float32) * 0.1f + level; + values[1] = arange(w, ScalarType.Float32).sin(); + var observed = ones([2, w], ScalarType.Bool); + observed[0, TensorIndex.Slice(ctxPatches * p)] = false; + var roles = tensor(new long[] { Role.Target, Role.KnownCovariate }); + return f.Model.PredictBlock(values, observed, zeros(2, ScalarType.Int64), roles, ctxPatches, horizonPatches)[1]; + } + Assert.True(Known(0f).allclose(Known(7f))); + } + + /// Negating every input reflects the averaged forecast exactly — the defining + /// property of the reference's symmetric averaging. + [Fact] + public void SymmetricAveragingIsOdd() + { + var f = Forecaster(); + f.ClipNonNegative = false; + var a = f.ForecastJoint(Series(), 3, Context, Horizon, Known); + var b = f.ForecastJoint(Series().Select(x => -x).ToArray(), 3, Context, Horizon, Known); + var q = a.GetLength(1); + for (var h = 0; h < Horizon; h++) + for (var k = 0; k < q; k++) + Assert.Equal(a[h, k], -b[h, q - 1 - k], 4); + } + + [Fact] + public void NonNegativeHistoryIsClippedAtZero() + { + var f = Forecaster(); + var s = Enumerable.Range(0, 40).Select(t => t % 5 == 0 ? 0f : 0.01f).ToArray(); + Assert.All(f.ForecastJoint(s, 1, 30, 10).Cast(), x => Assert.True(x >= 0)); + } + + [Fact] + public void AnyContextLengthAndMissingValuesAreAccepted() + { + var f = Forecaster(); + for (var context = 1; context <= 17; context += 4) + { + var s = Enumerable.Range(0, context + 3).Select(t => t == 0 && context > 2 ? float.NaN : 10f + t).ToArray(); + Assert.All(f.ForecastJoint(s, 1, context, 3).Cast(), x => Assert.True(double.IsFinite(x))); + } + } + + [Fact] + public void QuantilesAreOrdered() + { + var q = Forecaster().ForecastJoint(Series(), 3, Context, Horizon, Known); + for (var h = 0; h < Horizon; h++) + for (var k = 1; k < q.GetLength(1); k++) + Assert.True(q[h, k] >= q[h, k - 1]); + } + + [Fact] + public void ArgumentsAreValidated() + { + var f = Forecaster(); + var s = Series(); + Assert.Throws(() => f.ForecastJoint(s, 3, Context, Horizon, [true, false, true])); + Assert.Throws(() => f.ForecastJoint(s, 3, Context, f.MaxHorizon + 1)); + Assert.Throws(() => f.ForecastJoint(s, 3, Context, Horizon + 1, Known)); + var tooLong = new float[Tiny.ContextLength + 1]; + Assert.Throws(() => f.ForecastJoint(tooLong, 1, tooLong.Length, 1)); + Assert.Throws(() => f.Forecast(new float[7], 1)); + } + + [Fact] + public void ConfigRefusesUnknownKeys() + { + const string json = """{"patch_size":32,"d_model":192,"num_heads":3,"mlp_hidden_dim":512,"block_schedule":["time","time_variate"],"quantile_knots":[0.1,0.5,0.9],"block_size":32,"context_length":32768,"quantile_width_factor":3.0}"""; + Assert.Equal(3.0, TimestralConfig.Parse(json).QuantileWidthFactor); + Assert.Throws(() => TimestralConfig.Parse(json.Replace("\"block_size\"", "\"new_knob\":1,\"block_size\""))); + Assert.Throws(() => TimestralConfig.Parse(json.Replace("[0.1,0.5,0.9]", "[0.1,0.9]"))); + } +} diff --git a/tests/Tsfm.Forecasting.Timestral.Tests/Tsfm.Forecasting.Timestral.Tests.csproj b/tests/Tsfm.Forecasting.Timestral.Tests/Tsfm.Forecasting.Timestral.Tests.csproj new file mode 100644 index 0000000..b7a9da9 --- /dev/null +++ b/tests/Tsfm.Forecasting.Timestral.Tests/Tsfm.Forecasting.Timestral.Tests.csproj @@ -0,0 +1,19 @@ + + + + false + Tsfm.Forecasting.Timestral.Tests + + + + + + + + + + + + + + From eea6fe420e52668c0f30297a0de3855fedb228e5 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 14:46:33 +0000 Subject: [PATCH 5/9] docs: name Timestral in the README byline Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01F1dTUhvq2UnFWdBmH7AA8G --- README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index d55843c..1ad3dc9 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,9 @@ # Tsfm.Forecasting Probabilistic forecasting with native .NET inference for time-series foundation models, on CPU or GPU via -TorchSharp. Currently [Kronos](https://github.com/shiyu-coder/Kronos), [TimesFM 3.0](https://github.com/google-research/timesfm) -and [t0](https://huggingface.co/theforecastingcompany/t0-beta) (t0-alpha, t0-beta). +TorchSharp. Currently [Kronos](https://github.com/shiyu-coder/Kronos), [TimesFM 3.0](https://github.com/google-research/timesfm), +[t0](https://huggingface.co/theforecastingcompany/t0-beta) (t0-alpha, t0-beta) +and [Timestral](https://huggingface.co/collections/mistralai/timestral) (5M, 28M, 123M, 209M). [![ci](https://github.com/markhammond/tsfm-forecasting/actions/workflows/ci.yml/badge.svg)](https://github.com/markhammond/tsfm-forecasting/actions/workflows/ci.yml) From 6427736ab9b32da457b19d8d707e81b70aa7845e Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 19 Oct 2026 14:56:29 +0000 Subject: [PATCH 6/9] fix(timestral): package the weights under their own LICENSE.txt Mistral's Modified MIT licence is not MIT-compatible, so each Timestral weights package now declares LICENSE.txt and packs its own LICENSE.txt at the package root: the Mistral licence verbatim with that package's pinned revision, then the MIT notice for the small loader class. The repository's MIT LICENSE is no longer packed into those two packages (PackRepositoryLicense=false), so they carry one licence, the right one. Every other package is unchanged. verify-packages.sh now checks the weights packages declare LICENSE.txt, carry no stray MIT LICENSE, and require licence acceptance. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01F1dTUhvq2UnFWdBmH7AA8G --- Directory.Build.props | 3 +- LICENSE-Timestral-Weights | 30 ---------- NOTICE | 2 +- README.md | 2 +- scripts/verify-packages.sh | 13 ++++- .../LICENSE.txt | 57 +++++++++++++++++++ ...m.Forecasting.Timestral.Weights.28M.csproj | 14 ++++- .../LICENSE.txt | 57 +++++++++++++++++++ ...fm.Forecasting.Timestral.Weights.5M.csproj | 14 ++++- 9 files changed, 151 insertions(+), 41 deletions(-) delete mode 100644 LICENSE-Timestral-Weights create mode 100644 src/Tsfm.Forecasting.Timestral.Weights.28M/LICENSE.txt create mode 100644 src/Tsfm.Forecasting.Timestral.Weights.5M/LICENSE.txt diff --git a/Directory.Build.props b/Directory.Build.props index 30e2302..b23a8f8 100644 --- a/Directory.Build.props +++ b/Directory.Build.props @@ -40,7 +40,8 @@ - + diff --git a/LICENSE-Timestral-Weights b/LICENSE-Timestral-Weights deleted file mode 100644 index b9396bf..0000000 --- a/LICENSE-Timestral-Weights +++ /dev/null @@ -1,30 +0,0 @@ -Timestral weights — Mistral AI Modified MIT License -=================================================== - -This package redistributes, unmodified, the Timestral model weights published by Mistral AI -on the Hugging Face Hub. The weights are NOT under this repository's MIT or Apache-2.0 -licences. They are under the licence below, reproduced verbatim as its condition 1 requires. - -Note condition 2: no rights under it are granted to anyone whose company (or employer) had -global consolidated monthly revenue above USD 20 million in the preceding month. That -applies to every consumer of this package, not only to its publisher. - -The small loader class compiled into this package is this repository's own code, MIT -licensed; see LICENSE. - -Source revisions: - Timestral-5M https://huggingface.co/mistralai/Timestral-5M @ bdd20cfc15a3fea0e4cb5ccef37f8d3ca7aa335f - Timestral-28M https://huggingface.co/mistralai/Timestral-28M @ 180698bbf5a50dda3516aa6419798437fb03a430 - --------------------------------------------------------------------------------- - -Modified MIT License - -Attribution notice: 2026 - Mistral AI - -Permission is hereby granted, free of charge, to any person obtaining a copy of the weights of this model and associated documentation files (the “Model”), to deal in the Model without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Model, and to permit persons to whom the Model is furnished to do so, subject to the following conditions: - -1. The above attribution notice and this permission notice shall be included in all copies or substantial portions of the Model. -2. You are not authorized to exercise any rights under this license if the global consolidated monthly revenue of your company (or that of your employer) exceeds $20 million (or its equivalent in another currency) for the preceding month. This restriction in (b) applies to the Model and any derivatives, modifications, or combined works based on it, whether provided by Mistral AI or by a third party. You may contact Mistral AI (sales@mistral.ai) to request a commercial license, which Mistral AI may grant you at its sole discretion, or choose to use the Model on Mistral AI's hosted services available at https://mistral.ai/. - -THE MODEL IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL MISTRAL AI BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE MODEL OR THE USE OR OTHER DEALINGS IN THE MODEL. diff --git a/NOTICE b/NOTICE index ef5e785..b6cae7c 100644 --- a/NOTICE +++ b/NOTICE @@ -110,7 +110,7 @@ Tsfm.Forecasting.Timestral.Weights.5M and .28M — Mistral AI Modified MIT These packages redistribute, unmodified, the Timestral-5M and Timestral-28M weights published by Mistral AI. The weights are under Mistral AI's Modified MIT licence, -reproduced verbatim in LICENSE-Timestral-Weights inside each package. It is NOT an +reproduced verbatim in LICENSE.txt, each package's licence file. It is NOT an open-source licence: its condition 2 grants no rights to anyone whose company (or employer) had global consolidated monthly revenue above USD 20 million in the preceding month. Attribution notice: 2026 - Mistral AI. diff --git a/README.md b/README.md index 1ad3dc9..e34485b 100644 --- a/README.md +++ b/README.md @@ -204,7 +204,7 @@ What differs from TimesFM: > licence whose condition 2 withholds *every* right from anyone whose company — or > employer — had consolidated monthly revenue above USD 20 million in the preceding month. > That binds each consumer, not only whoever publishes a package. The weights packages -> carry the licence verbatim (`LICENSE-Timestral-Weights`, and `Timestral5M.License` at +> carry the licence verbatim (`LICENSE.txt`, the package licence, and `Timestral5M.License` at > run time) and require its acceptance on install. The code here is Apache-2.0 and grants > nothing in the weights. diff --git a/scripts/verify-packages.sh b/scripts/verify-packages.sh index b21c58a..67ade8a 100755 --- a/scripts/verify-packages.sh +++ b/scripts/verify-packages.sh @@ -15,10 +15,19 @@ for pkg in "$DIR"/*.nupkg; do contents=$(unzip -l "$pkg") required_files="LICENSE README.md" - # Revenue-restricted weights must carry their own licence, and demand acceptance of it. + # Revenue-restricted weights carry their own LICENSE.txt as the package licence — and not + # the repository's MIT LICENSE, which does not cover them — and demand its acceptance. case "$name" in Tsfm.Forecasting.Timestral.Weights.*) - required_files="$required_files LICENSE-Timestral-Weights" + required_files="LICENSE.txt README.md" + if grep -qE '[[:space:]]LICENSE$' <<< "$contents"; then + echo " STRAY carries the repository's MIT LICENSE beside the weights' licence"; fail=1 + fi + if unzip -p "$pkg" '*.nuspec' | grep -q 'LICENSE.txt'; then + echo " ok licence file is LICENSE.txt" + else + echo " MISSING LICENSE.txt"; fail=1 + fi if unzip -p "$pkg" '*.nuspec' | grep -q 'true'; then echo " ok requires licence acceptance" else diff --git a/src/Tsfm.Forecasting.Timestral.Weights.28M/LICENSE.txt b/src/Tsfm.Forecasting.Timestral.Weights.28M/LICENSE.txt new file mode 100644 index 0000000..3ad3812 --- /dev/null +++ b/src/Tsfm.Forecasting.Timestral.Weights.28M/LICENSE.txt @@ -0,0 +1,57 @@ +Tsfm.Forecasting.Timestral.Weights.28M +====================================== + +This package contains two works under two licences. + +1. The Timestral-28M model weights, redistributed unmodified from Mistral AI: + https://huggingface.co/mistralai/Timestral-28M @ 180698bbf5a50dda3516aa6419798437fb03a430 + They are under Mistral AI's Modified MIT License, reproduced verbatim below as its + condition 1 requires. It is NOT the MIT License and is not an open-source licence: + condition 2 grants no rights to anyone whose company (or employer) had global + consolidated monthly revenue above USD 20 million in the preceding month. That + applies to every user of this package, not only to its publisher. + +2. A small loader class (Tsfm.Forecasting.Timestral.Weights.Timestral28M), which is this + repository's own code under the MIT License, reproduced at the end of this file. + The MIT License grants nothing in the weights. + +================================================================================ +1. Timestral-28M weights — Mistral AI Modified MIT License +================================================================================ + +Modified MIT License + +Attribution notice: 2026 - Mistral AI + +Permission is hereby granted, free of charge, to any person obtaining a copy of the weights of this model and associated documentation files (the “Model”), to deal in the Model without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Model, and to permit persons to whom the Model is furnished to do so, subject to the following conditions: + +1. The above attribution notice and this permission notice shall be included in all copies or substantial portions of the Model. +2. You are not authorized to exercise any rights under this license if the global consolidated monthly revenue of your company (or that of your employer) exceeds $20 million (or its equivalent in another currency) for the preceding month. This restriction in (b) applies to the Model and any derivatives, modifications, or combined works based on it, whether provided by Mistral AI or by a third party. You may contact Mistral AI (sales@mistral.ai) to request a commercial license, which Mistral AI may grant you at its sole discretion, or choose to use the Model on Mistral AI's hosted services available at https://mistral.ai/. + +THE MODEL IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL MISTRAL AI BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE MODEL OR THE USE OR OTHER DEALINGS IN THE MODEL. + +================================================================================ +2. Loader code — MIT License +================================================================================ + +MIT License + +Copyright (c) 2026 Mark Hammond + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/src/Tsfm.Forecasting.Timestral.Weights.28M/Tsfm.Forecasting.Timestral.Weights.28M.csproj b/src/Tsfm.Forecasting.Timestral.Weights.28M/Tsfm.Forecasting.Timestral.Weights.28M.csproj index b85615e..5692c0d 100644 --- a/src/Tsfm.Forecasting.Timestral.Weights.28M/Tsfm.Forecasting.Timestral.Weights.28M.csproj +++ b/src/Tsfm.Forecasting.Timestral.Weights.28M/Tsfm.Forecasting.Timestral.Weights.28M.csproj @@ -4,13 +4,16 @@ Tsfm.Forecasting.Timestral.Weights.28M Tsfm.Forecasting.Timestral.Weights Tsfm.Forecasting.Timestral.Weights.28M - Timestral-28M (27.8M) weights for Tsfm.Forecasting.Timestral, embedded as assembly resources (~111 MB), redistributed unmodified from Mistral AI at the revision recorded in LICENSE-Timestral-Weights. NOT open source: Mistral's Modified MIT licence grants no rights to anyone whose company or employer exceeds USD 20M monthly revenue. + Timestral-28M (27.8M) weights for Tsfm.Forecasting.Timestral, embedded as assembly resources (~111 MB), redistributed unmodified from Mistral AI at the revision recorded in LICENSE.txt. NOT open source: Mistral's Modified MIT licence grants no rights to anyone whose company or employer exceeds USD 20M monthly revenue. timestral;mistral;timeseries;forecasting;weights;model true - LICENSE-Timestral-Weights + LICENSE.txt true + + false $(MSBuildThisFileDirectory)../../checkpoints $(TimestralCheckpointRoot)/Timestral-28M @@ -25,7 +28,12 @@ - + + + + - LICENSE-Timestral-Weights + LICENSE.txt true + + false $(MSBuildThisFileDirectory)../../checkpoints $(TimestralCheckpointRoot)/Timestral-5M @@ -25,7 +28,12 @@ - + + + +