diff --git a/THIRD_PARTY_LICENSES.md b/THIRD_PARTY_LICENSES.md index 686a4f9559..99fd0e0fdc 100644 --- a/THIRD_PARTY_LICENSES.md +++ b/THIRD_PARTY_LICENSES.md @@ -57,3 +57,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 and the `sdm.nn.MonotoneQuantiles` and `sdm.nn.CausalPatchNorm` modules 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..50926d37e3 --- /dev/null +++ b/benchmark/forecasting/timestral_parity.md @@ -0,0 +1,56 @@ +# 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-17 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 `1.05e-5` | +| CPU, 1,030-step horizon | Maximum absolute error `0.0635`; 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 now uses SDM's shared `RotaryEmbedding` with explicit positions. 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. + +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. 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 48306b949d..41764f61f1 100644 --- a/docs/source/api/models.rst +++ b/docs/source/api/models.rst @@ -50,6 +50,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 277e31f3d9..4982bb82cf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,6 +15,8 @@ license-files = [ "sdm/models/tabiclv2/LICENSE", "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 57a27fa914..734d184d7f 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 ICLModel @@ -16,6 +17,7 @@ from sdm.models.tabfm import TabFM from sdm.models.kumo import KumoTabular, KumoRelational, KumoForecast from sdm.models.timesfm3 import TimesFM3 +from sdm.models.timestral import Timestral __all__ = [ @@ -27,4 +29,5 @@ "KumoRelational", "KumoForecast", "TimesFM3", + "Timestral", ] diff --git a/sdm/models/timestral/LICENSE b/sdm/models/timestral/LICENSE new file mode 100644 index 0000000000..d645695673 --- /dev/null +++ b/sdm/models/timestral/LICENSE @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. 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 new file mode 100644 index 0000000000..93362a1348 --- /dev/null +++ b/sdm/models/timestral/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Timestral multivariate forecasting.""" + +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..4bb997f54c --- /dev/null +++ b/sdm/models/timestral/config.py @@ -0,0 +1,54 @@ +# 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 + split_context_horizon_attn: bool = True + attend_missing_patches: bool = True + quantile_width_factor: float = 2.0 + + @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..40ad38fb10 --- /dev/null +++ b/sdm/models/timestral/decode.py @@ -0,0 +1,324 @@ +# 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.stats(inputs, observation_mask) + anchor_scale = scale[..., -1:] * 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/layers.py b/sdm/models/timestral/layers.py new file mode 100644 index 0000000000..0c1dc42601 --- /dev/null +++ b/sdm/models/timestral/layers.py @@ -0,0 +1,215 @@ +# 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 + +import torch +from torch import Tensor, nn + +from sdm.models.timestral.config import LAYER_KIND_AXES +from sdm.nn import SDPA, RMSNorm, RotaryEmbedding, SwiGLU + + +@dataclass(frozen=True) +class _TwoStreamMask: + mask: Tensor + enc_len: int + + +class _Attention(nn.Module): + def __init__( + self, + embed_dim: int, + num_heads: int, + is_time: bool, + use_bias: bool, + gate: str, + ) -> None: + super().__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.head_dim = embed_dim // num_heads + self.is_time = is_time + self.sdpa = SDPA(num_query_heads=num_heads) + self.qkv_lin = nn.Linear(embed_dim, 3 * embed_dim, bias=use_bias) + self.out_lin = nn.Linear(embed_dim, embed_dim, bias=use_bias) + self.gate_lin = ( + nn.Linear(embed_dim, embed_dim, bias=use_bias) + if gate == "sdpa_sigmoid" + else None + ) + + def _qkv(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: + qkv = self.qkv_lin(x).unflatten(-1, (3, self.num_heads, self.head_dim)) + return qkv.unbind(-3) + + def _attend( + self, + query: Tensor, + key: Tensor, + value: Tensor, + mask: Tensor | None, + is_causal: bool = False, + ) -> Tensor: + return self.sdpa( + query=query, + key=key, + value=value, + attn_mask=mask, + is_causal=is_causal, + ).flatten(-2) + + def forward( + self, + x: Tensor, + mask: Tensor | _TwoStreamMask | None, + positions: Tensor, + rope: RotaryEmbedding, + ) -> Tensor: + rows, num_patches, _ = x.shape + inputs = x if self.is_time else x.permute(1, 0, 2) + q, k, v = self._qkv(inputs) + if self.is_time: + # Queries and keys share positions and rotary frequencies. + q, k = rope( + torch.stack([q, k]), + position_ids=positions, + ).unbind(0) + if isinstance(mask, _TwoStreamMask): + n = mask.enc_len + out_enc = self._attend( + query=q[:, :n], + key=k[:, :n], + value=v[:, :n], + mask=mask.mask[:, :, :n, :n], + ) + out_dec = self._attend(q[:, n:], k, v, mask.mask[:, :, n:, :]) + out = torch.cat([out_enc, out_dec], dim=1) + else: + out = self._attend( + query=q, + key=k, + value=v, + mask=mask, + is_causal=self.is_time and mask is None, + ) + if self.gate_lin is not None: + out = out * self.gate_lin(inputs).sigmoid() + if self.is_time: + out = out.reshape(rows, num_patches, self.embed_dim) + else: + out = out.reshape(num_patches, rows, self.embed_dim).permute( + 1, 0, 2 + ) + return self.out_lin(out) + + +class _TransformerBlock(nn.Module): + def __init__( + self, + kind: str, + embed_dim: int, + num_heads: int, + mlp_hidden_dim: int, + norm_eps: float, + attn_bias: bool, + mlp_bias: bool, + time_attn_gate: str, + variate_attn_gate: str, + ) -> None: + super().__init__() + axes = LAYER_KIND_AXES[kind] + self.attn_norms = nn.ModuleList( + [RMSNorm(embed_dim, eps=norm_eps) for _ in axes] + ) + self.attns = nn.ModuleList( + [ + _Attention( + embed_dim=embed_dim, + num_heads=num_heads, + is_time=is_time, + use_bias=attn_bias, + gate=time_attn_gate if is_time else variate_attn_gate, + ) + for is_time in axes + ] + ) + self.mlp_norm = RMSNorm(embed_dim, eps=norm_eps) + self.mlp = SwiGLU(embed_dim, mlp_hidden_dim, bias=mlp_bias) + + def forward( + self, + x: Tensor, + time_mask: Tensor | _TwoStreamMask | None, + variate_mask: Tensor | None, + positions: Tensor, + rope: RotaryEmbedding, + ) -> Tensor: + for norm, attn in zip(self.attn_norms, self.attns): + x = x + attn( + x=norm(x), + mask=time_mask if attn.is_time else variate_mask, + positions=positions, + rope=rope, + ) + return x + self.mlp(self.mlp_norm(x)) + + +class _Transformer(nn.Module): + def __init__( + self, + block_schedule: tuple[str, ...], + embed_dim: int, + num_heads: int, + mlp_hidden_dim: int, + rope_theta: float, + rope_dim: int, + norm_eps: float, + attn_bias: bool, + mlp_bias: bool, + time_attn_gate: str, + variate_attn_gate: str, + ) -> None: + super().__init__() + self.layers = nn.ModuleList( + [ + _TransformerBlock( + kind=kind, + embed_dim=embed_dim, + num_heads=num_heads, + mlp_hidden_dim=mlp_hidden_dim, + norm_eps=norm_eps, + attn_bias=attn_bias, + mlp_bias=mlp_bias, + time_attn_gate=time_attn_gate, + variate_attn_gate=variate_attn_gate, + ) + for kind in block_schedule + ] + ) + self.rope = RotaryEmbedding( + channels=embed_dim // num_heads, + layout="split_half", + theta=rope_theta, + partial_rotary_factor=rope_dim / (embed_dim // num_heads), + requires_grad=False, + ) + self.out_norm = RMSNorm(embed_dim, eps=norm_eps) + + def forward( + self, + x: Tensor, + time_mask: Tensor | _TwoStreamMask | None, + variate_mask: Tensor | None, + positions: Tensor, + ) -> Tensor: + for layer in self.layers: + x = layer( + x=x, + time_mask=time_mask, + variate_mask=variate_mask, + positions=positions, + rope=self.rope, + ) + return self.out_norm(x) 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..7995593d4f --- /dev/null +++ b/sdm/models/timestral/model.py @@ -0,0 +1,333 @@ +# 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) + state["transformer.rope.inv_freq"] = ( + self.model.transformer.rope.inv_freq.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 :], + ] + ) + for key, value in state.items(): + key = ( + key.replace("hidden_layer", "hidden_lin") + .replace("output_layer", "out_lin") + .replace("residual_layer", "residual_lin") + ) + 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") + key = key.replace(".wO.", ".out_lin.").replace("._gate.", ".gate_lin.") + if key.endswith(".scale"): + key = key[:-5] + "weight" + if ".mlp.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.w3.", ".mlp.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..7a73dc4aea --- /dev/null +++ b/sdm/models/timestral/network.py @@ -0,0 +1,311 @@ +# 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.config import _TimestralConfig +from sdm.models.timestral.layers import _Transformer, _TwoStreamMask +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 CausalPatchNorm, MonotoneQuantiles, ResidualMLP + +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 = ResidualMLP( + 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.hidden_lin.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 = ResidualMLP( + 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 = CausalPatchNorm(c.patch_size, c.minimum_scale, c.use_asinh) + self.representation = _PatchRepresentation(c.patch_size, c.d_model) + self.transformer = _Transformer( + block_schedule=c.block_schedule, + embed_dim=c.d_model, + num_heads=c.num_heads, + mlp_hidden_dim=c.mlp_hidden_dim, + rope_theta=c.rope_theta, + rope_dim=c.rope_dim, + norm_eps=c.norm_eps, + attn_bias=c.attn_bias, + mlp_bias=c.mlp_bias, + time_attn_gate=c.time_attn_gate, + variate_attn_gate=c.variate_attn_gate, + ) + self.head = _QuantileHead(c) + + def scale( + self, inputs: Tensor, observed: Tensor + ) -> tuple[Tensor, Tensor, Tensor]: + return self.norm(inputs, observed) + + 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) + mask: Tensor | _TwoStreamMask = time_mask + if self.config.split_context_horizon_attn: + mask = _TwoStreamMask(mask=time_mask, enc_len=x_enc.shape[1]) + h = self.transformer(x, mask, variate_mask, 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..913d012d13 --- /dev/null +++ b/sdm/models/timestral/scaler.py @@ -0,0 +1,168 @@ +# 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 + + +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 _welford_combine( + w_a: Tensor, + mean_a: Tensor, + m2_a: Tensor, + w_b: Tensor, + mean_b: Tensor, + m2_b: Tensor, +) -> tuple[Tensor, Tensor, Tensor]: + w = w_a + w_b + w_safe = w.clamp(min=torch.finfo(w.dtype).tiny) + delta = mean_b - mean_a + mean = mean_a + delta * w_b / w_safe + m2 = m2_a + m2_b + delta.square() * w_a * w_b / w_safe + is_zero = w == 0 + mean = torch.where(is_zero, torch.zeros_like(mean), mean) + m2 = torch.where(is_zero, torch.zeros_like(m2), m2) + return (w, mean, m2) + + +def _cumulative_welford( + w: Tensor, mean: Tensor, m2: Tensor +) -> tuple[Tensor, Tensor, Tensor]: + n = w.shape[1] + d = 1 + while d < n: + pad = (d, 0) + w_l = torch.nn.functional.pad(w[:, : n - d], pad) + mean_l = torch.nn.functional.pad(mean[:, : n - d], pad) + m2_l = torch.nn.functional.pad(m2[:, : n - d], pad) + w, mean, m2 = _welford_combine(w_l, mean_l, m2_l, w, mean, m2) + d *= 2 + return (w, mean, m2) + + +def _advancing_block_frames( + context: Tensor, + observation_mask: Tensor, + sinh_median: Tensor, + minimum_scale: float, +) -> tuple[Tensor, Tensor]: + r, n_pred, ps = sinh_median.shape + blank = sinh_median.new_zeros(r, n_pred, ps) + values = torch.cat( + [context.unflatten(-1, (-1, ps)).to(torch.float64), blank], dim=1 + ) + weights = torch.cat( + [observation_mask.unflatten(-1, (-1, ps)).to(torch.float64), blank], + dim=1, + ) + cnt = weights.sum(-1) + mean = (values * weights).sum(-1) / cnt.clamp(min=1.0) + m2 = (weights * (values - mean.unsqueeze(-1)).square()).sum(-1) + is_zero = cnt == 0 + mean = torch.where(is_zero, torch.zeros_like(mean), mean) + m2 = torch.where(is_zero, torch.zeros_like(m2), m2) + cnt, mean, m2 = ( + s[:, -n_pred - 1] for s in _cumulative_welford(cnt, mean, m2) + ) + offset = mean + clip = torch.finfo(torch.float32).max ** 0.5 / 1000.0 + locs, scales = ([], []) + for h in range(n_pred): + var = (m2 / (cnt - 1.0).clamp(min=1.0)).clamp(min=0.0) + scale = var.sqrt().clamp(min=minimum_scale) + locs.append(mean) + scales.append(scale) + folded = ( + sinh_median[:, h] * scale[:, None] + + mean[:, None] + - offset[:, None] + ).clamp(-clip, clip) + folded = folded + offset[:, None] + b_mean = folded.mean(-1) + b_m2 = (folded - b_mean[:, None]).square().sum(-1) + cnt, mean, m2 = _welford_combine( + w_a=cnt, + mean_a=mean, + m2_a=m2, + w_b=torch.full_like(cnt, ps), + mean_b=b_mean, + m2_b=b_m2, + ) + 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 + stat_dtype = torch.promote_types(orig_dtype, torch.float32) + x = values.to(stat_dtype) + mask_f = observation_mask.to(stat_dtype) + n = mask_f.sum(dim=-1, keepdim=True).clamp(min=1.0) + loc = (x * mask_f).sum(dim=-1, keepdim=True) / n + var = (((x - loc) * mask_f) ** 2).sum(dim=-1, keepdim=True) / n + scale = var.sqrt().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 & observation_mask.any(dim=-1, keepdim=True) + 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/nn/__init__.py b/sdm/nn/__init__.py index b10072f067..79cc7d90ee 100644 --- a/sdm/nn/__init__.py +++ b/sdm/nn/__init__.py @@ -6,9 +6,12 @@ from sdm.nn.rope import RotaryEmbedding from sdm.nn.sinusoidal import SinusoidalEmbedding from sdm.nn.glu import SwiGLU +from sdm.nn.residual import ResidualMLP +from sdm.nn.quantile import MonotoneQuantiles from sdm.nn.softplus import SoftplusScale from sdm.nn.rms_norm import RMSNorm from sdm.nn.revin import RevIN +from sdm.nn.causal_norm import CausalPatchNorm from sdm.nn.scaling import QueryScaling, QASSMax, LogScale, GatedLogScale from sdm.nn.attention import SDPA, Attention, TransformerBlock from sdm.nn.set_transformer import InducedTransformerBlock @@ -18,9 +21,12 @@ "RotaryEmbedding", "SinusoidalEmbedding", "SwiGLU", + "ResidualMLP", + "MonotoneQuantiles", "SoftplusScale", "RMSNorm", "RevIN", + "CausalPatchNorm", "QueryScaling", "QASSMax", "LogScale", diff --git a/sdm/nn/causal_norm.py b/sdm/nn/causal_norm.py new file mode 100644 index 0000000000..7b4c072833 --- /dev/null +++ b/sdm/nn/causal_norm.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 + +import torch +from torch import Tensor + + +class CausalPatchNorm(torch.nn.Module): + """Normalize patches using running observed means and sample deviations. + + Each patch uses observations through its last step. Statistics accumulate + in float64 and return in the input dtype. Unobserved values become zero. + + Args: + patch_size: Number of time steps per patch. + minimum_scale: Floor for the running standard deviation. + use_asinh: Apply an inverse hyperbolic sine after normalization. + """ + + def __init__( + self, + patch_size: int, + minimum_scale: float = 1e-6, + use_asinh: bool = False, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.minimum_scale = minimum_scale + self.use_asinh = use_asinh + + @torch.no_grad() + def stats(self, x: Tensor, observed: Tensor) -> tuple[Tensor, Tensor]: + """Return running patch means and scales, broadcast to time steps. + + Args: + x: Values ``[..., T]``, with ``T`` divisible by the patch size. + observed: Boolean observation mask with the same shape as ``x``. + """ + values = x.double().masked_fill(~observed, 0) + weights = observed.double() + count = weights.cumsum(-1).clamp_min(1) + mean = values.cumsum(-1) / count + previous = torch.nn.functional.pad(mean[..., :-1], (1, 0)) + # Welford increments preserve accuracy for large offsets. + increments = (values - previous) * (values - mean) * weights + variance = increments.cumsum(-1) / (count - 1).clamp_min(1) + scale = variance.sqrt().clamp_min(self.minimum_scale) + mean = mean.to(x.dtype).unflatten(-1, (-1, self.patch_size))[..., -1] + scale = scale.to(x.dtype).unflatten(-1, (-1, self.patch_size))[..., -1] + return ( + mean.repeat_interleave(self.patch_size, dim=-1), + scale.repeat_interleave(self.patch_size, dim=-1), + ) + + def forward( + self, x: Tensor, observed: Tensor + ) -> tuple[Tensor, Tensor, Tensor]: + """Return normalized values, patch means, and patch scales. + + Args: + x: Values ``[..., T]``, with ``T`` divisible by the patch size. + observed: Boolean observation mask with the same shape as ``x``. + """ + mean, scale = self.stats(x, observed) + dtype = torch.promote_types(x.dtype, torch.float32) + out = (x.to(dtype) - mean.to(dtype)) / scale.to(dtype).clamp_min( + torch.finfo(dtype).eps + ) + out = out.masked_fill(~observed, 0) + if self.use_asinh: + out = out.asinh() + return out.to(x.dtype), mean, scale diff --git a/sdm/nn/quantile.py b/sdm/nn/quantile.py new file mode 100644 index 0000000000..db09c258f0 --- /dev/null +++ b/sdm/nn/quantile.py @@ -0,0 +1,55 @@ +# 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 collections.abc import Sequence + +import torch +from torch import Tensor + + +class MonotoneQuantiles(torch.nn.Module): + """Construct ordered quantiles from a median and positive increments. + + Args: + levels: Increasing probability levels containing ``0.5``. + width_factor: Multiplier for probability-weighted softplus increments. + """ + + widths: Tensor + + def __init__( + self, levels: Sequence[float], width_factor: float = 1.0 + ) -> None: + super().__init__() + self.median_index = list(levels).index(0.5) + widths = [ + levels[i + 1] - levels[i] + if i < self.median_index + else levels[i] - levels[i - 1] + if i > self.median_index + else 0.0 + for i in range(len(levels)) + ] + self.register_buffer("widths", torch.tensor(widths), persistent=False) + self.width_factor = width_factor + + def forward(self, x: Tensor) -> Tensor: + """Return ordered quantiles with the same shape as ``x``. + + Args: + x: Raw parameters ``[..., Q]``, with ``Q`` probability levels. + """ + m = self.median_index + increments = ( + self.width_factor + * self.widths.to(x) + * torch.nn.functional.softplus(x) + ) + median = x[..., m : m + 1] + up = increments[..., m + 1 :].cumsum(-1) + down = increments[..., :m].flip(-1).cumsum(-1).flip(-1) + limit = torch.finfo(x.dtype).max + return torch.cat([median - down, median, median + up], dim=-1).clamp( + -limit, limit + ) diff --git a/sdm/nn/residual.py b/sdm/nn/residual.py new file mode 100644 index 0000000000..035a5252e5 --- /dev/null +++ b/sdm/nn/residual.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from torch import Tensor +from torch.nn import Linear + + +class ResidualMLP(torch.nn.Module): + """A two-layer SiLU MLP with a learned linear residual projection. + + Args: + in_channels: Number of input channels. + hidden_channels: Number of hidden channels. + out_channels: Number of output channels. + bias: Whether the projections learn additive biases. + device: The device. + dtype: The dtype. + """ + + def __init__( + self, + in_channels: int, + hidden_channels: int, + out_channels: int, + bias: bool = True, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> None: + super().__init__() + kwargs = {"bias": bias, "device": device, "dtype": dtype} + self.hidden_lin = Linear(in_channels, hidden_channels, **kwargs) + self.out_lin = Linear(hidden_channels, out_channels, **kwargs) + self.residual_lin = Linear(in_channels, out_channels, **kwargs) + + def forward(self, x: Tensor) -> Tensor: + """Transform the last dimension of ``x``. + + Args: + x: Input with shape ``[..., C]``, where ``C`` is input channels. + """ + return self.out_lin( + torch.nn.functional.silu(self.hidden_lin(x)) + ) + self.residual_lin(x) diff --git a/sdm/nn/rope.py b/sdm/nn/rope.py index 9db127ef69..2274032520 100644 --- a/sdm/nn/rope.py +++ b/sdm/nn/rope.py @@ -66,6 +66,7 @@ def __init__( def forward( self, x: Tensor, # [..., S, H, C] + position_ids: Tensor | None = None, # [..., S] ) -> Tensor: # [..., S, H, C] """The forward pass. @@ -73,6 +74,8 @@ def forward( x: Tensor with shape ``[..., S, H, C]``. ``S`` is the query sequence length, ``H`` is the number of attention heads, and ``C`` is the channels per head. + position_ids: Positions with shape ``[..., S]``, broadcastable + over the input batch dimensions. Defaults to ``0..S-1``. Returns: Tensor with shape ``[..., S, H, C]``. @@ -81,10 +84,15 @@ def forward( raise ValueError( f"Expected {self.channels} channels (got {x.size(-1)})" ) - seq = torch.arange(x.size(-3), device=x.device, dtype=torch.float32) - freq = seq.view(-1, 1) * self.inv_freq.view(1, -1) # [S, C // 2] - sin = freq.sin()[:, None, :].to(x.dtype) # [S, 1, C // 2] - cos = freq.cos()[:, None, :].to(x.dtype) # [S, 1, C // 2] + if position_ids is None: + position_ids = torch.arange( + x.size(-3), + device=x.device, + dtype=torch.float32, + ) + freq = position_ids.float().unsqueeze(-1) * self.inv_freq + sin = freq.sin().unsqueeze(-2).to(x.dtype) # [..., S, 1, C_rot // 2] + cos = freq.cos().unsqueeze(-2).to(x.dtype) rotary = x[..., : self.rotary_channels] if self.layout == "interleaved": diff --git a/sdm/processing/categorical/align.py b/sdm/processing/categorical/align.py index 0d76bd808f..c79aeadb73 100644 --- a/sdm/processing/categorical/align.py +++ b/sdm/processing/categorical/align.py @@ -23,22 +23,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 @@ -76,13 +82,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, @@ -96,6 +117,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]. @@ -226,6 +251,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 ) @@ -237,6 +263,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 ) @@ -278,8 +305,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( @@ -300,8 +328,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 @@ -309,25 +338,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 @@ -431,6 +469,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 table.categorical.categories ] @@ -503,17 +544,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, ) @@ -524,7 +591,7 @@ def _align_to_categories( categories=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: @@ -533,6 +600,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/nn/test_causal_norm.py b/test/nn/test_causal_norm.py new file mode 100644 index 0000000000..83d3e519e5 --- /dev/null +++ b/test/nn/test_causal_norm.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from sdm.nn import CausalPatchNorm +from sdm.testing import withCUDA + + +@withCUDA +def test_causal_patch_norm(device: torch.device) -> None: + x = torch.tensor([[1.0, 3.0, float("nan"), 7.0, 9.0, 11.0]], device=device) + observed = x.isfinite() + norm = CausalPatchNorm(patch_size=2, use_asinh=True) + out, mean, scale = norm(x, observed) + for end in (2, 4, 6): + values = x[0, :end][observed[0, :end]] + torch.testing.assert_close( + mean[0, end - 2 : end], values.mean().expand(2) + ) + torch.testing.assert_close( + scale[0, end - 2 : end], values.std().expand(2) + ) + assert out[0, 2] == 0 + torch.testing.assert_close( + out.sinh() * scale + mean, x.masked_fill(~observed, mean[0, 2]) + ) + changed = x.clone() + changed[..., 4:] += 100 + torch.testing.assert_close( + norm(changed, observed)[0][..., :4], out[..., :4] + ) diff --git a/test/nn/test_compile.py b/test/nn/test_compile.py index d61f15a436..c835049084 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, position_ids=positions) + out = fullgraph(module)(x, position_ids=positions) torch.testing.assert_close(out, expected) diff --git a/test/nn/test_quantile.py b/test/nn/test_quantile.py new file mode 100644 index 0000000000..2fc3f9feb5 --- /dev/null +++ b/test/nn/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.nn 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/nn/test_residual.py b/test/nn/test_residual.py new file mode 100644 index 0000000000..e96f611360 --- /dev/null +++ b/test/nn/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.nn import ResidualMLP +from sdm.testing import withCUDA + + +@withCUDA +def test_residual_mlp(device: torch.device) -> None: + model = ResidualMLP(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_rope.py b/test/nn/test_rope.py index 6485fcb439..7887caa1ad 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), position_ids=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, position_ids=positions) + recovered = module(rotated, position_ids=-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 b40bf306e4..60ad7ef1e3 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 @@ -504,3 +505,157 @@ def test_align_categories_ensemble_matches_member_fits( expected_query = reference.transform(query) assert query_output[member_id].equal(expected_query) 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) diff --git a/test/processing/test_contract.py b/test/processing/test_contract.py index 7525b3b8c7..68cd6fa374 100644 --- a/test/processing/test_contract.py +++ b/test/processing/test_contract.py @@ -135,6 +135,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()),