diff --git a/THIRD_PARTY_LICENSES.md b/THIRD_PARTY_LICENSES.md index 7dac166b03..724dfd4820 100644 --- a/THIRD_PARTY_LICENSES.md +++ b/THIRD_PARTY_LICENSES.md @@ -43,16 +43,6 @@ Its weights are not bundled with this project and may be downloaded only after t The [`sdm/models/timesfm3/`](sdm/models/timesfm3/) implementation contains code derived from [`TimesFM`](https://github.com/google-research/timesfm). Its weights are not bundled with this project and may be downloaded only after the user accepts their separate license. -## Timestral - -- Source: https://github.com/mistralai/timestral -- Code license: Apache License 2.0 -- License terms: [`sdm/models/timestral/LICENSE`](sdm/models/timestral/LICENSE) -- Optional pretrained weights: https://huggingface.co/collections/mistralai/timestral -- Weights license: [Modified MIT License](https://huggingface.co/mistralai/Timestral-209M/blob/main/LICENSE) - -The [`sdm/models/timestral/`](sdm/models/timestral/) implementation contain code derived from [`Timestral`](https://github.com/mistralai/timestral). - ## DPM-Solver - Source: https://github.com/LuChengTHU/dpm-solver @@ -76,3 +66,12 @@ Portions of [`CONTRIBUTING.md`](CONTRIBUTING.md) are adapted from the [PyTorch c ## Contributor Covenant [`CODE_OF_CONDUCT.md`](CODE_OF_CONDUCT.md) is adapted from [Contributor Covenant version 1.4](https://www.contributor-covenant.org/version/1/4/code-of-conduct/), distributed under the MIT License. The complete copyright notice and license terms are distributed in [`third_party/contributor-covenant/LICENSE`](third_party/contributor-covenant/LICENSE). + +## Timestral + +- Source: [mistralai/timestral](https://github.com/mistralai/timestral), revision `a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff`. +- Source license: [Apache-2.0](sdm/models/timestral/LICENSE). +- Weights license: [Modified MIT](sdm/models/timestral/WEIGHTS_LICENSE), including a company revenue restriction. +- Checkpoints: `mistralai/Timestral-5M` at `bdd20cfc15a3fea0e4cb5ccef37f8d3ca7aa335f`, `Timestral-28M` at `180698bbf5a50dda3516aa6419798437fb03a430`, `Timestral-123M` at `5a147287addda85f6acd8da3fb8858fe9af7e8f7`, and `Timestral-209M` at `9874d252ab0aae16f7f7c999f98f98d6ad7d8365`. + +The `sdm/models/timestral/` numerical implementation, `sdm.nn.MonotoneQuantiles`, and the running-statistics computation in `sdm.nn.RevIN` contain code derived from Timestral. The SDM wrapper preserves the upstream default forecast, including categorical encoding, sign averaging, advancing normalization frames, and nonnegative clipping. Prediction parity is checked with the pinned 5M weights; larger checkpoints share the implementation but have not been parity-tested. diff --git a/benchmark/forecasting/timestral_parity.md b/benchmark/forecasting/timestral_parity.md new file mode 100644 index 0000000000..ff22e961b3 --- /dev/null +++ b/benchmark/forecasting/timestral_parity.md @@ -0,0 +1,77 @@ +# Timestral 5M parity + +Run from the repository root with the upstream reference installed: + +```bash +pip install git+https://github.com/mistralai/timestral.git@a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff +python -m benchmark.forecasting.timestral_parity --device cpu --repeats 20 +python -m benchmark.forecasting.timestral_parity --device cuda --repeats 20 +``` + +The script verifies the installed source revision and uses `mistralai/Timestral-5M` revision `bdd20cfc15a3fea0e4cb5ccef37f8d3ca7aa335f`. Downloading weights is subject to their [Modified MIT license](../../sdm/models/timestral/WEIGHTS_LICENSE). + +It compares upstream `TimestralForecaster.forecast` with SDM's public `Timestral.forward`, using default sign averaging, nonnegative clipping, and normalization refinement. It also compares SDM's fitted `predict` with its one-shot output. Deterministic inputs cover a short patch, univariate and multivariate targets, independent batches, missing and entirely missing targets, constant and positive targets, numerical and binary covariates, categorical past and known-future covariates, unseen and missing categories, missing target history with categorical covariates, a 1,030-step forecast, and truncation of a 32,775-step context. + +Recorded on 2026-10-18 with PyTorch `2.14.0+cu130`, a Cortex-X925 CPU using two PyTorch threads, and an NVIDIA GB10 GPU: + +| Comparison | Result | +| ----------------------------- | --------------------------------------------------------------- | +| CPU, horizons up to 65 steps | Maximum absolute error `3.25e-5` | +| CPU, 1,030-step horizon | Maximum absolute error `0.0801`; maximum scaled error `1.82e-5` | +| CPU, all 15 cases | Pass `atol=rtol=2e-5` | +| CUDA, horizons up to 65 steps | Maximum absolute error `1.96e-5` | +| CUDA, 1,030-step horizon | Maximum absolute error `0.0547`; maximum scaled error `1.66e-5` | +| CUDA, all 15 cases | Pass `atol=rtol=1e-4` | + +Scaled error is `abs(actual - reference) / (1 + abs(reference))`. The long-horizon example produces large tail quantiles; its absolute error should be interpreted alongside its scaled error. Timestral uses SDM's shared `TransformerBlock`, including `Attention`, `RotaryEmbedding` with explicit positions, and input-dependent sigmoid output gating. Time and variate attention are separate residual blocks, with one normalized `SwiGLU` MLP after the final attention step of each upstream layer. A single full mask replaces upstream's split context/forecast attention calls. Checkpoint loading permutes query/key channels to match SDM's partial-rotation layout while leaving values unchanged. This preserves the attention computation mathematically; reordered floating-point reductions introduce small differences on CPU and CUDA. + +Only 5M was checked against pretrained weights. The 28M, 123M, and 209M variants use pinned configurations and checkpoint loading but have not been parity-tested. Parity is scoped to float32; lower precision is unvalidated. + +## Shared normalization + +Timestral uses `RevIN.running_stats(correction=1)` for causal context normalization and advancing forecast statistics, and `RevIN.stats(correction=0)` for known-future covariates. The model keeps its scale floor, binary-covariate handling, missing-value replacement, and `asinh` transform. The separate `CausalPatchNorm` module and private Welford implementations have been removed. + +After merging main at `50ff254e8`, all 15 cases pass again through public `forward` and fitted `predict`, using the same pinned upstream source and 5M checkpoint. CPU passes `atol=rtol=2e-5`; CUDA passes `atol=rtol=1e-4`. For horizons up to 65 steps, maximum absolute errors are `3.24e-5` on CPU and `2.19e-5` on CUDA. For the 1,030-step horizon, maximum scaled errors are `1.81e-5` on CPU and `1.91e-5` on CUDA; maximum absolute errors are `0.0801` and `0.0671`, respectively. The runtime measurements below predate this normalization change. + +The integration needs no upstream package at runtime. Optional pretrained pytest comparisons skip when upstream or cached 5M weights are unavailable. The ordinary public-contract and neural-block tests use random weights and run offline. SDM represents categorical values with category indices; for multivariate targets those indices become the raw numerical covariates, matching upstream when supplied the same codes. + +The default recipe uses `AlignCategories(min_frequency=0, unknown="append")` to align independently tensorized vocabularies. Context codes and unused dictionary entries are retained; new query categories are appended to the returned vocabulary without changing fitted state. Numeric vocabularies use tensor lookups; string vocabularies use SDM's batched relational join, with cuDF on CUDA when installed and the existing CPU fallback otherwise. Forecast tensor execution and target encoding stay on the input device. + +## Runtime versus upstream + +The [runtime benchmark](timestral_runtime.py) checks prediction parity before timing each workload. Reproduce it with: + +```bash +python -m benchmark.forecasting.timestral_runtime --device cpu --output benchmark/forecasting/timestral_runtime_cpu.json +python -m benchmark.forecasting.timestral_runtime --device cuda --output benchmark/forecasting/timestral_runtime_cuda.json +``` + +Measured before moving category alignment into the recipe, on the same CPU/GPU and PyTorch version described above, with two CPU threads, float32, eager execution, one series per call, and the default forecast options. Values are median milliseconds across 30 calls after five warmups. Each pair alternates execution order; CUDA is synchronized before and after each individual timed call. These are public API latencies, including preprocessing and output construction, with inputs already on the target device. Loading, initial table construction, and warmup are excluded. CPU and CUDA benchmarks ran sequentially. + +| Workload | Context / horizon | CPU SDM | CPU upstream | CUDA SDM | CUDA upstream | +| ---------------------- | ----------------- | ------- | ------------ | -------- | ------------- | +| Univariate | 128 / 96 | 79.03 | 68.79 | 9.25 | 9.09 | +| Univariate | 1,024 / 96 | 82.16 | 71.09 | 10.03 | 9.84 | +| Three targets | 1,024 / 96 | 85.71 | 73.62 | 10.56 | 10.32 | +| Numerical covariates | 1,024 / 96 | 92.55 | 80.50 | 12.06 | 11.47 | +| Categorical covariates | 1,024 / 96 | 96.02 | 82.22 | 14.92 | 11.90 | +| Long horizon | 79 / 1,030 | 163.30 | 143.34 | 23.54 | 23.65 | + +SDM takes roughly 14–17% longer on CPU in these workloads. On CUDA, most workloads are within 5% of upstream; the categorical case takes roughly 25% longer. The categorical comparison includes SDM's dictionary alignment, while upstream accepts numeric category codes. Neither implementation is compiled or uses lower precision. No attribution of the remaining runtime difference to individual operations has been profiled. + +The [CPU report](timestral_runtime_cpu.json) and [CUDA report](timestral_runtime_cuda.json) include all individual samples, quartiles, revisions, and benchmark settings. + +## Shared transformer refactor + +Remeasured on 2026-10-18 after replacing the private attention and residual blocks with `sdm.nn.TransformerBlock`. The same benchmark settings apply, with 20 timed calls after five warmups (`--repeats 20`). The CUDA baseline was measured immediately before this refactor; CPU and CUDA runs were sequential. All values below are median milliseconds. + +| Workload | CPU SDM | CPU upstream | CUDA SDM before | CUDA SDM after | CUDA upstream after | +| ----------------------------------- | ------- | ------------ | --------------- | -------------- | ------------------- | +| Univariate (128 / 96) | 80.70 | 68.87 | 9.91 | 9.43 | 9.12 | +| Univariate (1,024 / 96) | 84.05 | 71.73 | 10.72 | 10.33 | 9.84 | +| Three targets (1,024 / 96) | 85.74 | 74.24 | 11.37 | 10.84 | 10.31 | +| Numerical covariates (1,024 / 96) | 92.79 | 80.73 | 12.88 | 12.11 | 11.41 | +| Categorical covariates (1,024 / 96) | 94.89 | 82.06 | 14.39 | 13.47 | 11.86 | +| Long horizon (79 / 1,030) | 163.31 | 139.79 | 24.77 | 23.81 | 23.58 | + +The full-mask implementation showed no latency regression in these workloads. Before/after timings are separate runs, and upstream CUDA timings also decreased, so these measurements do not isolate a speedup caused by the refactor. The public 5M parity checks still pass on both devices at the tolerances above. Larger batches, longer covariate timelines, and lower precision remain unbenchmarked. diff --git a/benchmark/forecasting/timestral_parity.py b/benchmark/forecasting/timestral_parity.py new file mode 100644 index 0000000000..cea4828266 --- /dev/null +++ b/benchmark/forecasting/timestral_parity.py @@ -0,0 +1,217 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compare public Timestral forecasts against the pinned upstream 5M model. + +Install git+https://github.com/mistralai/timestral.git@a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff +then run python -m benchmark.forecasting.timestral_parity --device cpu +or --device cuda. This downloads the 5M weights under their Modified MIT terms. +""" + +import argparse +import json +from importlib.metadata import distribution +from time import perf_counter + +import torch + +from sdm import CategoricalTensor, Stype, TableTensor +from sdm.models import Timestral + +SOURCE_REVISION = "a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff" +WEIGHT_REVISION = "bdd20cfc15a3fea0e4cb5ccef37f8d3ca7aa335f" + + +def make_case( + name: str, + device: str, + *, + context_length: int = 79, + prediction_length: int = 65, +) -> tuple[TableTensor, TableTensor, TableTensor, dict]: + t, h, v, batch = context_length, prediction_length, 1, 1 + if name == "multivariate": + v = 3 + if name == "batched": + batch = 2 + if name == "long_horizon": + h = 1030 + if name == "context_cap": + t, h = 32775, 1 + if name == "single_patch": + t, h = 7, 1 + time = torch.arange(t + h, device=device).float() + target = (time[:t] / 9).sin().expand(batch, v, t).clone() + target += torch.arange(v, device=device)[None, :, None] * 0.2 + target += torch.arange(batch, device=device)[:, None, None] * 0.3 + if name == "positive": + target = target.abs() + 1 + if name == "constant": + target.fill_(2) + if name in {"missing", "categorical_missing"}: + target[..., :8] = torch.nan + target[..., 12:15] = torch.nan + if name == "all_missing": + target.fill_(torch.nan) + reference = {} + context = TableTensor(size=(batch, t), device=device) + query = TableTensor(size=(batch, h), device=device) + if name in { + "numeric", + "indicator", + "categorical", + "categorical_multi", + "categorical_missing", + }: + if name == "categorical_multi": + target = target.expand(batch, 2, t).clone() + v = 2 + past = (time[:t] / 13).cos().expand(batch, 1, t).clone() + known = (time / 7).sin().expand(batch, 1, t + h).clone() + if name == "indicator": + known = (known > 0).float() + if name.startswith("categorical"): + past = (time[:t].long() % 3).expand(batch, 1, t).clone() + known = (time.long() % 4).expand(batch, 1, t + h).clone() + past[..., 3:6] = -1 + known[..., 3:6] = -1 + known[..., t + 4 :] = 4 # Category first seen in the horizon. + + def table(codes: torch.Tensor, columns: list[str]) -> TableTensor: + return TableTensor( + columns={Stype.categorical: columns}, + categorical=CategoricalTensor( + codes.transpose(-2, -1), + [torch.arange(5, device=device) for _ in columns], + ), + ) + + context = table( + torch.cat([known[..., :t], past], dim=1), ["known", "past"] + ) + query = table(known[..., t:], ["known"]) + reference = { + "categorical_past_covariates": past.float().masked_fill( + past < 0, torch.nan + ), + "categorical_known_covariates": known.float().masked_fill( + known < 0, torch.nan + ), + } + else: + past[..., 4:6] = torch.nan + context = TableTensor.from_tensor( + torch.cat([known[..., :t], past], dim=1).transpose(-2, -1), + columns=["known", "past"], + ) + query = TableTensor.from_tensor( + known[..., t:].transpose(-2, -1), columns=["known"] + ) + reference = {"past_covariates": past, "known_covariates": known} + targets = TableTensor.from_tensor( + target.transpose(-2, -1), columns=[f"target{i}" for i in range(v)] + ) + return context, targets, query, reference + + +def main() -> None: + from timestral import TimestralForecaster # noqa: PLC0415 + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", default="cpu") + parser.add_argument("--repeats", type=int, default=5) + args = parser.parse_args() + installed = json.loads( + distribution("timestral").read_text("direct_url.json") or "{}" + ) + assert installed.get("vcs_info", {}).get("commit_id") == SOURCE_REVISION + torch.set_num_threads(2) + model = Timestral(device=args.device, accept_license=True) + upstream = TimestralForecaster.from_pretrained( + "mistralai/Timestral-5M", + revision=WEIGHT_REVISION, + device=args.device, + ) + cases = [ + "single_patch", + "univariate", + "multivariate", + "batched", + "missing", + "all_missing", + "constant", + "positive", + "numeric", + "indicator", + "categorical", + "categorical_multi", + "categorical_missing", + "long_horizon", + "context_cap", + ] + report = {} + for name in cases: + x, y, query, covariates = make_case(name, args.device) + target = y.numerical.transpose(-2, -1) + expected = ( + upstream.forecast(target, query.size(-2), **covariates) + .quantiles.transpose(-3, -2) + .flatten(-2) + ) + actual = model(x, y, query, num_estimators=1).numerical + tol = 1e-4 if args.device.startswith("cuda") else 2e-5 + torch.testing.assert_close(actual, expected, atol=tol, rtol=tol) + model.fit( + x, + y, + context_only_columns=["past"] + if "past" in x.column_names + else None, + num_estimators=1, + ) + torch.testing.assert_close(model.predict(query).numerical, actual) + report[name] = { + "max_absolute_error": (actual - expected).abs().max().item(), + "max_scaled_error": ( + (actual - expected).abs() / (1 + expected.abs()) + ) + .max() + .item(), + } + print(name, report[name], flush=True) + x, y, query, _ = make_case("numeric", args.device) + for _ in range(3): + model(x, y, query, num_estimators=1) + if args.device.startswith("cuda"): + start, end = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), + ) + start.record() + for _ in range(args.repeats): + model(x, y, query, num_estimators=1) + end.record() + end.synchronize() + elapsed = start.elapsed_time(end) / args.repeats + else: + start_time = perf_counter() + for _ in range(args.repeats): + model(x, y, query, num_estimators=1) + elapsed = 1000 * (perf_counter() - start_time) / args.repeats + print( + json.dumps( + { + "device": args.device, + "torch": torch.__version__, + "source_revision": SOURCE_REVISION, + "weight_revision": WEIGHT_REVISION, + "cases": report, + "numeric_forward_ms": elapsed, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/benchmark/forecasting/timestral_runtime.py b/benchmark/forecasting/timestral_runtime.py new file mode 100644 index 0000000000..8b26a1fc25 --- /dev/null +++ b/benchmark/forecasting/timestral_runtime.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark public SDM and upstream Timestral-5M forecast latency.""" + +import argparse +import json +import statistics +from importlib.metadata import distribution +from pathlib import Path +from time import perf_counter +from typing import Any + +import torch + +from benchmark.forecasting.timestral_parity import ( + SOURCE_REVISION, + WEIGHT_REVISION, + make_case, +) +from sdm import TableTensor +from sdm.models import Timestral + + +def main() -> None: + from timestral import TimestralForecaster # noqa: PLC0415 + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", default="cpu") + parser.add_argument("--threads", type=int, default=2) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--repeats", type=int, default=30) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + installed = json.loads( + distribution("timestral").read_text("direct_url.json") or "{}" + ) + assert installed.get("vcs_info", {}).get("commit_id") == SOURCE_REVISION + torch.set_num_threads(args.threads) + device = torch.device(args.device) + model = Timestral(device=device, accept_license=True) + upstream = TimestralForecaster.from_pretrained( + "mistralai/Timestral-5M", + revision=WEIGHT_REVISION, + device=device, + ) + cases = [ + ("univariate", 128, 96), + ("univariate", 1024, 96), + ("multivariate", 1024, 96), + ("numeric", 1024, 96), + ("categorical", 1024, 96), + ("long_horizon", 79, 1030), + ] + results = [] + for case, context_length, horizon in cases: + x, y, query, covariates = make_case( + case, + args.device, + context_length=context_length, + prediction_length=horizon, + ) + target = y.numerical.transpose(-2, -1) + + def sdm_forecast( + x: TableTensor = x, + y: TableTensor = y, + query: TableTensor = query, + ) -> torch.Tensor: + return model(x, y, query, num_estimators=1).numerical + + def upstream_forecast( + target: torch.Tensor = target, + query: TableTensor = query, + covariates: dict = covariates, + ) -> torch.Tensor: + return ( + upstream.forecast(target, query.size(-2), **covariates) + .quantiles.transpose(-3, -2) + .flatten(-2) + ) + + # Check exactly the paths that will be timed, including their defaults. + actual, expected = sdm_forecast(), upstream_forecast() + tol = 1e-4 if device.type == "cuda" else 2e-5 + torch.testing.assert_close(actual, expected, atol=tol, rtol=tol) + for _ in range(args.warmup): + sdm_forecast() + upstream_forecast() + calls = {"sdm": sdm_forecast, "upstream": upstream_forecast} + samples: dict[str, list[float]] = {name: [] for name in calls} + for repeat in range(args.repeats): + # Alternate order to reduce systematic thermal/order bias. + order = ( + ("sdm", "upstream") if repeat % 2 == 0 else ("upstream", "sdm") + ) + for name in order: + if device.type == "cuda": + torch.cuda.synchronize(device) + start = perf_counter() + calls[name]() + if device.type == "cuda": + torch.cuda.synchronize(device) + samples[name].append(1000 * (perf_counter() - start)) + row: dict[str, Any] = { + "case": case, + "batch": target.size(0), + "targets": target.size(1), + "context": target.size(2), + "horizon": query.size(-2), + "median_ms": { + name: statistics.median(values) + for name, values in samples.items() + }, + "p25_p75_ms": { + name: [ + statistics.quantiles(values, n=4)[0], + statistics.quantiles(values, n=4)[2], + ] + for name, values in samples.items() + }, + "samples_ms": samples, + } + row["upstream_over_sdm"] = ( + row["median_ms"]["upstream"] / row["median_ms"]["sdm"] + ) + results.append(row) + print( + json.dumps( + { + key: value + for key, value in row.items() + if key != "samples_ms" + } + ), + flush=True, + ) + report = { + "device": str(device), + "torch": torch.__version__, + "gpu": torch.cuda.get_device_name(device) + if device.type == "cuda" + else None, + "threads": args.threads, + "warmup": args.warmup, + "repeats": args.repeats, + "source_revision": SOURCE_REVISION, + "weight_revision": WEIGHT_REVISION, + "timing": ( + "Per-call wall time, CUDA synchronized before and after each " + "call; alternating implementation order" + ), + "results": results, + } + if args.output is not None: + args.output.write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/benchmark/forecasting/timestral_runtime_cpu.json b/benchmark/forecasting/timestral_runtime_cpu.json new file mode 100644 index 0000000000..9fc51023d7 --- /dev/null +++ b/benchmark/forecasting/timestral_runtime_cpu.json @@ -0,0 +1,541 @@ +{ + "device": "cpu", + "torch": "2.14.0+cu130", + "gpu": null, + "threads": 2, + "warmup": 5, + "repeats": 30, + "source_revision": "a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff", + "weight_revision": "bdd20cfc15a3fea0e4cb5ccef37f8d3ca7aa335f", + "timing": "Per-call wall time, CUDA synchronized before and after each call; alternating implementation order", + "results": [ + { + "case": "univariate", + "batch": 1, + "targets": 1, + "context": 128, + "horizon": 96, + "median_ms": { + "sdm": 79.0257619228214, + "upstream": 68.78633901942521 + }, + "p25_p75_ms": { + "sdm": [ + 77.78316689655185, + 80.50416322657838 + ], + "upstream": [ + 67.9091740748845, + 69.40999050857499 + ] + }, + "samples_ms": { + "sdm": [ + 79.19317693449557, + 77.561428071931, + 77.08019996061921, + 78.10774422250688, + 84.13261291570961, + 82.12078781798482, + 79.68216482549906, + 80.3943679202348, + 77.75886706076562, + 77.06551207229495, + 77.79126684181392, + 79.89341183565557, + 80.83354914560914, + 79.9292039591819, + 76.4934360049665, + 78.88826704584062, + 81.75170188769698, + 92.77647105045617, + 96.02617681957781, + 78.14724813215435, + 80.9503011405468, + 78.43443006277084, + 78.99527391418815, + 77.55344407632947, + 77.63341208919883, + 79.86942795105278, + 78.43983895145357, + 77.87694712169468, + 79.90493206307292, + 79.05624993145466 + ], + "upstream": [ + 66.47958792746067, + 69.25779208540916, + 68.79457901231945, + 68.0084410123527, + 71.67719886638224, + 69.78031597100198, + 67.86293792538345, + 69.16708801873028, + 68.55523702688515, + 67.86180194467306, + 68.0078009609133, + 69.64372494257987, + 69.18945605866611, + 69.67503693886101, + 65.80588105134666, + 68.17397591657937, + 68.35237494669855, + 68.27527098357677, + 70.55048714391887, + 68.77809902653098, + 69.03088092803955, + 69.09921788610518, + 66.9796490110457, + 67.92458612471819, + 70.49784809350967, + 69.33207903057337, + 66.40857993625104, + 66.91612908616662, + 69.29513579234481, + 70.4225511290133 + ] + }, + "upstream_over_sdm": 0.8704293048968478 + }, + { + "case": "univariate", + "batch": 1, + "targets": 1, + "context": 1024, + "horizon": 96, + "median_ms": { + "sdm": 82.15700392611325, + "upstream": 71.09309104271233 + }, + "p25_p75_ms": { + "sdm": [ + 79.77258152095601, + 83.39686307590455 + ], + "upstream": [ + 69.72350127762184, + 72.8778667980805 + ] + }, + "samples_ms": { + "sdm": [ + 79.51856683939695, + 85.11587115935981, + 83.80360808223486, + 83.65208888426423, + 84.30776512250304, + 82.32350694015622, + 80.58929513208568, + 78.69143597781658, + 79.237368889153, + 79.60837497375906, + 77.9005300719291, + 82.58923306129873, + 83.31178780645132, + 80.98790096119046, + 83.21602805517614, + 87.11580885574222, + 84.08405398949981, + 84.10386298783123, + 83.31071585416794, + 83.1941890064627, + 83.17615697160363, + 80.27132996357977, + 78.65570904687047, + 81.35895407758653, + 81.99050091207027, + 82.38417794927955, + 80.61806415207684, + 80.8339340146631, + 78.16569716669619, + 79.82731703668833 + ], + "upstream": [ + 69.01040184311569, + 72.56405707448721, + 74.11977299489081, + 72.84879102371633, + 72.11330812424421, + 72.96509412117302, + 71.83146988973022, + 73.23006796650589, + 68.55190987698734, + 69.22955298796296, + 69.19940910302103, + 70.71904605254531, + 70.12168993242085, + 70.06797892972827, + 71.67556695640087, + 71.58151990734041, + 73.36976286023855, + 73.12397309578955, + 73.60731298103929, + 70.71527000516653, + 73.99452710524201, + 71.46713603287935, + 69.72230179235339, + 70.03677915781736, + 70.68979903124273, + 69.099185988307, + 69.91170812398195, + 69.72390110604465, + 71.58196810632944, + 69.4147190079093 + ] + }, + "upstream_over_sdm": 0.8653320793763233 + }, + { + "case": "multivariate", + "batch": 1, + "targets": 3, + "context": 1024, + "horizon": 96, + "median_ms": { + "sdm": 85.71017894428223, + "upstream": 73.6150179291144 + }, + "p25_p75_ms": { + "sdm": [ + 84.83623014762998, + 88.64118671044707 + ], + "upstream": [ + 73.1091175111942, + 76.13008347107098 + ] + }, + "samples_ms": { + "sdm": [ + 84.84089816920459, + 85.07428807206452, + 84.60681908763945, + 84.85659305006266, + 86.49947703815997, + 91.91129403188825, + 88.20980205200613, + 88.58645497821271, + 86.57776494510472, + 84.84428911469877, + 89.36590515077114, + 89.19081906788051, + 85.20870213396847, + 85.53882804699242, + 83.8657848071307, + 85.88152984157205, + 86.56807010993361, + 84.82222608290613, + 85.01732791773975, + 89.17497890070081, + 82.55300996825099, + 88.19298492744565, + 84.31536611169577, + 89.36756988987327, + 88.95373204723, + 83.48993211984634, + 83.3710681181401, + 88.80538190715015, + 86.38075808994472, + 84.98129690997303 + ], + "upstream": [ + 73.51174601353705, + 73.14539700746536, + 71.41971308737993, + 74.77357797324657, + 78.53345503099263, + 76.94010599516332, + 78.47723085433245, + 78.03150615654886, + 72.78440683148801, + 72.6885530166328, + 75.88462601415813, + 76.06166391633451, + 73.00027902238071, + 76.5295010060072, + 73.17654881626368, + 73.6792329698801, + 74.48905077762902, + 73.49565112963319, + 76.33534213528037, + 74.51597996987402, + 73.21865996345878, + 75.68171597085893, + 73.2433320954442, + 74.70016996376216, + 73.5508028883487, + 72.13916396722198, + 72.76888797059655, + 70.64904691651464, + 76.44566218368709, + 73.32104397937655 + ] + }, + "upstream_over_sdm": 0.8588830269152681 + }, + { + "case": "numeric", + "batch": 1, + "targets": 1, + "context": 1024, + "horizon": 96, + "median_ms": { + "sdm": 92.54666732158512, + "upstream": 80.50090540200472 + }, + "p25_p75_ms": { + "sdm": [ + 89.77919502649456, + 94.72719486802816 + ], + "upstream": [ + 78.21136974962428, + 82.27557619102299 + ] + }, + "samples_ms": { + "sdm": [ + 92.32090879231691, + 88.0122990347445, + 89.65628803707659, + 91.10029391013086, + 89.58889683708549, + 95.28047102503479, + 95.14522412791848, + 94.07439990900457, + 93.98020897060633, + 97.5752230733633, + 87.3493761755526, + 94.58785178139806, + 95.5016859807074, + 96.54335910454392, + 91.07443713583052, + 94.19150301255286, + 93.39101286605, + 89.3671689555049, + 90.6835610512644, + 91.45315503701568, + 94.47468491271138, + 91.73187310807407, + 89.78534303605556, + 95.22999194450676, + 94.53298896551132, + 92.77242585085332, + 92.02839899808168, + 88.62316003069282, + 89.76075099781156, + 97.16121107339859 + ], + "upstream": [ + 77.52325385808945, + 78.84410885162652, + 80.75236715376377, + 82.1600048802793, + 82.73321692831814, + 80.95222199335694, + 83.29918212257326, + 76.68153289705515, + 78.83425289765, + 81.65157679468393, + 79.43652104586363, + 81.36501908302307, + 80.24187618866563, + 82.83988805487752, + 78.29131302423775, + 81.47783484309912, + 80.4031218867749, + 77.82123703509569, + 78.65587109699845, + 77.42578303441405, + 80.59868891723454, + 77.97153992578387, + 99.25994696095586, + 103.70279592461884, + 82.7017140109092, + 81.92648692056537, + 79.07130802050233, + 77.42432691156864, + 77.64269388280809, + 82.62229012325406 + ] + }, + "upstream_over_sdm": 0.869841213431022 + }, + { + "case": "categorical", + "batch": 1, + "targets": 1, + "context": 1024, + "horizon": 96, + "median_ms": { + "sdm": 96.017403062433, + "upstream": 82.22335006576031 + }, + "p25_p75_ms": { + "sdm": [ + 93.20792782818899, + 98.86425902368501 + ], + "upstream": [ + 79.87546169897541, + 83.82622647332028 + ] + }, + "samples_ms": { + "sdm": [ + 93.84049801155925, + 93.91022589989007, + 92.87218609824777, + 95.87670699693263, + 92.75975497439504, + 98.25857891701162, + 97.74616709910333, + 99.81162403710186, + 95.88374802842736, + 101.01620806381106, + 93.52270909585059, + 95.33309610560536, + 95.6008539069444, + 100.5875701084733, + 99.73247302696109, + 91.38992382213473, + 89.66528018936515, + 97.7443908341229, + 97.16276288963854, + 92.28431805968285, + 93.28541508875787, + 99.20641197822988, + 98.75020803883672, + 100.40187602862716, + 96.99946804903448, + 98.53371395729482, + 99.82532006688416, + 96.15105809643865, + 92.97546604648232, + 92.50572510063648 + ], + "upstream": [ + 80.15352417714894, + 83.5324281360954, + 78.34516884759068, + 77.65130186453462, + 81.25374000519514, + 79.98101296834648, + 80.22121991962194, + 84.51171685010195, + 82.16971810907125, + 83.87775300070643, + 77.84928404726088, + 79.0355489589274, + 83.88228109106421, + 82.27698202244937, + 87.27123285643756, + 77.20589800737798, + 82.31701306067407, + 81.50188298895955, + 81.72328886575997, + 86.03324298746884, + 82.98880001530051, + 84.14394385181367, + 82.54950703121722, + 83.80905096419156, + 82.98716903664172, + 83.0429750494659, + 82.16743008233607, + 86.77493408322334, + 79.55880789086223, + 76.87604404054582 + ] + }, + "upstream_over_sdm": 0.8563379912732754 + }, + { + "case": "long_horizon", + "batch": 1, + "targets": 1, + "context": 79, + "horizon": 1030, + "median_ms": { + "sdm": 163.29976543784142, + "upstream": 143.33941892255098 + }, + "p25_p75_ms": { + "sdm": [ + 160.56545224273577, + 165.6599523848854 + ], + "upstream": [ + 139.7089892416261, + 144.74900579079986 + ] + }, + "samples_ms": { + "sdm": [ + 160.94030905514956, + 155.0492630340159, + 159.47321592830122, + 165.23974994197488, + 160.5940719600767, + 166.8330510146916, + 162.22178819589317, + 164.0700630377978, + 159.79750105179846, + 160.64944793470204, + 159.97357992455363, + 164.32394878938794, + 168.736238963902, + 162.88210381753743, + 163.44563593156636, + 163.63017889671028, + 159.9401410203427, + 160.0661240518093, + 161.48451413027942, + 168.34240080788732, + 166.05512192472816, + 163.8794899918139, + 160.47959309071302, + 165.93013098463416, + 163.91811310313642, + 165.9674111288041, + 163.15389494411647, + 165.5698928516358, + 168.2706749998033, + 160.62261699698865 + ], + "upstream": [ + 143.2235548272729, + 137.12006295099854, + 139.53510206192732, + 143.84142286144197, + 139.5603660494089, + 144.71775316633284, + 144.81199299916625, + 144.03428602963686, + 139.72818898037076, + 140.38642379455268, + 144.16490914300084, + 143.45528301782906, + 145.48881887458265, + 140.1501540094614, + 147.37343112938106, + 145.30371711589396, + 146.38862200081348, + 139.55775811336935, + 145.08176711387932, + 142.7995760459453, + 144.72801005467772, + 138.99077009409666, + 139.65139002539217, + 137.80487491749227, + 140.83440601825714, + 140.76175005175173, + 143.50710692815483, + 144.70423385500908, + 146.21166419237852, + 142.24271499551833 + ] + }, + "upstream_over_sdm": 0.8777686761412516 + } + ] +} diff --git a/benchmark/forecasting/timestral_runtime_cuda.json b/benchmark/forecasting/timestral_runtime_cuda.json new file mode 100644 index 0000000000..7de6b44f46 --- /dev/null +++ b/benchmark/forecasting/timestral_runtime_cuda.json @@ -0,0 +1,541 @@ +{ + "device": "cuda", + "torch": "2.14.0+cu130", + "gpu": "NVIDIA GB10", + "threads": 2, + "warmup": 5, + "repeats": 30, + "source_revision": "a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff", + "weight_revision": "bdd20cfc15a3fea0e4cb5ccef37f8d3ca7aa335f", + "timing": "Per-call wall time, CUDA synchronized before and after each call; alternating implementation order", + "results": [ + { + "case": "univariate", + "batch": 1, + "targets": 1, + "context": 128, + "horizon": 96, + "median_ms": { + "sdm": 9.248486021533608, + "upstream": 9.086007601581514 + }, + "p25_p75_ms": { + "sdm": [ + 9.17865033261478, + 9.280570084229112 + ], + "upstream": [ + 9.03238810133189, + 9.148046956397593 + ] + }, + "samples_ms": { + "sdm": [ + 9.297325974330306, + 9.27718193270266, + 9.181966073811054, + 9.338669944554567, + 9.207551134750247, + 9.288318920880556, + 9.235118050128222, + 9.3220139387995, + 9.184845956042409, + 9.28772590123117, + 9.202189976349473, + 9.274621959775686, + 9.146078955382109, + 9.278318146243691, + 9.185262955725193, + 9.304029867053032, + 9.161982918158174, + 9.255565935745835, + 9.140111040323973, + 9.278350044041872, + 9.168366901576519, + 9.268797934055328, + 9.168703109025955, + 9.241406107321382, + 9.199518943205476, + 9.266766952350736, + 9.16615896858275, + 9.28723020479083, + 9.143598843365908, + 9.27454186603427 + ], + "upstream": [ + 9.213807061314583, + 9.052592096850276, + 9.14259091950953, + 9.02188802137971, + 9.146254975348711, + 9.033088106662035, + 9.123407071456313, + 9.059280157089233, + 9.126031072810292, + 9.002608945593238, + 9.216655045747757, + 9.025776060298085, + 9.112735046073794, + 9.037695825099945, + 9.146366966888309, + 9.009199915453792, + 9.135551052168012, + 9.030288085341454, + 9.130894904956222, + 9.050735970959067, + 9.153086924925447, + 9.01137595064938, + 9.177310857921839, + 9.027295978739858, + 9.167599026113749, + 9.044975973665714, + 9.162415051832795, + 9.055423084646463, + 9.15487902238965, + 9.04475199058652 + ] + }, + "upstream_over_sdm": 0.9824318899792042 + }, + { + "case": "univariate", + "batch": 1, + "targets": 1, + "context": 1024, + "horizon": 96, + "median_ms": { + "sdm": 10.031183948740363, + "upstream": 9.844193933531642 + }, + "p25_p75_ms": { + "sdm": [ + 9.954813634976745, + 10.06162480916828 + ], + "upstream": [ + 9.791534976102412, + 9.891232941299677 + ] + }, + "samples_ms": { + "sdm": [ + 10.084616020321846, + 10.085687972605228, + 9.965016972273588, + 10.059609077870846, + 9.961609030142426, + 10.113383876159787, + 9.950666222721338, + 10.053895879536867, + 9.946457808837295, + 10.038327891379595, + 9.963738033547997, + 10.027304058894515, + 10.002905037254095, + 10.084087960422039, + 9.97639400884509, + 10.076679987832904, + 9.934121975675225, + 10.06767200306058, + 9.924746118485928, + 10.04020101390779, + 9.955529822036624, + 10.08711219765246, + 9.952665073797107, + 10.035063838586211, + 9.957114001736045, + 10.050377110019326, + 9.92228090763092, + 10.035288985818624, + 9.91602591238916, + 10.04956103861332 + ], + "upstream": [ + 9.866057895123959, + 9.777611121535301, + 9.871354093775153, + 9.805290028452873, + 9.922154014930129, + 9.79789812117815, + 9.919144911691546, + 9.815819095820189, + 9.92314494214952, + 9.787706891074777, + 9.881048928946257, + 9.815770899876952, + 9.876154130324721, + 9.75348288193345, + 9.889112785458565, + 9.802763117477298, + 9.910921100527048, + 9.792811004444957, + 9.911464992910624, + 9.787147166207433, + 9.947720915079117, + 9.798139100894332, + 9.888089960440993, + 9.78057086467743, + 9.890232933685184, + 9.773835074156523, + 9.89021803252399, + 9.822329971939325, + 9.894232964143157, + 9.76561103016138 + ] + }, + "upstream_over_sdm": 0.9813591280785752 + }, + { + "case": "multivariate", + "batch": 1, + "targets": 3, + "context": 1024, + "horizon": 96, + "median_ms": { + "sdm": 10.562653536908329, + "upstream": 10.318222921341658 + }, + "p25_p75_ms": { + "sdm": [ + 10.4778892127797, + 10.60572499409318 + ], + "upstream": [ + 10.222919285297394, + 10.348035080824047 + ] + }, + "samples_ms": { + "sdm": [ + 10.619108099490404, + 10.630645090714097, + 10.493444977328181, + 10.563493007794023, + 10.444054147228599, + 10.625236900523305, + 10.442086029797792, + 10.567429009824991, + 10.449477937072515, + 10.615220060572028, + 10.490837972611189, + 10.60738111846149, + 10.521892923861742, + 10.605172952637076, + 10.46554115600884, + 10.596373118460178, + 10.451462119817734, + 10.561814066022635, + 10.471781948581338, + 10.603205068036914, + 10.487797902897, + 10.64733206294477, + 10.505909100174904, + 10.566293029114604, + 10.505061829462647, + 10.616020997986197, + 10.469733970239758, + 10.576516855508089, + 10.479924967512488, + 10.603348957374692 + ], + "upstream": [ + 10.339798871427774, + 10.21970296278596, + 10.316886939108372, + 10.229190811514854, + 10.319558903574944, + 10.222967015579343, + 10.360389947891235, + 10.549044935032725, + 10.376389836892486, + 10.219848016276956, + 10.345398914068937, + 10.23067906498909, + 10.335878934711218, + 10.219767922535539, + 10.347095085307956, + 10.216550901532173, + 10.335270082578063, + 10.279992129653692, + 10.324694914743304, + 10.196136077865958, + 10.351925855502486, + 10.222776094451547, + 10.350855067372322, + 10.24859887547791, + 10.328406002372503, + 10.254471097141504, + 10.363718029111624, + 10.209175990894437, + 10.352998971939087, + 10.23144694045186 + ] + }, + "upstream_over_sdm": 0.9768589763251654 + }, + { + "case": "numeric", + "batch": 1, + "targets": 1, + "context": 1024, + "horizon": 96, + "median_ms": { + "sdm": 12.056970968842506, + "upstream": 11.474702041596174 + }, + "p25_p75_ms": { + "sdm": [ + 11.974626977462322, + 12.100886087864637 + ], + "upstream": [ + 11.381447373423725, + 11.495898070279509 + ] + }, + "samples_ms": { + "sdm": [ + 12.15472910553217, + 12.108250055462122, + 11.975466972216964, + 12.084122048690915, + 11.945738922804594, + 12.079898035153747, + 11.985786957666278, + 12.129082111641765, + 12.055563041940331, + 12.13936903513968, + 11.972106993198395, + 12.070953845977783, + 11.995146982371807, + 12.079034000635147, + 11.966651072725654, + 12.037466978654265, + 11.983274947851896, + 12.058378895744681, + 11.944635072723031, + 12.10073009133339, + 11.958315037190914, + 12.089337920770049, + 11.983867036178708, + 12.109690112993121, + 11.95301883853972, + 12.101354077458382, + 12.004635063931346, + 12.106857961043715, + 11.942267883569002, + 12.088218005374074 + ], + "upstream": [ + 11.488302843645215, + 11.360174976289272, + 11.508413823321462, + 11.388176120817661, + 11.525037931278348, + 11.344191851094365, + 11.475438019260764, + 11.403790907934308, + 11.491726152598858, + 11.35590416379273, + 11.491038836538792, + 11.385214980691671, + 12.017147149890661, + 11.387423146516085, + 11.486207135021687, + 11.513869976624846, + 11.488605989143252, + 11.3729911390692, + 11.478045955300331, + 11.36590400710702, + 11.486781993880868, + 11.351632187142968, + 11.473966063931584, + 11.374464025720954, + 11.516813887283206, + 11.383775155991316, + 11.545598041266203, + 11.404030956327915, + 11.533710174262524, + 11.392703047022223 + ] + }, + "upstream_over_sdm": 0.9517068649538076 + }, + { + "case": "categorical", + "batch": 1, + "targets": 1, + "context": 1024, + "horizon": 96, + "median_ms": { + "sdm": 14.91776597686112, + "upstream": 11.896683485247195 + }, + "p25_p75_ms": { + "sdm": [ + 14.794546994380653, + 14.958197949454188 + ], + "upstream": [ + 11.811736156232655, + 11.984022916294634 + ] + }, + "samples_ms": { + "sdm": [ + 14.959926018491387, + 14.979509869590402, + 14.773702947422862, + 14.919094042852521, + 14.790454879403114, + 15.025766100734472, + 14.798406977206469, + 15.013189986348152, + 14.829894062131643, + 14.925622148439288, + 14.780007069930434, + 14.957621926441789, + 14.82211914844811, + 14.927814016118646, + 14.807575149461627, + 15.010918024927378, + 14.757063006982207, + 14.963157940655947, + 14.905990101397038, + 14.936005929484963, + 14.790054876357317, + 14.916437910869718, + 14.822231139987707, + 14.931958867236972, + 14.761031139642, + 14.954373938962817, + 14.74954397417605, + 14.922661939635873, + 14.795911032706499, + 14.97079012915492 + ], + "upstream": [ + 11.957210954278708, + 11.781228007748723, + 12.040410889312625, + 11.81250810623169, + 11.937530944123864, + 11.824764078482985, + 12.020057998597622, + 11.836411897093058, + 11.962379794567823, + 11.777980020269752, + 11.98833086527884, + 11.776748113334179, + 11.982586933299899, + 11.834236094728112, + 11.969371000304818, + 11.810380034148693, + 12.004938907921314, + 11.812188196927309, + 11.939306976273656, + 11.855052085593343, + 12.00135494582355, + 11.797867948189378, + 11.955899186432362, + 11.80577208288014, + 11.95151498541236, + 11.774236103519797, + 12.026121839880943, + 11.840892024338245, + 12.154922122135758, + 11.855836026370525 + ] + }, + "upstream_over_sdm": 0.797484254927989 + }, + { + "case": "long_horizon", + "batch": 1, + "targets": 1, + "context": 79, + "horizon": 1030, + "median_ms": { + "sdm": 23.541401023976505, + "upstream": 23.653536452911794 + }, + "p25_p75_ms": { + "sdm": [ + 23.424018116202205, + 23.573976999614388 + ], + "upstream": [ + 23.558445158414543, + 23.73058773810044 + ] + }, + "samples_ms": { + "sdm": [ + 23.61173601821065, + 23.56804092414677, + 23.36952998302877, + 23.54908105917275, + 23.532121907919645, + 23.630055831745267, + 23.412794107571244, + 23.663287982344627, + 23.450265871360898, + 23.545193020254374, + 23.375466000288725, + 23.537609027698636, + 23.438937962055206, + 23.58682407066226, + 23.384073982015252, + 23.574408842250705, + 23.42740213498473, + 23.676472017541528, + 23.38964189402759, + 23.591272998601198, + 23.429673863574862, + 23.565097013488412, + 23.476761067286134, + 23.566728923469782, + 23.411146132275462, + 23.553977021947503, + 23.45932088792324, + 23.5626008361578, + 23.413866059854627, + 23.57383305206895 + ], + "upstream": [ + 23.716296069324017, + 23.644503904506564, + 23.732743924483657, + 23.545081028714776, + 23.732471046969295, + 23.56932102702558, + 23.759831907227635, + 23.58583197928965, + 23.717367788776755, + 23.540536873042583, + 23.729959968477488, + 23.579096887260675, + 23.70728785172105, + 23.558809189125896, + 23.662569001317024, + 23.582744877785444, + 23.761958815157413, + 23.559560999274254, + 23.68732006289065, + 23.53330492042005, + 23.71304715052247, + 23.60141696408391, + 23.72499997727573, + 23.54863309301436, + 23.790406994521618, + 23.557353066280484, + 23.739191936329007, + 23.517576977610588, + 23.742598947137594, + 23.53461692109704 + ] + }, + "upstream_over_sdm": 1.0047633286065294 + } + ] +} diff --git a/docs/source/api/models.rst b/docs/source/api/models.rst index 9a7b49bc25..947a0965eb 100644 --- a/docs/source/api/models.rst +++ b/docs/source/api/models.rst @@ -56,6 +56,13 @@ Overview - `Apache-2.0 `__ - `TimesFM Non-Commercial v1.0 `__ + * - | :class:`~sdm.models.Timestral` + | (`Source `__) + - — + - 5M–209M + - `Apache-2.0 `__ + - `Modified MIT `__ + Model API --------- diff --git a/examples/timeseries/timestral_categorical.py b/examples/timeseries/timestral_categorical.py new file mode 100644 index 0000000000..50137d53a3 --- /dev/null +++ b/examples/timeseries/timestral_categorical.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pandas as pd +import torch + +import sdm + +# Forecast daily sales conditioned on future event labels and past footfall. +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +time = torch.arange(96).float() +events = ["weekend" if i % 7 >= 5 else "weekday" for i in range(110)] +context = sdm.TableTensor.from_pandas( + pd.DataFrame( + { + "sales": 20 + 3 * (time / 7).sin(), + "footfall": 100 + time, + "event": pd.Categorical(events[:96]), + } + ), + stypes={ + "sales": "numerical", + "footfall": "numerical", + "event": "categorical", + }, + device=device, +) +query = sdm.TableTensor.from_pandas( + pd.DataFrame({"event": pd.Categorical(events[96:])}), + stypes={"event": "categorical"}, + device=device, +) +model = sdm.models.Timestral(model_size="5M", device=device) +forecast = model( + x_context=context.drop_columns("sales"), + y_context=context["sales"], + x_query=query, +) +print(forecast["sales__q50"].numerical.squeeze(-1).cpu()) diff --git a/pyproject.toml b/pyproject.toml index ee24ad0381..9ee0edc177 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ license-files = [ "sdm/models/tabfm/LICENSE", "sdm/models/timesfm3/LICENSE", "sdm/models/timestral/LICENSE", + "sdm/models/timestral/WEIGHTS_LICENSE", "sdm/models/kumo/relational/NOTICE", "sdm/models/kumo/timeseries/anomaly/NOTICE", "third_party/contributor-covenant/LICENSE", diff --git a/sdm/models/__init__.py b/sdm/models/__init__.py index 30e914f36c..50abecf58c 100644 --- a/sdm/models/__init__.py +++ b/sdm/models/__init__.py @@ -8,6 +8,7 @@ "sdm.models.tabfm", "sdm.models.kumo", "sdm.models.timesfm3", + "sdm.models.timestral", } from sdm.models.base import Callback, Model, ICLModel, ECOC @@ -20,6 +21,7 @@ KumoAnomaly, ) from sdm.models.timesfm3 import TimesFM3 +from sdm.models.timestral import Timestral __all__ = [ @@ -34,4 +36,5 @@ "KumoForecast", "KumoAnomaly", "TimesFM3", + "Timestral", ] diff --git a/sdm/models/timestral/WEIGHTS_LICENSE b/sdm/models/timestral/WEIGHTS_LICENSE new file mode 100644 index 0000000000..b6ae15066a --- /dev/null +++ b/sdm/models/timestral/WEIGHTS_LICENSE @@ -0,0 +1,10 @@ +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/sdm/models/timestral/__init__.py b/sdm/models/timestral/__init__.py index 94822d579d..93362a1348 100644 --- a/sdm/models/timestral/__init__.py +++ b/sdm/models/timestral/__init__.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Timestral.""" +"""Timestral multivariate forecasting.""" -__all__ = [] +from sdm.models.timestral.model import Timestral + +__all__ = ["Timestral"] diff --git a/sdm/models/timestral/config.py b/sdm/models/timestral/config.py new file mode 100644 index 0000000000..6b36eea9d0 --- /dev/null +++ b/sdm/models/timestral/config.py @@ -0,0 +1,73 @@ +# Portions derived from mistralai/timestral (Apache-2.0). +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass + +LAYER_KIND_AXES: dict[str, tuple[bool, ...]] = { + "time": (True,), + "variate": (False,), + "time_variate": (True, False), + "variate_time": (False, True), +} + + +@dataclass(frozen=True) +class _TimestralConfig: + patch_size: int + d_model: int + num_heads: int + mlp_hidden_dim: int + block_schedule: tuple[str, ...] + quantile_knots: tuple[float, ...] + block_size: int + context_length: int + rope_theta: float = 10000.0 + rope_partial_factor: float = 0.5 + norm_eps: float = 1e-06 + attn_bias: bool = False + mlp_bias: bool = False + time_attn_gate: str = "sdpa_sigmoid" + variate_attn_gate: str = "sdpa_sigmoid" + minimum_scale: float = 1e-06 + use_asinh: bool = True + stabilize_with_global: bool = False + covariate_enabled: bool = True + # Upstream execution hint; SDM uses the equivalent full attention mask. + split_context_horizon_attn: bool = True + attend_missing_patches: bool = True + quantile_width_factor: float = 2.0 + + @property + def column_every(self) -> int: + interval = next( + ( + i + 1 + for i, kind in enumerate(self.block_schedule) + if kind == "time_variate" + ), + len(self.block_schedule) + 1, + ) + expected = tuple( + "time_variate" if (i + 1) % interval == 0 else "time" + for i in range(len(self.block_schedule)) + ) + if tuple(self.block_schedule) != expected: + raise ValueError("Unsupported Timestral block schedule") + return interval + + @property + def head_dim(self) -> int: + return self.d_model // self.num_heads + + @property + def rope_dim(self) -> int: + return int(self.head_dim * self.rope_partial_factor) + + @property + def num_quantiles(self) -> int: + return len(self.quantile_knots) + + @property + def median_index(self) -> int: + return self.quantile_knots.index(0.5) diff --git a/sdm/models/timestral/decode.py b/sdm/models/timestral/decode.py new file mode 100644 index 0000000000..79d217d441 --- /dev/null +++ b/sdm/models/timestral/decode.py @@ -0,0 +1,329 @@ +# Portions derived from mistralai/timestral (Apache-2.0). +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import math + +import torch +from torch import Tensor + +from sdm.models.timestral.network import ( + ROLE_KNOWN_COV, + ROLE_PAST_COV, + ROLE_TARGET, + _TimestralModel, +) +from sdm.models.timestral.scaler import ( + _advancing_block_frames, + _clamp_nonfinite, +) + + +def _ffill_with_mask(inputs: Tensor, mask: Tensor) -> Tensor: + idx = torch.arange(inputs.shape[-1], device=inputs.device).expand_as( + inputs + ) + idx = torch.where(mask, idx, torch.zeros_like(idx)) + idx = idx.cummax(dim=-1).values + return torch.nan_to_num(torch.gather(inputs, dim=-1, index=idx), nan=0.0) + + +def _pad_to_patch( + inputs: Tensor, mask: Tensor, patch_size: int +) -> tuple[Tensor, Tensor]: + ctx_len = inputs.shape[-1] + if ctx_len % patch_size != 0: + pad_len = patch_size - ctx_len % patch_size + inputs = torch.nn.functional.pad(inputs, (pad_len, 0), value=0.0) + mask = torch.nn.functional.pad(mask, (pad_len, 0), value=False) + return (inputs, mask) + + +def _replace_extreme_values( + x: Tensor, threshold: float, replacement: float = 0.0 +) -> Tensor: + return torch.where( + torch.isfinite(x) & (x.abs() < threshold), + x, + torch.full_like(x, replacement), + ) + + +class _Decoder: + def __init__( + self, + model: _TimestralModel, + max_context: int | None = None, + per_shot_horizon_patches: int = 32, + quantile_real_cap_k: float = 10000.0, + extreme_value_threshold: float = 1e30, + ) -> None: + self.model = model + self.config = model.config + self.max_context = ( + max_context + if max_context is not None + else model.config.context_length + ) + self.per_shot_horizon_patches = min( + per_shot_horizon_patches, model.config.block_size + ) + self.quantile_real_cap_k = quantile_real_cap_k + self.extreme_value_threshold = extreme_value_threshold + self.median_index = model.config.median_index + + @property + def patch_size(self) -> int: + return self.config.patch_size + + def _forecast_raw( + self, + inputs: Tensor, + observation_mask: Tensor, + prediction_length: int, + past_covariates: Tensor | None = None, + known_covariates: Tensor | None = None, + ) -> Tensor: + rows, mask, role_ids, known_future, known_future_mask, n_tgt = ( + self._assemble_rows( + inputs=inputs, + observation_mask=observation_mask, + prediction_length=prediction_length, + past_covariates=past_covariates, + known_covariates=known_covariates, + ) + ) + quantiles = self._block_ar_decode( + inputs=rows, + observation_mask=mask, + prediction_length=prediction_length, + role_ids=role_ids, + known_future=known_future, + known_future_mask=known_future_mask, + ) + quantiles = quantiles[:, :n_tgt] + if self.quantile_real_cap_k > 0: + quantiles = self._cap_quantiles( + quantiles=quantiles, + inputs=rows[:, :n_tgt], + observation_mask=mask[:, :n_tgt], + ) + return quantiles + + def _cap_quantiles( + self, quantiles: Tensor, inputs: Tensor, observation_mask: Tensor + ) -> Tensor: + not_obs = ~observation_mask + ctx_max = torch.nan_to_num( + inputs.masked_fill(not_obs, float("-inf")).amax(-1, keepdim=True), + neginf=0.0, + ) + ctx_min = torch.nan_to_num( + inputs.masked_fill(not_obs, float("inf")).amin(-1, keepdim=True), + posinf=0.0, + ) + _, scale, _ = self.model.norm.running_stats( + x=inputs.unsqueeze(-2), + mask=not_obs.unsqueeze(-2), + correction=1, + ) + scale = scale.squeeze(-2).clamp_min(self.config.minimum_scale) + anchor_scale = scale.to(inputs.dtype) * self.quantile_real_cap_k + return quantiles.clamp( + (ctx_min - anchor_scale).unsqueeze(-1), + (ctx_max + anchor_scale).unsqueeze(-1), + ) + + def _assemble_rows( + self, + inputs: Tensor, + observation_mask: Tensor, + prediction_length: int, + past_covariates: Tensor | None, + known_covariates: Tensor | None, + ) -> tuple[ + Tensor, Tensor, Tensor | None, Tensor | None, Tensor | None, int + ]: + ps = self.patch_size + device = inputs.device + b, n_tgt, _ctx0 = inputs.shape + horizon = math.ceil(prediction_length / ps) * ps + ctx_rows = [_ffill_with_mask(inputs, observation_mask)] + ctx_masks = [observation_mask] + roles = [ROLE_TARGET] * n_tgt + n_past = 0 + if past_covariates is not None: + n_past = past_covariates.shape[1] + ctx_masks.append(~torch.isnan(past_covariates)) + ctx_rows.append(torch.nan_to_num(past_covariates, nan=0.0)) + roles += [ROLE_PAST_COV] * n_past + n_known = 0 + known_future_vals: Tensor | None = None + known_future_obs: Tensor | None = None + if known_covariates is not None: + n_known = known_covariates.shape[1] + kc_ctx, kc_fut = ( + known_covariates[..., :-prediction_length], + known_covariates[..., -prediction_length:], + ) + ctx_masks.append(~torch.isnan(kc_ctx)) + ctx_rows.append(torch.nan_to_num(kc_ctx, nan=0.0)) + roles += [ROLE_KNOWN_COV] * n_known + known_future_obs = ~torch.isnan(kc_fut) + known_future_vals = torch.nan_to_num(kc_fut, nan=0.0) + rows = torch.cat(ctx_rows, dim=1) + mask = torch.cat(ctx_masks, dim=1) + if rows.shape[-1] > self.max_context: + rows = rows[:, :, -self.max_context :] + mask = mask[:, :, -self.max_context :] + rows, mask = _pad_to_patch(rows, mask, ps) + v_all = rows.shape[1] + if n_past == 0 and n_known == 0: + return (rows, mask, None, None, None, n_tgt) + role_ids = torch.tensor(roles, dtype=torch.long, device=device).repeat( + b + ) + known_future: Tensor | None = None + known_future_mask: Tensor | None = None + if known_future_vals is not None and known_future_obs is not None: + k0 = n_tgt + n_past + plen = known_future_vals.shape[-1] + kf = torch.zeros( + b, v_all, horizon, device=device, dtype=rows.dtype + ) + kfm = torch.zeros( + b, v_all, horizon, device=device, dtype=torch.bool + ) + kf[:, k0 : k0 + n_known, :plen] = known_future_vals + kfm[:, k0 : k0 + n_known, :plen] = known_future_obs + known_future = kf.reshape(b * v_all, horizon) + known_future_mask = kfm.reshape(b * v_all, horizon) + return (rows, mask, role_ids, known_future, known_future_mask, n_tgt) + + def _block_ar_decode( + self, + inputs: Tensor, + observation_mask: Tensor, + prediction_length: int, + role_ids: Tensor | None, + known_future: Tensor | None, + known_future_mask: Tensor | None, + ) -> Tensor: + ps = self.patch_size + b, v, ctx_len = inputs.shape + context = inputs.reshape(b * v, ctx_len) + mask = observation_mask.reshape(b * v, ctx_len) + group_ids = torch.arange(b, device=inputs.device).repeat_interleave(v) + known_rows = past_rows = target_rows = None + if role_ids is not None: + known_rows = role_ids == ROLE_KNOWN_COV + past_rows = (role_ids != ROLE_TARGET) & ~known_rows + target_rows = role_ids == ROLE_TARGET + blocks: list[Tensor] = [] + remaining = math.ceil(prediction_length / ps) + offset = 0 + while remaining > 0: + block = min(self.per_shot_horizon_patches, remaining) + blen = block * ps + kf_block = ( + known_future[:, offset : offset + blen] + if known_future is not None + else None + ) + kfm_block = ( + known_future_mask[:, offset : offset + blen] + if known_future_mask is not None + else None + ) + quantiles = self.model.predict_block( + inputs=context, + observed=mask, + group_ids=group_ids, + horizon_patches=block, + num_variates=v, + role_ids=role_ids, + known_future=kf_block, + known_future_mask=kfm_block, + ) + quantiles = self._refine_block( + quantiles=quantiles, + context=context, + mask=mask, + target_rows=target_rows, + ) + blocks.append(quantiles) + remaining -= block + offset += blen + if remaining <= 0: + break + feedback = _replace_extreme_values( + x=quantiles[..., self.median_index], + threshold=self.extreme_value_threshold, + ) + feedback_mask = torch.ones_like(feedback, dtype=torch.bool) + if role_ids is not None: + assert known_rows is not None + assert past_rows is not None + if kf_block is not None: + assert kfm_block is not None + feedback = torch.where( + known_rows[:, None], kf_block, feedback + ) + feedback_mask = torch.where( + known_rows[:, None], kfm_block, feedback_mask + ) + feedback = torch.where( + past_rows[:, None], torch.zeros_like(feedback), feedback + ) + feedback_mask = feedback_mask & ~past_rows[:, None] + context = torch.cat([context, feedback], dim=-1) + mask = torch.cat([mask, feedback_mask], dim=-1) + if context.shape[-1] > self.max_context: + keep = max(ps, self.max_context // ps * ps) + context = context[..., -keep:] + mask = mask[..., -keep:] + packed = torch.cat(blocks, dim=-2)[..., :prediction_length, :] + return packed.reshape(b, v, prediction_length, -1) + + def _refine_block( + self, + quantiles: Tensor, + context: Tensor, + mask: Tensor, + target_rows: Tensor | None, + ) -> Tensor: + ps = self.patch_size + n_pred = quantiles.shape[1] // ps + if n_pred <= 1: + return quantiles + _, loc, scale = self.model.scale(context, mask) + loc0, scale0 = (loc[:, -1], scale[:, -1]) + median = ( + quantiles[..., self.median_index] + .unflatten(-1, (n_pred, ps)) + .to(torch.float64) + ) + sinh_median = (median - loc0[:, None, None]) / scale0[:, None, None] + loc_j, scale_j = _advancing_block_frames( + context=context, + observation_mask=mask, + sinh_median=sinh_median, + minimum_scale=self.config.minimum_scale, + ) + first = torch.zeros(loc_j.shape, dtype=torch.bool, device=loc_j.device) + first[:, 0] = True + loc_j = torch.where( + first, loc0[:, None], loc_j.to(quantiles.dtype) + ).repeat_interleave(ps, dim=-1) + scale_j = torch.where( + first, scale0[:, None], scale_j.to(quantiles.dtype) + ).repeat_interleave(ps, dim=-1) + sinh_norm = (quantiles - loc0[:, None, None]) / scale0[:, None, None] + refined = _clamp_nonfinite( + vals=sinh_norm * scale_j[..., None] + loc_j[..., None] + ) + if target_rows is not None: + refined = torch.where( + target_rows[:, None, None], refined, quantiles + ) + return refined diff --git a/sdm/models/timestral/encoding.py b/sdm/models/timestral/encoding.py new file mode 100644 index 0000000000..b3af4272f0 --- /dev/null +++ b/sdm/models/timestral/encoding.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from torch import Tensor + + +def _encode_categories( + target: Tensor, codes: Tensor, max_context: int, num_categories: int +) -> Tensor: + b, c, _length = codes.shape + t = target.size(-1) + valid = target.isfinite() + valid &= torch.arange(t, device=target.device) >= t - max_context + y = target.double().nan_to_num() + mean = (y * valid).sum(-1) / valid.sum(-1).clamp_min(1) + # Each column has an extra slot for missing values. + slots = num_categories + 1 + rows = torch.arange(b * c, device=codes.device).view(b, c, 1) + index = (rows * slots + codes.clamp_min(-1) + 1).expand_as(codes) + sums = y.new_zeros(b * c * slots) + counts = y.new_zeros(b * c * slots) + sums.scatter_add_( + 0, + index[..., :t].reshape(-1), + (y[:, None] * valid[:, None]).expand(b, c, t).reshape(-1), + ) + counts.scatter_add_( + 0, + index[..., :t].reshape(-1), + valid[:, None].expand(b, c, t).double().reshape(-1), + ) + prior = mean[:, None, None].expand(b, c, slots).reshape(-1) + out = ((prior + sums) / (1 + counts))[index].float() + out[..., :t] = out[..., :t].masked_fill( + ~valid[:, None] & (codes[..., :t] < 0), torch.nan + ) + return out diff --git a/sdm/models/timestral/masks.py b/sdm/models/timestral/masks.py new file mode 100644 index 0000000000..5c27ef8b9c --- /dev/null +++ b/sdm/models/timestral/masks.py @@ -0,0 +1,140 @@ +# Portions derived from mistralai/timestral (Apache-2.0). +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from torch import Tensor + + +def _compute_enc_key_validity( + obs_per_patch: Tensor, attend_missing: bool +) -> Tensor: + if not attend_missing: + return obs_per_patch + started = obs_per_patch.cummax(dim=-1).values + return obs_per_patch | started + + +def _enc_causal_mask(key_valid: Tensor) -> Tensor: + _, ll = key_valid.shape + device = key_valid.device + causal = torch.tril(torch.ones(ll, ll, dtype=torch.bool, device=device)) + eye = torch.eye(ll, dtype=torch.bool, device=device) + return causal[None] & key_valid[:, None, :] | eye[None] + + +def _build_forecast_time_mask( + context_len: int, horizon: int, key_valid: Tensor, device: torch.device +) -> Tensor: + rows = key_valid.shape[0] + enc = _enc_causal_mask(key_valid) + enc_dec = torch.zeros( + rows, context_len, horizon, dtype=torch.bool, device=device + ) + dec_enc = key_valid[:, None, :].expand(rows, horizon, context_len) + dec_dec = torch.tril( + torch.ones(horizon, horizon, dtype=torch.bool, device=device) + )[None].expand(rows, horizon, horizon) + top = torch.cat([enc, enc_dec], dim=2) + bottom = torch.cat([dec_enc, dec_dec], dim=2) + return torch.cat([top, bottom], dim=1).unsqueeze(1) + + +def _build_forecast_positions( + context_len: int, horizon: int, first_obs: Tensor +) -> Tensor: + p = torch.arange(context_len + horizon, device=first_obs.device) + return p[None, :] - first_obs[:, None] + + +def _compute_block_layout( + min_context: Tensor, n_patches: int, block_size: int +) -> tuple[Tensor, Tensor]: + device = min_context.device + p = torch.arange(n_patches, device=device)[None, :] + s = min_context[:, None] + active = p >= s + offset = (p - s).clamp(min=0) + anchor = s + block_size * (offset // block_size) + anchor = torch.where(active, anchor, torch.zeros_like(anchor)) + return (anchor, active) + + +def _build_two_stream_time_mask( + min_context: Tensor, + n_patches: int, + block_size: int, + key_valid: Tensor, + bidirectional_rows: Tensor | None, +) -> Tensor: + anchor, active = _compute_block_layout(min_context, n_patches, block_size) + rows, n = active.shape + device = active.device + causal = torch.tril(torch.ones(n, n, dtype=torch.bool, device=device)) + key_col = key_valid[:, None, :] + j = torch.arange(n, device=device).view(1, 1, n) + eye = torch.eye(n, dtype=torch.bool, device=device) + enc_enc = _enc_causal_mask(key_valid) + if bidirectional_rows is not None: + enc_bi = key_valid[:, None, :].expand(rows, n, n) | eye[None] + enc_enc = torch.where( + bidirectional_rows[:, None, None], enc_bi, enc_enc + ) + enc_dec = torch.zeros(rows, n, n, dtype=torch.bool, device=device) + dec_enc = active[:, :, None] & (j < anchor[:, :, None]) & key_col + same_block = anchor[:, :, None] == anchor[:, None, :] + dec_dec = ( + same_block & causal[None] & active[:, :, None] & active[:, None, :] + ) + dec_dec = dec_dec | eye[None] + top = torch.cat([enc_enc, enc_dec], dim=2) + bottom = torch.cat([dec_enc, dec_dec], dim=2) + return torch.cat([top, bottom], dim=1).unsqueeze(1) + + +def _variate_firewall( + r: int, + device: torch.device, + target_rows: Tensor | None, + bidirectional_rows: Tensor | None, +) -> Tensor | None: + if target_rows is None and bidirectional_rows is None: + return None + forbid = torch.zeros(r, r, dtype=torch.bool, device=device) + if target_rows is not None: + forbid = forbid | ~target_rows[:, None] & target_rows[None, :] + if bidirectional_rows is not None: + forbid = ( + forbid + | bidirectional_rows[:, None] & (~bidirectional_rows)[None, :] + ) + return forbid + + +def _build_two_stream_variate_mask( + group_ids: Tensor, + enc_key_valid: Tensor, + n_dec: int, + target_rows: Tensor | None = None, + bidirectional_rows: Tensor | None = None, +) -> Tensor | None: + r, _ = enc_key_valid.shape + if r == 1: + return None + device = enc_key_valid.device + row_eq = group_ids[:, None] == group_ids[None, :] + dec_valid = torch.ones(r, n_dec, dtype=torch.bool, device=device) + slot_valid = torch.cat([enc_key_valid, dec_valid], dim=1) + key_valid = slot_valid.permute(1, 0) + mask = row_eq[None, None] & key_valid[:, None, None, :] + forbid = _variate_firewall(r, device, target_rows, bidirectional_rows) + if forbid is not None: + mask = mask & ~forbid[None, None] + eye = torch.eye(r, dtype=torch.bool, device=device) + return mask | eye[None, None] + + +def _build_two_stream_positions(n_patches: int, first_obs: Tensor) -> Tensor: + p = torch.arange(n_patches, device=first_obs.device) + row = p[None, :] - first_obs[:, None] + return torch.cat([row, row], dim=-1) diff --git a/sdm/models/timestral/model.py b/sdm/models/timestral/model.py new file mode 100644 index 0000000000..d0619c95e3 --- /dev/null +++ b/sdm/models/timestral/model.py @@ -0,0 +1,351 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, ClassVar, Literal, cast + +import torch +from safetensors.torch import load_file +from torch import Tensor + +import sdm.processing as sp +from sdm import ( + Recipe, + RelatedTables, + Stype, + TableTensor, + Task, +) +from sdm.cache import Cache +from sdm.models._huggingface import download_checkpoint +from sdm.models.base import ICLModel +from sdm.models.timestral.config import LAYER_KIND_AXES, _TimestralConfig +from sdm.models.timestral.decode import _Decoder, _ffill_with_mask +from sdm.models.timestral.encoding import _encode_categories +from sdm.models.timestral.network import _TimestralModel + +_REVISIONS = { + "5M": "bdd20cfc15a3fea0e4cb5ccef37f8d3ca7aa335f", + "28M": "180698bbf5a50dda3516aa6419798437fb03a430", + "123M": "5a147287addda85f6acd8da3fb8858fe9af7e8f7", + "209M": "9874d252ab0aae16f7f7c999f98f98d6ad7d8365", +} +_LEVELS = (0.01, 0.05, *[i / 100 for i in range(10, 95, 5)], 0.95, 0.99) + + +class Timestral(ICLModel): + """Multivariate quantile forecasting with Timestral. + + See `Timestral `_. + + Rows are ordered time steps. Target columns are forecast jointly. Features + present in the query are known-future covariates; other features are + historical covariates. Categorical features use target encoding for one + target and category indices for multiple targets. Missing categories + form their own category. Encoding learns from the context only. The + default recipe aligns category dictionaries while preserving unseen values. + + Returns 21 numerical columns per target, named ``__q1`` through + ``__q99`` at the checkpoint's quantile levels. Missing numerical + values are allowed. The most recent 32,768 context steps are used; horizons + beyond 1,024 steps use median feedback. The default forecast includes sign + averaging and clipping to zero for nonnegative target histories. + + Derived from upstream revision + ``a534bdbd244d64f3d9d743f46c7bf6fc330ed4ff``. Pretrained weights use + `Modified MIT license `_ + terms, including a company revenue restriction. + + Args: + model_size: Released checkpoint size. + pretrained: Whether to load pretrained weights. + accept_license: Accept the weights license without a prompt. + device: The device. + """ + + supported_feature_stypes: ClassVar[frozenset[Stype]] = frozenset( + {Stype.numerical, Stype.categorical} + ) + supported_target_stypes: ClassVar[frozenset[Stype]] = frozenset( + {Stype.numerical} + ) + supports_multi_target: ClassVar[bool] = True + supports_context_only_columns: ClassVar[bool] = True + supports_related_tables: ClassVar[bool] = False + + def __init__( + self, + model_size: Literal["5M", "28M", "123M", "209M"] = "5M", + pretrained: bool = True, + accept_license: bool = False, + device: torch.device | str | None = None, + ) -> None: + super().__init__(task=Task.regression) + repo = f"mistralai/Timestral-{model_size}" + revision = _REVISIONS[model_size] + if pretrained: + config_path = download_checkpoint( + repo_id=repo, + filename="config.json", + revision=revision, + ) + config = _TimestralConfig( + **json.loads(Path(config_path).read_text()) + ) + else: + # Released dimensions for offline initialization. + dims = { + "5M": (192, 3, 512, 8), + "28M": (384, 6, 1024, 12), + "123M": (640, 10, 1712, 20), + "209M": (768, 12, 2048, 24), + } + width, heads, hidden, depth = dims[model_size] + config = _TimestralConfig( + patch_size=32, + d_model=width, + num_heads=heads, + mlp_hidden_dim=hidden, + block_schedule=tuple( + "time_variate" if i % 4 == 3 else "time" + for i in range(depth) + ), + quantile_knots=_LEVELS, + block_size=32, + context_length=32768, + quantile_width_factor=3.0, + ) + with torch.device(device or torch.get_default_device()): + self.model = _TimestralModel(config) + if pretrained: + path = download_checkpoint( + repo_id=repo, + filename="model.safetensors", + revision=revision, + license_prompt=None + if accept_license + else ( + "Timestral weights use Modified MIT terms, including a " + "company revenue restriction. Review " + f"https://huggingface.co/{repo}/blob/{revision}/LICENSE" + ), + ) + state = load_file( + path, device=str(next(self.model.parameters()).device) + ) + state = _remap_checkpoint(state, config) + for key, value in self.model.state_dict().items(): + if key.endswith(".inv_freq"): + state[key] = value.detach() + self.model.load_state_dict(state, strict=True) + self.eval() + + @classmethod + def default_recipe(cls) -> Recipe: # noqa: D102 + return Recipe( + features=sp.AlignCategories(min_frequency=0, unknown="append"), + output=sp.AverageEstimators(), + ) + + def _forward( + self, + x_context: TableTensor | None, + y_context: TableTensor | None, + x_query: TableTensor | None, + x_context_only: TableTensor | None, + related_context_tables: RelatedTables[TableTensor] | None, + related_query_tables: RelatedTables[TableTensor] | None, + cache: Cache | None, + generator: torch.Generator | None, + symmetric_averaging: bool = True, + clip_nonnegative: bool = True, + **kwargs: Any, + ) -> TableTensor: + if cache is not None and not cache.is_recording: + x_context = cast(TableTensor, cache["context"]) + y_context = cast(TableTensor, cache["target"]) + x_context_only = cast(TableTensor | None, cache["past"]) + assert x_context is not None + assert y_context is not None + if cache is not None and cache.is_recording: + cache["context"] = x_context + cache["target"] = y_context + cache["past"] = x_context_only + names = y_context.columns[Stype.numerical] + columns = [ + f"{name}__q{round(q * 100)}" + for name in names + for q in self.model.config.quantile_knots + ] + if x_query is None or x_query.size(-2) == 0: + return TableTensor.from_tensor( + y_context.numerical.new_empty( + (*y_context.shape[:-2], 0, len(columns)) + ), + columns=columns, + ) + y = y_context.numerical.transpose(-2, -1) + batch = y.shape[:-2] + y = y.reshape(-1, *y.shape[-2:]).float() + observed = y.isfinite() + y = y.masked_fill(~observed, 0) + decoder = _Decoder(self.model) + past = self._covariates( + context=x_context_only, + query=None, + target=y, + observed=observed, + max_context=decoder.max_context, + ) + known = self._covariates( + context=x_context, + query=x_query, + target=y, + observed=observed, + max_context=decoder.max_context, + ) + horizon = x_query.size(-2) + out = decoder._forecast_raw(y, observed, horizon, past, known) + if symmetric_averaging: + mirror = decoder._forecast_raw( + inputs=-y, + observation_mask=observed, + prediction_length=horizon, + past_covariates=-past if past is not None else None, + known_covariates=-known if known is not None else None, + ) + out = (out - mirror.flip(-1)) * 0.5 + if clip_nonnegative: + eligible = observed.any(-1) & ((y >= 0) | ~observed).all(-1) + out = torch.where(eligible[..., None, None], out.clamp_min(0), out) + out = out.transpose(-3, -2).reshape(*batch, horizon, len(columns)) + return TableTensor.from_tensor(out, columns=columns) + + @staticmethod + def _covariates( + context: TableTensor | None, + query: TableTensor | None, + target: Tensor, + observed: Tensor, + max_context: int, + ) -> Tensor | None: + if context is None or context.size(-1) == 0: + return None + numeric = context.numerical + codes = context.categorical.code + num_categories = max( + (c.numel() for c in context.categorical.categories), default=0 + ) + if query is not None: + numeric = torch.cat([numeric, query.numerical], dim=-2) + codes = torch.cat([codes, query.categorical.code], dim=-2) + num_categories = max( + num_categories, + max( + (c.numel() for c in query.categorical.categories), + default=0, + ), + ) + length = numeric.size(-2) + numeric = ( + numeric.transpose(-2, -1) + .reshape(target.size(0), -1, length) + .float() + ) + codes = codes.transpose(-2, -1).reshape(target.size(0), -1, length) + if codes.size(1) == 0: + return numeric + if target.size(1) == 1: + started = observed[:, 0].cummax(-1).values + filled = _ffill_with_mask(target, observed)[:, 0].masked_fill( + ~started, torch.nan + ) + categorical = _encode_categories( + target=filled, + codes=codes, + max_context=max_context, + num_categories=num_categories, + ) + else: + categorical = codes.float().masked_fill(codes < 0, torch.nan) + return torch.cat([numeric, categorical], dim=1) + + +def _remap_checkpoint( + state: dict[str, Tensor], config: _TimestralConfig +) -> dict[str, Tensor]: + out = {} + heads = config.num_heads + time_attns = { + f"transformer.layers.{layer}.attns.{axis}" + for layer, kind in enumerate(config.block_schedule) + for axis, is_time in enumerate(LAYER_KIND_AXES[kind]) + if is_time + } + channels = torch.arange( + config.head_dim, + device=next(iter(state.values())).device, + ) + half, rotated_half = config.head_dim // 2, config.rope_dim // 2 + # Move the active pairs from the full head halves into a contiguous prefix. + permutation = torch.cat( + [ + channels[:rotated_half], + channels[half : half + rotated_half], + channels[rotated_half:half], + channels[half + rotated_half :], + ] + ) + column_every = config.column_every + for key, value in state.items(): + key = ( + key.replace("hidden_layer", "mlp.0") + .replace("output_layer", "mlp.2") + .replace("residual_layer", "res") + ) + if ".wQKV." in key: + # Upstream interleaves heads within each head channel. + head_dim = value.size(0) // (3 * heads) + trailing = value.shape[1:] + qkv = value.reshape(3, head_dim, heads, *trailing).transpose(1, 2) + if key.rsplit(".", 2)[0] in time_attns: + qkv = torch.cat( + [ + qkv[:2].index_select(2, permutation), + qkv[2:], + ] + ) + value = qkv.reshape_as(value) + key = key.replace("wQKV", "qkv_lin") + if key.startswith("transformer.layers."): + parts = key.split(".") + layer = int(parts[2]) + if parts[3] in {"attns", "attn_norms"}: + column = int(parts[4]) == 1 + module = "attn" if parts[3] == "attns" else "query_norm" + suffix = ".".join(parts[5:]) + else: + column = (layer + 1) % column_every == 0 + module = "mlp.0" if parts[3] == "mlp_norm" else "mlp.1" + suffix = ".".join(parts[4:]) + stack = "col_blocks" if column else "row_blocks" + index = layer // column_every if column else layer + key = f"transformer.{stack}.{index}.{module}.{suffix}" + key = key.replace(".wO.", ".out_lin.").replace("._gate.", ".gate.0.") + key = key.replace("transformer.out_norm.", "transformer.norm.") + if key.endswith(".scale"): + key = key[:-5] + "weight" + if ".query_norm." in key: + out[key.replace(".query_norm.", ".key_value_norm.")] = value + if ".mlp.1.w12." in key: + # Upstream stores the linear branch first, the SiLU branch second. + up, gate = value.chunk(2, dim=0) + out[key.replace("w12", "up_lin")] = up + out[key.replace("w12", "gate_lin")] = gate + else: + out[key.replace(".mlp.1.w3.", ".mlp.1.down_lin.")] = value + return out diff --git a/sdm/models/timestral/network.py b/sdm/models/timestral/network.py new file mode 100644 index 0000000000..c735eb34b3 --- /dev/null +++ b/sdm/models/timestral/network.py @@ -0,0 +1,326 @@ +# Portions derived from mistralai/timestral (Apache-2.0). +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from torch import Tensor, nn + +from sdm.models.timestral.block import MonotoneQuantiles, ResidualBlock +from sdm.models.timestral.config import _TimestralConfig +from sdm.models.timestral.icl import ICLBlock +from sdm.models.timestral.masks import ( + _build_forecast_positions, + _build_forecast_time_mask, + _build_two_stream_positions, + _build_two_stream_time_mask, + _build_two_stream_variate_mask, + _compute_enc_key_validity, +) +from sdm.models.timestral.scaler import ( + _apply_affine_asinh_masked, + _denormalize_quantiles, + _instance_norm_stats, +) +from sdm.nn import RevIN + +ROLE_TARGET = 0 +ROLE_PAST_COV = 1 +ROLE_KNOWN_COV = 2 + + +class _PatchRepresentation(nn.Module): + def __init__(self, patch_size: int, d_model: int) -> None: + super().__init__() + self.patch_size = patch_size + self.patch_proj = ResidualBlock( + in_channels=2 * patch_size, + hidden_channels=4 * d_model, + out_channels=d_model, + bias=True, + ) + self.forecast_query = nn.Parameter(torch.zeros(d_model)) + + def forward(self, scaled: Tensor, observed: Tensor) -> Tensor: + r, t = scaled.shape + n = t // self.patch_size + values = scaled.reshape(r, n, self.patch_size) + missing = (~observed).reshape(r, n, self.patch_size).to(scaled.dtype) + x = torch.cat([values, missing], dim=-1).to( + self.patch_proj.mlp[0].weight.dtype + ) + return self.patch_proj(x) + + def forecast_tokens(self, rows: int, n_patches: int) -> Tensor: + return self.forecast_query.view(1, 1, -1).expand(rows, n_patches, -1) + + +class _QuantileHead(nn.Module): + def __init__(self, config: _TimestralConfig) -> None: + super().__init__() + self.quantiles = MonotoneQuantiles( + config.quantile_knots, config.quantile_width_factor + ) + self.patch_size = config.patch_size + self.num_quantiles = config.num_quantiles + self.projection = ResidualBlock( + in_channels=config.d_model, + hidden_channels=4 * config.d_model, + out_channels=config.patch_size * config.num_quantiles, + bias=False, + ) + + def forward(self, x: Tensor) -> Tensor: + raw = self.projection(x) + *leading, n_patches, _ = raw.shape + preds = raw.reshape( + *leading, n_patches, self.patch_size, self.num_quantiles + ) + preds = preds.to(torch.promote_types(preds.dtype, torch.float32)) + preds = self.quantiles(preds) + return preds.reshape( + *leading, n_patches * self.patch_size, self.num_quantiles + ) + + +class _TimestralModel(nn.Module): + def __init__(self, config: _TimestralConfig) -> None: + super().__init__() + self.config = config + c = config + self.norm = RevIN(eps=0) + self.representation = _PatchRepresentation(c.patch_size, c.d_model) + self.transformer = ICLBlock( + channels=c.d_model, + num_layers=len(c.block_schedule), + column_every=c.column_every, + num_heads=c.num_heads, + hidden_channels=c.mlp_hidden_dim, + ) + self.head = _QuantileHead(c) + + def scale( + self, inputs: Tensor, observed: Tensor + ) -> tuple[Tensor, Tensor, Tensor]: + x = inputs.unflatten(-1, (-1, self.config.patch_size)) + mask = ~observed.unflatten(-1, (-1, self.config.patch_size)) + mean, scale, _ = self.norm.running_stats(x, mask, correction=1) + mean = mean.to(inputs.dtype) + scale = scale.clamp_min(self.config.minimum_scale).to(inputs.dtype) + + dtype = torch.promote_types(inputs.dtype, torch.float32) + out = self.norm( + x=x.to(dtype), + mean=mean, + scale=scale.to(dtype).clamp_min(torch.finfo(dtype).eps), + ).masked_fill(mask, 0) + if self.config.use_asinh: + out = out.asinh() + return ( + out.flatten(-2).to(inputs.dtype), + mean.expand_as(x).flatten(-2), + scale.expand_as(x).flatten(-2), + ) + + def _known_cov_scaling( + self, + inputs: Tensor, + observed: Tensor, + scaled: Tensor, + loc: Tensor, + scale: Tensor, + known_rows: Tensor, + ) -> tuple[Tensor, Tensor, Tensor]: + loc_c, scale_c, is_binary = _instance_norm_stats( + values=inputs, + observation_mask=observed, + minimum_scale=self.config.minimum_scale, + ) + known_scaled = _apply_affine_asinh_masked( + data=inputs, + loc=loc_c, + scale=scale_c, + observation_mask=observed, + use_asinh=False, + ) + if self.config.use_asinh: + work = torch.promote_types(known_scaled.dtype, torch.float32) + squashed = known_scaled.to(work).asinh().to(known_scaled.dtype) + known_scaled = torch.where( + observed & ~is_binary, squashed, known_scaled + ) + sel = known_rows[:, None] + return ( + torch.where(sel, known_scaled, scaled), + torch.where(sel, loc_c, loc), + torch.where(sel, scale_c, scale), + ) + + def _run_two_stream( + self, + x_enc: Tensor, + x_dec: Tensor, + time_mask: Tensor, + variate_mask: Tensor | None, + positions: Tensor, + ) -> Tensor: + x = torch.cat([x_enc, x_dec], dim=1) + h = self.transformer( + x=x, + row_mask=time_mask, + col_mask=variate_mask, + pos=positions, + ) + return self.head(h[:, x_enc.shape[1] :, :]) + + def predict_block( + self, + inputs: Tensor, + observed: Tensor, + group_ids: Tensor, + horizon_patches: int, + num_variates: int, + role_ids: Tensor | None = None, + known_future: Tensor | None = None, + known_future_mask: Tensor | None = None, + ) -> Tensor: + if role_ids is None or not self.config.covariate_enabled: + return self._predict_context_only( + inputs=inputs, + observed=observed, + group_ids=group_ids, + horizon_patches=horizon_patches, + num_variates=num_variates, + ) + r = inputs.shape[0] + horizon_len = horizon_patches * self.config.patch_size + if known_future is None: + known_future = torch.zeros( + r, horizon_len, device=inputs.device, dtype=inputs.dtype + ) + known_future_mask = torch.zeros( + r, horizon_len, device=inputs.device, dtype=torch.bool + ) + assert known_future_mask is not None + full_inputs = torch.cat([inputs, known_future], dim=-1) + full_observed = torch.cat([observed, known_future_mask], dim=-1) + return self._predict_window( + full_inputs=full_inputs, + full_observed=full_observed, + group_ids=group_ids, + horizon_patches=horizon_patches, + num_variates=num_variates, + role_ids=role_ids, + ) + + def _predict_context_only( + self, + inputs: Tensor, + observed: Tensor, + group_ids: Tensor, + horizon_patches: int, + num_variates: int, + ) -> Tensor: + c = self.config + ps = c.patch_size + rows, t = inputs.shape + ll, h = (t // ps, horizon_patches) + scaled, loc, scale = self.scale(inputs, observed) + x_enc = self.representation(scaled, observed) + x_dec = self.representation.forecast_tokens(rows, h) + obs_per_patch = observed.unflatten(-1, (ll, ps)).any(dim=-1) + key_valid = _compute_enc_key_validity( + obs_per_patch=obs_per_patch, + attend_missing=c.attend_missing_patches, + ) + time_mask = _build_forecast_time_mask(ll, h, key_valid, inputs.device) + variate_mask = _build_two_stream_variate_mask(group_ids, key_valid, h) + # Packed rows are contiguous [batch, variate, patch]. + first_obs = obs_per_patch.unflatten(0, (-1, num_variates)).any(1) + first_obs = first_obs.int().argmax(-1).repeat_interleave(num_variates) + positions = _build_forecast_positions(ll, h, first_obs) + preds = self._run_two_stream( + x_enc=x_enc, + x_dec=x_dec, + time_mask=time_mask, + variate_mask=variate_mask, + positions=positions, + ) + anchor_loc = loc.unflatten(-1, (ll, ps))[..., 0][:, ll - 1 : ll] + anchor_scale = scale.unflatten(-1, (ll, ps))[..., 0][:, ll - 1 : ll] + return _denormalize_quantiles( + preds=preds, + loc=anchor_loc.unsqueeze(-1), + scale=anchor_scale.unsqueeze(-1), + use_asinh=c.use_asinh, + ) + + def _predict_window( + self, + full_inputs: Tensor, + full_observed: Tensor, + group_ids: Tensor, + horizon_patches: int, + num_variates: int, + role_ids: Tensor, + ) -> Tensor: + c = self.config + ps = c.patch_size + rows, w = full_inputs.shape + h = horizon_patches + n = w // ps + ll = n - h + target_rows = role_ids == ROLE_TARGET + known_rows = role_ids == ROLE_KNOWN_COV + scaled, loc, scale = self.scale(full_inputs, full_observed) + scaled, loc, scale = self._known_cov_scaling( + inputs=full_inputs, + observed=full_observed, + scaled=scaled, + loc=loc, + scale=scale, + known_rows=known_rows, + ) + x_enc = self.representation(scaled, full_observed) + x_dec = self.representation.forecast_tokens(rows, n) + obs_per_patch = full_observed.unflatten(-1, (n, ps)).any(dim=-1) + key_valid = _compute_enc_key_validity( + obs_per_patch=obs_per_patch, + attend_missing=c.attend_missing_patches, + ) + s_rows = torch.full( + (rows,), ll, device=full_inputs.device, dtype=torch.long + ) + time_mask = _build_two_stream_time_mask( + min_context=s_rows, + n_patches=n, + block_size=c.block_size, + key_valid=key_valid, + bidirectional_rows=known_rows, + ) + variate_mask = _build_two_stream_variate_mask( + group_ids=group_ids, + enc_key_valid=key_valid, + n_dec=n, + target_rows=target_rows, + bidirectional_rows=known_rows, + ) + # Packed rows are contiguous [batch, variate, patch]. + first_obs = obs_per_patch.unflatten(0, (-1, num_variates)).any(1) + first_obs = first_obs.int().argmax(-1).repeat_interleave(num_variates) + positions = _build_two_stream_positions(n, first_obs) + preds = self._run_two_stream( + x_enc=x_enc, + x_dec=x_dec, + time_mask=time_mask, + variate_mask=variate_mask, + positions=positions, + ).unflatten(1, (n, ps)) + horizon_preds = preds[:, ll:n].reshape(rows, h * ps, preds.shape[-1]) + anchor_loc = loc.unflatten(-1, (n, ps))[..., 0][:, ll - 1 : ll] + anchor_scale = scale.unflatten(-1, (n, ps))[..., 0][:, ll - 1 : ll] + return _denormalize_quantiles( + preds=horizon_preds, + loc=anchor_loc.unsqueeze(-1), + scale=anchor_scale.unsqueeze(-1), + use_asinh=c.use_asinh, + ) diff --git a/sdm/models/timestral/scaler.py b/sdm/models/timestral/scaler.py new file mode 100644 index 0000000000..0e256e65bd --- /dev/null +++ b/sdm/models/timestral/scaler.py @@ -0,0 +1,113 @@ +# Portions derived from mistralai/timestral (Apache-2.0). +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import functools + +import torch +from torch import Tensor + +from sdm.nn import RevIN + + +def _fp32_io(*dtypes: torch.dtype) -> tuple[torch.dtype, torch.dtype]: + out_dtype = functools.reduce(torch.promote_types, dtypes) + return (out_dtype, torch.promote_types(out_dtype, torch.float32)) + + +def _apply_affine_asinh_masked( + data: Tensor, + loc: Tensor, + scale: Tensor, + observation_mask: Tensor, + use_asinh: bool, +) -> Tensor: + out_dtype, work = _fp32_io(data.dtype, loc.dtype, scale.dtype) + scaled = (data.to(work) - loc.to(work)) / scale.to(work).clamp_min( + torch.finfo(work).eps + ) + scaled = scaled.masked_fill(~observation_mask, 0.0) + if use_asinh: + scaled = torch.where(observation_mask, scaled.asinh(), scaled) + return scaled.to(out_dtype) + + +def _invert_affine_asinh( + x: Tensor, loc: Tensor, scale: Tensor, use_asinh: bool +) -> Tensor: + out_dtype, work = _fp32_io(x.dtype, loc.dtype, scale.dtype) + x = x.to(work) + if use_asinh: + x = x.sinh() + return (x * scale.to(work) + loc.to(work)).to(out_dtype) + + +def _clamp_nonfinite(vals: Tensor) -> Tensor: + clamped = torch.where( + vals == float("inf"), + torch.where(vals.isfinite(), vals, -float("inf")).amax( + dim=-1, keepdim=True + ), + torch.where( + vals == -float("inf"), + torch.where(vals.isfinite(), vals, float("inf")).amin( + dim=-1, keepdim=True + ), + vals, + ), + ) + return torch.where(clamped.isinf(), torch.zeros_like(clamped), clamped) + + +def _denormalize_quantiles( + preds: Tensor, loc: Tensor, scale: Tensor, use_asinh: bool +) -> Tensor: + return _clamp_nonfinite(_invert_affine_asinh(preds, loc, scale, use_asinh)) + + +def _advancing_block_frames( + context: Tensor, + observation_mask: Tensor, + sinh_median: Tensor, + minimum_scale: float, +) -> tuple[Tensor, Tensor]: + n_pred = sinh_median.size(-2) + state = RevIN.running_stats( + x=context.unsqueeze(-2), + mask=~observation_mask.unsqueeze(-2), + correction=1, + ) + offset = state[0][:, 0, :] + clip = torch.finfo(torch.float32).max ** 0.5 / 1000.0 + locs, scales = ([], []) + for h in range(n_pred): + mean = state[0][:, 0, 0] + scale = state[1][:, 0, 0].clamp_min(minimum_scale) + locs.append(mean) + scales.append(scale) + folded = ( + sinh_median[:, h] * scale[:, None] + mean[:, None] - offset + ).clamp(-clip, clip) + folded = folded + offset + state = RevIN.running_stats( + x=folded.unsqueeze(-2), + state=state, + correction=1, + ) + return (torch.stack(locs, dim=1), torch.stack(scales, dim=1)) + + +def _instance_norm_stats( + values: Tensor, observation_mask: Tensor, minimum_scale: float +) -> tuple[Tensor, Tensor, Tensor]: + orig_dtype = values.dtype + loc, scale = RevIN.stats(values, mask=~observation_mask, correction=0) + observed = observation_mask.any(dim=-1, keepdim=True) + scale = scale.where(observed, 0.0).clamp_min(minimum_scale) + is_zero_one = ((values == 0.0) | (values == 1.0) | ~observation_mask).all( + dim=-1, keepdim=True + ) + is_binary = is_zero_one & observed + loc = torch.where(is_binary, torch.zeros_like(loc), loc) + scale = torch.where(is_binary, torch.ones_like(scale), scale) + return (loc.to(orig_dtype), scale.to(orig_dtype), is_binary) diff --git a/sdm/processing/categorical/align.py b/sdm/processing/categorical/align.py index 3a4cbcf924..971384e1c4 100644 --- a/sdm/processing/categorical/align.py +++ b/sdm/processing/categorical/align.py @@ -25,22 +25,28 @@ class AlignCategories(EnsembleProcessor): - """Align categorical columns to categories observed during fitting. + """Align categorical columns to a fitted category vocabulary. A categorical column stores each value as an integer code indexing an ordered list of categories. Separately created tables can use different codes for the same value. Fitting learns the category list for each column, - and transforming remaps another table to use it. Missing values and - categories not retained during fitting receive code ``-1``. + and transforming remaps another table to use it. Missing values receive + code ``-1``. By default, categories absent from the fitted vocabulary + also receive code ``-1``. Args: sort_by: How to order fitted categories. - ``"code"`` keeps observed categories in original order. + ``"code"`` keeps retained categories in original order. ``"frequency"`` orders observed categories by descending frequency. ``"value"`` orders observed categories by ascending value. min_frequency: Minimum number of observations required to retain a category. Values of rarer categories receive code ``-1``. - Must be positive. + Zero retains every dictionary entry, including unobserved ones. + unknown: How to handle categories absent from the fitted vocabulary. + ``"missing"`` maps them to ``-1``. ``"append"`` appends them to + the returned vocabulary in input dictionary order, keeping them + distinct from missing values. It requires ``min_frequency=0``. + Transforming never updates the fitted vocabulary. >>> import pandas as pd >>> import sdm @@ -78,13 +84,28 @@ def __init__( sort_by: Literal["code", "frequency", "value"] = "code", *, min_frequency: int = 1, + unknown: Literal["missing", "append"] = "missing", ) -> None: super().__init__() - if min_frequency <= 0: - raise ValueError("min_frequency must be positive") + if min_frequency < 0: + raise ValueError("min_frequency must be nonnegative") + if unknown == "append" and min_frequency != 0: + raise ValueError( + "Appending unknown categories requires min_frequency=0" + ) self.sort_by = sort_by self.min_frequency = min_frequency + self.unknown = unknown self._categories: BufferList[BufferList[Tensor]] = BufferList() + self._table_ids: tuple[int, ...] = () + + def get_extra_state(self) -> tuple[int, ...]: + r""":meta private:""" # noqa: D415 + return self._table_ids + + def set_extra_state(self, state: object) -> None: + r""":meta private:""" # noqa: D415 + self._table_ids = cast(tuple[int, ...], state) def _fit_column( self, @@ -98,6 +119,10 @@ def _fit_column( aligned_codes = torch.full_like(codes, -1) if align_codes else None return (input_categories,) * batch_size, aligned_codes + if self.min_frequency == 0 and self.sort_by == "code": + categories = (input_categories,) * batch_size + return categories, codes if align_codes else None + mask = codes >= 0 indices = codes.clamp_min(0).long() # Count categories independently per batch: [B, N] -> [B, K]. @@ -228,6 +253,7 @@ def _fit( table, align_codes=False, ) + self._table_ids = tuple(range(len(fitted_categories))) self._categories = BufferList( BufferList(categories) for categories in fitted_categories ) @@ -239,6 +265,7 @@ def _fit_transform( generator: torch.Generator | None = None, ) -> TableTensor: fitted_categories, aligned_tables = self._fit_and_align(table) + self._table_ids = tuple(range(len(fitted_categories))) self._categories = BufferList( BufferList(categories) for categories in fitted_categories ) @@ -280,8 +307,9 @@ def _fit_ensemble( align_codes=False, ) fitted_categories.extend(group_categories) + self._table_ids = self._member_table_ids(ensemble_table) self._categories = BufferList( - BufferList(categories) for categories in fitted_categories + BufferList(fitted_categories[index]) for index in self._table_ids ) def _fit_transform_ensemble( @@ -302,8 +330,9 @@ def _fit_transform_ensemble( tables=aligned_tables, member_table_ids=member_table_ids, ) + self._table_ids = member_table_ids self._categories = BufferList( - BufferList(categories) for categories in fitted_categories + BufferList(fitted_categories[index]) for index in member_table_ids ) return output @@ -311,25 +340,34 @@ def _transform_ensemble( self, ensemble_table: EnsembleTable, ) -> EnsembleTable: - table_ids = self._member_table_ids(ensemble_table) - aligned_tables = [] - offset = 0 - for group in ensemble_table._iter_groups(): - end = offset + group.size(0) - aligned_tables.extend( - self._align_to_categories( - group, - tuple( - cast(Sequence[Tensor], self._categories[index]) - for index in range(offset, end) - ), + aligned_tables: list[TableTensor | None] = [None] * len(ensemble_table) + for group_id in range(ensemble_table.num_groups): + member_ids = [ + member_id + for member_id, (group, _) in enumerate( + ensemble_table._locations ) + if group == group_id + ] + tables = self._align_to_categories( + table=ensemble_table.expanded_group(group_id), + fitted_categories=tuple( + cast(Sequence[Tensor], self._categories[member_id]) + for member_id in member_ids + ), ) - offset = end - + for member_id, table in zip(member_ids, tables): + aligned_tables[member_id] = table + # Keep members shared when both their query and fitted context share + # a table. Downstream fitted processors depend on that assignment. + outputs: dict[tuple[tuple[int, int], int], int] = {} + member_table_ids = [] + for member_id, location in enumerate(ensemble_table._locations): + key = (location, self._table_ids[member_id]) + member_table_ids.append(outputs.setdefault(key, member_id)) return ensemble_table.replace_tables( - tables=aligned_tables, - member_table_ids=table_ids, + tables=cast(Sequence[TableTensor], aligned_tables), + member_table_ids=member_table_ids, ) @staticmethod @@ -392,6 +430,9 @@ def _align_to_categories( if single_table: codes = codes.unsqueeze(0) aligned_codes = torch.full_like(codes, -1) + output_categories = [ + list(categories) for categories in fitted_categories + ] lookups_by_column: list[dict[int, Tensor]] = [ {} for _ in range(codes.size(-1)) ] @@ -464,17 +505,43 @@ def _align_to_categories( if input_categories.numel() == 0: continue lookup_by_categories = lookups_by_column[column_index] - lookups = [ - lookup_by_categories[id(batch_categories[column_index])] - for batch_categories in fitted_categories - ] + extended_categories: dict[int, Tensor] = {} + lookups = [] + for batch_index, batch_categories in enumerate(fitted_categories): + categories = batch_categories[column_index] + identity = id(categories) + if ( + self.unknown == "append" + and identity not in extended_categories + ): + lookup = lookup_by_categories[identity] + unseen = lookup < 0 + extra = input_categories[unseen] + extended_categories[identity] = ( + extra + if categories.numel() == 0 + else torch.cat([categories, extra]) + ) + appended_codes = unseen.cumsum(0) - 1 + categories.numel() + lookup_by_categories[identity] = torch.where( + unseen, + appended_codes.to(lookup.dtype), + lookup, + ) + if self.unknown == "append": + output_categories[batch_index][column_index] = ( + extended_categories[identity] + ) + else: + assert self.unknown == "missing" + lookups.append(lookup_by_categories[identity]) column_codes = codes[..., column_index] mask = column_codes >= 0 lookup = torch.stack(lookups) indices = column_codes.clamp_min(0).long() aligned_codes[..., column_index] = torch.where( mask, - lookup.gather(1, indices), + lookup.gather(1, indices.flatten(1)).view_as(indices), -1, ) @@ -485,7 +552,7 @@ def _align_to_categories( categories=tuple(batch_categories), ), ) - for batch_index, batch_categories in enumerate(fitted_categories) + for batch_index, batch_categories in enumerate(output_categories) ) def __repr__(self, *, indent: int = 0) -> str: @@ -494,6 +561,8 @@ def __repr__(self, *, indent: int = 0) -> str: arguments.append(f"sort_by={self.sort_by!r}") if self.min_frequency != 1: arguments.append(f"min_frequency={self.min_frequency!r}") + if self.unknown != "missing": + arguments.append(f"unknown={self.unknown!r}") if not arguments: return super().__repr__(indent=indent) return ( diff --git a/test/models/timestral/test_model.py b/test/models/timestral/test_model.py new file mode 100644 index 0000000000..9a19d2dad0 --- /dev/null +++ b/test/models/timestral/test_model.py @@ -0,0 +1,196 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pandas as pd +import pytest +import torch +from huggingface_hub import hf_hub_download +from huggingface_hub.utils import LocalEntryNotFoundError + +from sdm import CategoricalTensor, Stype, TableTensor +from sdm.models import Timestral +from sdm.testing import withCUDA + +WEIGHT_REVISION = "bdd20cfc15a3fea0e4cb5ccef37f8d3ca7aa335f" + + +@withCUDA +@pytest.mark.parametrize("targets", [1, 2]) +def test_forward_and_predict(device: torch.device, targets: int) -> None: + model = Timestral(pretrained=False, device=device) + x = TableTensor.from_pandas( + pd.DataFrame( + { + "past": [0.0, 1, 2, 3, 4, 5, 6], + "event": pd.Categorical(["a", "b", None, "a", "b", "a", "b"]), + } + ), + stypes={"past": Stype.numerical, "event": Stype.categorical}, + device=device, + ) + query = TableTensor.from_pandas( + pd.DataFrame( + { + "event": pd.Categorical( + ["c", "b", "a"], categories=["c", "b", "a"] + ) + } + ), + stypes={"event": Stype.categorical}, + device=device, + ) + y = TableTensor.from_tensor( + torch.arange(7 * targets, device=device).float().view(7, targets), + columns=[f"y{i}" for i in range(targets)], + ) + out = model(x, y, query) + assert out.size() == (3, targets * 21) + assert out.device == device + assert out.numerical.isfinite().all() + quantiles = out.numerical.unflatten(-1, (targets, 21)) + assert (quantiles[..., 1:] >= quantiles[..., :-1]).all() + assert (out.numerical >= 0).all() + assert "y0__q1" in out.column_names + assert "y0__q50" in out.column_names + assert "y0__q99" in out.column_names + if targets == 1: + # Missing learns 2.5, while unseen c uses the context mean of 3. + encoded_context = TableTensor.from_tensor( + torch.tensor( + [ + [0, 2.75], + [1, 3.5], + [2, 2.5], + [3, 2.75], + [4, 3.5], + [5, 2.75], + [6, 3.5], + ], + device=device, + ), + columns=["past", "event"], + ) + encoded_query = TableTensor.from_tensor( + torch.tensor([[3], [3.5], [2.75]], device=device), + columns=["event"], + ) + torch.testing.assert_close( + model(encoded_context, y, encoded_query).numerical, + out.numerical, + ) + model.fit(x, y, context_only_columns=["past"]) + torch.testing.assert_close(model.predict(query).numerical, out.numerical) + + +@withCUDA +@pytest.mark.parametrize("batch_shape", [(), (2,), (2, 1)]) +def test_batch_independence( + device: torch.device, batch_shape: tuple[int, ...] +) -> None: + model = Timestral(pretrained=False, device=device) + target = torch.randn(*batch_shape, 9, 2, device=device) + x = TableTensor(size=(*batch_shape, 9), device=device) + query = TableTensor(size=(*batch_shape, 3), device=device) + out = model(x, TableTensor.from_tensor(target), query, num_estimators=1) + expected = [] + for series in target.reshape(-1, 9, 2): + expected.append( + model( + TableTensor(size=(9,), device=device), + TableTensor.from_tensor(series), + TableTensor(size=(3,), device=device), + ).numerical + ) + torch.testing.assert_close( + out.numerical, torch.stack(expected).reshape(*batch_shape, 3, 42) + ) + + +@withCUDA +@pytest.mark.parametrize( + "case", + [ + "univariate", + "multivariate", + "missing", + "all_missing", + "constant", + "positive", + "numeric", + "indicator", + "categorical", + "categorical_multi", + "long_horizon", + ], +) +def test_pretrained_parity(device: torch.device, case: str) -> None: + upstream = pytest.importorskip("timestral") + # Run the parity script first to cache weights for offline tests. + try: + hf_hub_download( + "mistralai/Timestral-5M", + "model.safetensors", + revision=WEIGHT_REVISION, + local_files_only=True, + ) + except LocalEntryNotFoundError: + pytest.skip( + "Run the Timestral parity script to cache the 5M checkpoint" + ) + model = Timestral(device=device) + reference = upstream.TimestralForecaster.from_pretrained( + "mistralai/Timestral-5M", + revision=WEIGHT_REVISION, + device=device, + local_files_only=True, + ) + t, horizon = 79, 1030 if case == "long_horizon" else 65 + v = 2 if case in {"multivariate", "categorical_multi"} else 1 + time = torch.arange(t + horizon, device=device).float() + target = (time[:t] / 9).sin().expand(1, v, t).clone() + if case == "missing": + target[..., :8] = torch.nan + target[..., 12:15] = torch.nan + if case == "all_missing": + target.fill_(torch.nan) + if case == "constant": + target.fill_(2) + if case == "positive": + target = target.abs() + 1 + x = TableTensor(size=(1, t), device=device) + query = TableTensor(size=(1, horizon), device=device) + covariates = {} + if case in {"numeric", "indicator"}: + known = (time / 7).sin()[None, None] + if case == "indicator": + known = (known > 0).float() + covariates = {"known_covariates": known} + x = TableTensor.from_tensor(known[..., :t].transpose(-2, -1)) + query = TableTensor.from_tensor(known[..., t:].transpose(-2, -1)) + if case.startswith("categorical"): + codes = (time.long() % 4).view(1, -1, 1) + codes[:, 3:6] = -1 + codes[:, t + 4 :] = 4 + known = TableTensor( + columns={Stype.categorical: ["event"]}, + categorical=CategoricalTensor( + codes, [torch.arange(5, device=device)] + ), + ) + x, query = known[..., :t, :], known[..., t:, :] + covariates = { + "categorical_known_covariates": codes.transpose(-2, -1) + .float() + .masked_fill(codes.transpose(-2, -1) < 0, torch.nan) + } + y = TableTensor.from_tensor(target.transpose(-2, -1)) + expected = ( + reference.forecast( + y.numerical.transpose(-2, -1), query.size(-2), **covariates + ) + .quantiles.transpose(-3, -2) + .flatten(-2) + ) + actual = model(x, y, query, num_estimators=1).numerical + tol = 1e-4 if device.type == "cuda" else 2e-5 + torch.testing.assert_close(actual, expected, atol=tol, rtol=tol) diff --git a/test/models/timestral/test_quantile.py b/test/models/timestral/test_quantile.py new file mode 100644 index 0000000000..2af478b5c4 --- /dev/null +++ b/test/models/timestral/test_quantile.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from sdm.models.timestral.block import MonotoneQuantiles +from sdm.testing import withCUDA + + +@withCUDA +def test_monotone_quantiles(device: torch.device) -> None: + levels = [0.01, 0.1, 0.5, 0.8, 0.99] + head = MonotoneQuantiles(levels, width_factor=3).to(device) + x = torch.randn(2, 7, 5, device=device) + out = head(x) + torch.testing.assert_close(out[..., 2], x[..., 2]) + assert (out[..., 1:] >= out[..., :-1]).all() + expected_gap = ( + 3 * (levels[3] - levels[2]) * torch.nn.functional.softplus(x[..., 3]) + ) + torch.testing.assert_close(out[..., 3] - out[..., 2], expected_gap) diff --git a/test/models/timestral/test_residual.py b/test/models/timestral/test_residual.py new file mode 100644 index 0000000000..07132e2ee4 --- /dev/null +++ b/test/models/timestral/test_residual.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from sdm.models.timestral.block import ResidualBlock +from sdm.testing import withCUDA + + +@withCUDA +def test_residual_block(device: torch.device) -> None: + model = ResidualBlock(3, 7, 5, device=device) + x = torch.randn(2, 4, 3, device=device, requires_grad=True) + out = model(x) + assert out.size() == (2, 4, 5) + out.square().sum().backward() + assert x.grad is not None + assert x.grad.isfinite().all() diff --git a/test/nn/test_attention.py b/test/nn/test_attention.py index 83d43f82d1..3703532f97 100644 --- a/test/nn/test_attention.py +++ b/test/nn/test_attention.py @@ -14,6 +14,7 @@ SDPA, Attention, QASSMax, + RotaryEmbedding, SoftplusScale, TransformerBlock, ) @@ -747,3 +748,148 @@ def test_transformer_block_kv_cache() -> None: torch.testing.assert_close(cache_out, direct_out) torch.testing.assert_close(cached_out, direct_out) + + +@withCUDA +@pytest.mark.parametrize("cross_attention", [False, True]) +@pytest.mark.parametrize("block", [False, True]) +def test_attention_query_pos( + device: torch.device, cross_attention: bool, block: bool +) -> None: + rope = RotaryEmbedding(channels=4, layout="split_half", device=device) + if block: + module = TransformerBlock( + channels=4, + num_query_heads=1, + mlp=None, + query_norm=torch.nn.Identity(), + key_value_norm=torch.nn.Identity(), + query_transform=rope, + key_transform=rope, + device=device, + ) + attn = module.attn + else: + module = attn = Attention( + channels=4, + num_query_heads=1, + query_transform=rope, + key_transform=rope, + device=device, + ) + with torch.no_grad(): + attn.qkv_lin.weight.copy_(torch.eye(4, device=device).repeat(3, 1)) + attn.qkv_lin.bias.zero_() + attn.out_lin.weight.copy_(torch.eye(4, device=device)) + attn.out_lin.bias.zero_() + + query = torch.randn(2, 1, 6, 4, device=device) + positions = torch.tensor( + [[[0, 1, 2, 0, 1, 2], [2, 3, 4, 2, 3, 4], [0, 2, 4, 0, 2, 4]]], + device=device, + ) + key_value = ( + torch.randn(1, 3, 4, 4, device=device) if cross_attention else None + ) + key_positions = ( + torch.tensor([0, 2, 1, 3], device=device) if cross_attention else None + ) + keys = query if key_value is None else key_value + expected = reference_sdpa( + query=rope(query.unsqueeze(-2), pos=positions), + key=rope( + keys.unsqueeze(-2), + pos=positions if key_positions is None else key_positions, + ), + value=keys.unsqueeze(-2), + ).squeeze(-2) + if block: + expected = expected + query + + actual, cached = module( + query=query, + key_value=key_value, + query_pos=positions, + key_pos=key_positions, + return_key_value=True, + ) + torch.testing.assert_close(actual, expected) + # Cached keys have already received RoPE and must not be rotated again. + cached_out = module(query=query, key_value=cached, query_pos=positions) + torch.testing.assert_close(cached_out, expected) + + if block: + with torch.no_grad(): + chunked = module( + query=query, + key_value=key_value, + query_pos=positions, + key_pos=key_positions, + batch_size_limit=2, + ) + torch.testing.assert_close(chunked, expected) + + contiguous = torch.arange(query.size(-2), device=device) + torch.testing.assert_close( + module(query=query), module(query=query, query_pos=contiguous) + ) + + +@withCUDA +@pytest.mark.parametrize("block", [False, True]) +@pytest.mark.parametrize("gate_channels", [1, 8]) +def test_attention_gate( + device: torch.device, block: bool, gate_channels: int +) -> None: + gate = torch.nn.Sequential( + torch.nn.Linear(8, gate_channels, device=device), + torch.nn.Sigmoid(), + ) + norm = torch.nn.LayerNorm(8, device=device) + if block: + module = TransformerBlock( + channels=8, + num_query_heads=2, + mlp=None, + query_norm=norm, + gate=gate, + device=device, + ) + attn = module.attn + else: + module = attn = Attention( + channels=8, + num_query_heads=2, + gate=gate, + device=device, + ) + # A non-diagonal projection and nonzero bias distinguish gating before + # projection from gating after it. + with torch.no_grad(): + attn.out_lin.weight.normal_() + attn.out_lin.bias.normal_() + query = torch.randn(3, 4, 8, device=device, requires_grad=True) + inputs = norm(query) if block else query + q, k, v = attn.qkv_lin(inputs).chunk(3, dim=-1) + attended = reference_sdpa( + query=q.unflatten(-1, (2, 4)), + key=k.unflatten(-1, (2, 4)), + value=v.unflatten(-1, (2, 4)), + ).flatten(-2) + expected = attn.out_lin(attended * gate(inputs)) + if block: + expected = expected + query + actual, cached = module(query=query, return_key_value=True) + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(module(query=query, key_value=cached), expected) + parameters = (query, *gate.parameters()) + for actual_grad, expected_grad in zip( + torch.autograd.grad(actual.sum(), parameters), + torch.autograd.grad(expected.sum(), parameters), + ): + torch.testing.assert_close(actual_grad, expected_grad) + + if block: + with torch.no_grad(): + chunked = module(query=query, batch_size_limit=1) + torch.testing.assert_close(chunked, expected) diff --git a/test/nn/test_compile.py b/test/nn/test_compile.py index d61f15a436..6a062167b4 100644 --- a/test/nn/test_compile.py +++ b/test/nn/test_compile.py @@ -61,9 +61,11 @@ def test_qassmax_compile( @withCUDA @pytest.mark.parametrize("partial_rotary_factor", [1.0, 0.25]) +@pytest.mark.parametrize("explicit_positions", [False, True]) def test_rotary_embedding_compile( device: torch.device, partial_rotary_factor: float, + explicit_positions: bool, ) -> None: module = RotaryEmbedding( channels=16, @@ -73,8 +75,13 @@ def test_rotary_embedding_compile( ) x = torch.randn(2, 5, 3, 16, device=device) - expected = module(x) - out = fullgraph(module)(x) + positions = None + if explicit_positions: + positions = torch.tensor( + [[0, 2, -3, 2, 7], [-1, 4, 0, 0, 3]], device=device + ) + expected = module(x, pos=positions) + out = fullgraph(module)(x, pos=positions) torch.testing.assert_close(out, expected) diff --git a/test/nn/test_rope.py b/test/nn/test_rope.py index 6485fcb439..b1e847db8d 100644 --- a/test/nn/test_rope.py +++ b/test/nn/test_rope.py @@ -86,3 +86,62 @@ def test_partial_rope( assert torch.equal(output[..., 4:], x[..., 4:]) expected_prefix = reference(x[..., :4]) torch.testing.assert_close(output[..., :4], expected_prefix) + + +@withCUDA +@pytest.mark.parametrize("layout", ["split_half", "interleaved"]) +@pytest.mark.parametrize("partial_rotary_factor", [1.0, 0.5]) +@pytest.mark.parametrize("batched_positions", [False, True]) +def test_rope_explicit_positions( + device: torch.device, + layout: Literal["split_half", "interleaved"], + partial_rotary_factor: float, + batched_positions: bool, +) -> None: + module = RotaryEmbedding( + channels=16, + layout=layout, + partial_rotary_factor=partial_rotary_factor, + device=device, + ) + positions = torch.tensor([0, 7, 2, 7, 3], device=device) + if batched_positions: + positions = torch.stack([positions, positions.flip(0)]).unsqueeze(1) + # Repeat each vector across time so rotations can be selected from a + # default sequential call. Positions broadcast over the second batch axis. + x = torch.randn(2, 3, 1, 2, 16, device=device) + sequential = module(x.expand(2, 3, 8, 2, 16)) + index = positions[..., None, None].expand(2, 3, 5, 2, 16) + expected = sequential.gather(-3, index) + out = module(x.expand(2, 3, 5, 2, 16), pos=positions) + torch.testing.assert_close(out, expected) + rotated = int(16 * partial_rotary_factor) + torch.testing.assert_close( + out[..., rotated:], x[..., rotated:].expand_as(out[..., rotated:]) + ) + + +@withCUDA +@pytest.mark.parametrize("layout", ["split_half", "interleaved"]) +def test_rope_negative_positions_invert_rotation( + device: torch.device, + layout: Literal["split_half", "interleaved"], +) -> None: + module = RotaryEmbedding( + channels=16, + layout=layout, + partial_rotary_factor=0.5, + device=device, + ) + x = torch.randn(2, 5, 3, 16, device=device, requires_grad=True) + positions = torch.tensor( + [[0, 2, -3, 2, 7], [-1, 4, 0, 0, 3]], device=device + ) + rotated = module(x, pos=positions) + recovered = module(rotated, pos=-positions) + torch.testing.assert_close(recovered, x) + rotated.sum().backward() + assert x.grad is not None + assert x.grad.isfinite().all() + assert module.inv_freq.grad is not None + assert module.inv_freq.grad.isfinite().all() diff --git a/test/processing/categorical/test_align.py b/test/processing/categorical/test_align.py index 52cf0a6b95..62c9ef5454 100644 --- a/test/processing/categorical/test_align.py +++ b/test/processing/categorical/test_align.py @@ -7,6 +7,7 @@ import pytest import torch +import sdm.processing as sp from sdm import CategoricalTensor, EnsembleTable, StringTensor, TableTensor from sdm.processing import AlignCategories from sdm.testing import withCUDA @@ -506,6 +507,160 @@ def test_align_categories_ensemble_matches_member_fits( assert fitted_query_output[member_id].equal(expected_query) +@withCUDA +@pytest.mark.parametrize("dtype", [torch.int32, torch.int64]) +def test_append_categories_preserves_unseen_and_missing( + device: torch.device, + dtype: torch.dtype, +) -> None: + context = _table( + [[0], [1], [-1]], + categories=(("red", "blue", "unused"),), + device=device, + dtype=dtype, + ) + query = _table( + [[0], [1], [2], [-1]], + categories=(("green", "blue", "red"),), + device=device, + dtype=dtype, + ) + processor = AlignCategories(min_frequency=0, unknown="append") + assert processor.fit_transform(context).equal(context) + out = processor.transform(query) + assert out.categorical.code.dtype == dtype + assert out.categorical.code.squeeze(-1).tolist() == [3, 1, 0, -1] + assert out.categorical.categories[0].tolist() == [ + "red", + "blue", + "unused", + "green", + ] + other_query = _table( + [[0]], + categories=(("yellow",),), + device=device, + dtype=dtype, + ) + other = processor.transform(other_query) + assert other.categorical.code.squeeze(-1).tolist() == [3] + assert other.categorical.categories[0].tolist() == [ + "red", + "blue", + "unused", + "yellow", + ] + assert processor.transform(query).equal(out) + assert processor.transform(context).equal(context) + + +@withCUDA +def test_append_categories_empty_fitted_vocabulary( + device: torch.device, +) -> None: + context = TableTensor.from_pandas( + pd.DataFrame({"kind": [None, None]}), + stypes={"kind": "categorical"}, + device=device, + ) + query = TableTensor.from_pandas( + pd.DataFrame({"kind": ["red", None, "blue"]}), + stypes={"kind": "categorical"}, + device=device, + ) + processor = AlignCategories(min_frequency=0, unknown="append").fit(context) + assert processor.transform(query).equal(query) + assert processor.transform(context).equal(context) + + +@withCUDA +def test_append_categories_numeric_values(device: torch.device) -> None: + context = TableTensor( + categorical=CategoricalTensor( + code=torch.tensor([[0], [1], [-1]], device=device), + categories=(torch.tensor([20, 10, 30], device=device),), + ), + ) + query = TableTensor( + categorical=CategoricalTensor( + code=torch.tensor([[0], [1], [2], [-1]], device=device), + categories=(torch.tensor([10, 40, 20], device=device),), + ), + ) + out = ( + AlignCategories(min_frequency=0, unknown="append") + .fit(context) + .transform(query) + ) + assert out.categorical.code.squeeze(-1).tolist() == [1, 3, 0, -1] + assert out.categorical.categories[0].tolist() == [20, 10, 30, 40] + + +@withCUDA +@pytest.mark.parametrize("append", [False, True]) +def test_align_categories_shared_query_uses_each_fitted_vocabulary( + device: torch.device, + append: bool, +) -> None: + contexts = ( + _table( + [[0], [1]], categories=(("red", "blue", "unused"),), device=device + ), + _table([[0], [1]], categories=(("blue", "red"),), device=device), + ) + query = _table( + [[0], [1], [-1]], + categories=(("green", "red"),), + device=device, + ) + member_ids = (1, 0, 1, 0) + kwargs = {"min_frequency": 0, "unknown": "append"} if append else {} + processor = AlignCategories(**kwargs) + processor.fit_ensemble( + EnsembleTable.from_tables(tables=contexts, member_table_ids=member_ids) + ) + out = processor.transform_ensemble( + EnsembleTable.from_table(query, num_members=4) + ) + for i, member in enumerate(member_ids): + expected = ( + AlignCategories(**kwargs).fit(contexts[member]).transform(query) + ) + assert out[i].equal(expected) + + +@withCUDA +def test_alignment_preserves_fitted_ensemble_pipeline( + device: torch.device, +) -> None: + context = _table( + [[0], [1], [0]], + categories=(("red", "blue"),), + device=device, + ) + query = _table([[1], [0]], categories=(("blue", "red"),), device=device) + processor = sp.Sequential( + sp.AlignCategories(), + sp.ToNumerical(), + sp.Choice( + sp.Identity(), sp.FlipSign(probability=1), method="round_robin" + ), + sp.Standardize(), + ) + members = 4 + processor.fit_transform_ensemble( + EnsembleTable.from_table(context, num_members=members) + ) + out = processor.transform_ensemble( + EnsembleTable.from_table(query, num_members=members) + ) + # Each sign branch uses its own context-fitted statistics. + expected = torch.tensor([[-(2**-0.5)], [2**0.5]], device=device) + for member_id in range(len(out)): + sign = 1 if member_id % 2 == 0 else -1 + torch.testing.assert_close(out[member_id].numerical, sign * expected) + + def test_string_category_lookups_compile(compilation: None) -> None: fitted = ( StringTensor.from_list(["a", "b"]), diff --git a/test/processing/test_contract.py b/test/processing/test_contract.py index bd3ebe6e20..a736af0dc7 100644 --- a/test/processing/test_contract.py +++ b/test/processing/test_contract.py @@ -136,6 +136,10 @@ def _make_processor_pair( ProcessorCase(sp.RandomProjection(2)), ProcessorCase(sp.WarpTails(probability=1.0)), ProcessorCase(sp.AlignCategories(), _make_align_categories_table), + ProcessorCase( + sp.AlignCategories(min_frequency=0, unknown="append"), + _make_align_categories_table, + ), ProcessorCase(sp.ShuffleCategories()), ProcessorCase(sp.ImputeMode()), ProcessorCase(sp.AddCategoryCounts()),