diff --git a/examples/timesfm3_forecasting.py b/examples/timesfm3_forecasting.py new file mode 100644 index 000000000..b7d361de5 --- /dev/null +++ b/examples/timesfm3_forecasting.py @@ -0,0 +1,698 @@ +#!/usr/bin/env python +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Run the five-component TimesFM 3 ONNX forecasting pipeline. + +The example builds the pinned ``google/timesfm-3.0-pytorch`` checkpoint, or +loads an already-saved ModelPackage directory, and runs: + +``raw_preprocessor -> preprocessor -> model -> postprocessor -> stitcher``. + +Only the learned ``model`` component runs on CUDA. The control-flow-heavy +components stay on CPU, leaving one host/device boundary on either side of the +learned core. CUDA inference uses persistent I/O binding and fixed-shape +OrtValues; ``--cuda-graph`` additionally captures that core. + +The official checkpoint weights are subject to the TimesFM Non-Commercial +License v1.0. Review that license before downloading or using the weights. + +Usage:: + + # Build, save, and run the pinned checkpoint on CPU + python examples/timesfm3_forecasting.py --output-dir output/timesfm3 + + # Reuse a saved package and benchmark its learned core on CUDA + python examples/timesfm3_forecasting.py --model-dir output/timesfm3 \ + --device cuda --dtype f16 --cuda-graph --benchmark + + # Compare the ONNX core and end-user forecast with official PyTorch + python examples/timesfm3_forecasting.py --model-dir output/timesfm3 \ + --device cuda --dtype f16 --compare-pytorch --benchmark +""" + +from __future__ import annotations + +import argparse +import csv +import importlib.util +import math +import os +import statistics +import sys +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import onnxruntime as ort + +MODEL_ID = "google/timesfm-3.0-pytorch" +REVISION = "900fcab43d1bfe71733a33b3fec61a41fce28a27" +COMPONENTS = ( + "raw_preprocessor", + "preprocessor", + "model", + "postprocessor", + "stitcher", +) +QUANTILES = tuple(i / 10 for i in range(1, 10)) +ORT_NUMPY_DTYPES = { + "tensor(float)": np.dtype(np.float32), + "tensor(float16)": np.dtype(np.float16), +} + + +def _run_session( + session: ort.InferenceSession, feeds: dict[str, np.ndarray] +) -> dict[str, np.ndarray]: + names = [output.name for output in session.get_outputs()] + return dict(zip(names, session.run(names, feeds))) + + +def _model_path(root: Path, component: str) -> Path: + path = root / component / "model.onnx" + if not path.is_file(): + raise FileNotFoundError( + f"Missing TimesFM component {component!r}: {path}. " + "Expected //model.onnx." + ) + return path + + +def _percentile(samples: list[float], percentile: float) -> float: + """Return a linearly interpolated percentile without extra dependencies.""" + ordered = sorted(samples) + position = (len(ordered) - 1) * percentile + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower) + + +def _latency_summary(samples: list[float]) -> str: + return ( + f"p50/median={statistics.median(samples):.3f} ms, " + f"p95={_percentile(samples, 0.95):.3f} ms" + ) + + +@dataclass(frozen=True) +class ForecastResult: + point: np.ndarray + quantiles: np.ndarray + validity: np.ndarray + model_inputs: np.ndarray + patch_mask: np.ndarray + raw_logits: np.ndarray + component_ms: dict[str, float] + end_to_end_ms: float + + +class TimesFM3OrtDriver: + """Long-lived ORT sessions for a saved five-component TimesFM 3 package.""" + + def __init__( + self, + model_dir: str | os.PathLike[str], + *, + device: str = "cpu", + enable_cuda_graph: bool = False, + ) -> None: + self.model_dir = Path(model_dir) + self.device = device + self.enable_cuda_graph = enable_cuda_graph + self._binding: ort.IOBinding | None = None + self._binding_signature: tuple[tuple[int, ...], tuple[int, ...], str] | None = None + self._host_model_inputs: np.ndarray | None = None + self._host_patch_mask: np.ndarray | None = None + self._input_ortvalues: dict[str, ort.OrtValue] = {} + self._output_ortvalue: ort.OrtValue | None = None + self._run_options: ort.RunOptions | None = None + self._next_graph_id = 0 + + if device == "cuda" and hasattr(ort, "preload_dlls"): + # On Windows this also finds CUDA/cuDNN DLLs bundled with PyTorch or + # NVIDIA site packages, avoiding a dependency on their PATH ordering. + ort.preload_dlls() + torch_spec = importlib.util.find_spec("torch") if sys.platform == "win32" else None + if torch_spec is not None and torch_spec.origin is not None: + torch_lib = Path(torch_spec.origin).parent / "lib" + if torch_lib.is_dir(): + ort.preload_dlls(directory=str(torch_lib)) + available = ort.get_available_providers() + if "CPUExecutionProvider" not in available: + raise RuntimeError( + "CPUExecutionProvider is required for the TimesFM control components; " + f"available providers: {available}." + ) + if device == "cuda" and "CUDAExecutionProvider" not in available: + raise RuntimeError( + "CUDA requested, but CUDAExecutionProvider is unavailable. " + f"Available providers: {available}. Install an ONNX Runtime GPU package " + "compatible with your CUDA installation." + ) + if enable_cuda_graph and device != "cuda": + raise ValueError("--cuda-graph requires --device cuda.") + + self.sessions = { + name: self._create_session(name) + for name in COMPONENTS + if name != "model" or not self.enable_cuda_graph + } + context_meta = next( + value + for value in self.sessions["raw_preprocessor"].get_inputs() + if value.name == "context_values" + ) + try: + self.input_dtype = ORT_NUMPY_DTYPES[context_meta.type] + except KeyError as error: + raise TypeError(f"Unsupported TimesFM input type: {context_meta.type}.") from error + self._validate_session_providers() + + def _create_session( + self, + component: str, + fixed_dimensions: dict[str, int] | None = None, + ) -> ort.InferenceSession: + options = ort.SessionOptions() + for name, value in (fixed_dimensions or {}).items(): + options.add_free_dimension_override_by_name(name, value) + path = str(_model_path(self.model_dir, component)) + if component == "model" and self.device == "cuda": + cuda_options = {"enable_cuda_graph": "1"} if self.enable_cuda_graph else {} + providers: list[str | tuple[str, dict[str, str]]] = [ + ("CUDAExecutionProvider", cuda_options), + "CPUExecutionProvider", + ] + else: + providers = ["CPUExecutionProvider"] + return ort.InferenceSession(path, sess_options=options, providers=providers) + + def _validate_session_providers(self) -> None: + for name, session in self.sessions.items(): + providers = session.get_providers() + expected = ( + "CUDAExecutionProvider" + if name == "model" and self.device == "cuda" + else "CPUExecutionProvider" + ) + if not providers or providers[0] != expected: + raise RuntimeError( + f"{name!r} did not select {expected}; active providers: {providers}." + ) + + @staticmethod + def _timed_cpu( + session: ort.InferenceSession, feeds: dict[str, np.ndarray] + ) -> tuple[dict[str, np.ndarray], float]: + start = time.perf_counter() + outputs = _run_session(session, feeds) + return outputs, (time.perf_counter() - start) * 1_000 + + def _rebuild_cuda_binding(self, model_inputs: np.ndarray, patch_mask: np.ndarray) -> None: + model_inputs = np.ascontiguousarray(model_inputs) + patch_mask = np.ascontiguousarray(patch_mask) + if self.enable_cuda_graph: + batch, variates, patches = model_inputs.shape[:3] + self.sessions["model"] = self._create_session( + "model", + { + "component.model.batch": batch, + "component.model.variates": variates, + "component.model.patches": patches, + }, + ) + self._host_model_inputs = np.empty_like(model_inputs) + self._host_patch_mask = np.empty_like(patch_mask) + np.copyto(self._host_model_inputs, model_inputs) + np.copyto(self._host_patch_mask, patch_mask) + + self._input_ortvalues = { + "model_inputs": ort.OrtValue.ortvalue_from_numpy( + self._host_model_inputs, "cuda", 0 + ), + "patch_mask": ort.OrtValue.ortvalue_from_numpy(self._host_patch_mask, "cuda", 0), + } + output_meta = self.sessions["model"].get_outputs()[0] + output_width = output_meta.shape[-1] + if not isinstance(output_width, int): + raise TypeError(f"Expected a fixed raw_logits width, got {output_meta.shape}.") + output_shape = (*model_inputs.shape[:3], output_width) + self._output_ortvalue = ort.OrtValue.ortvalue_from_shape_and_type( + output_shape, model_inputs.dtype, "cuda", 0 + ) + + binding = self.sessions["model"].io_binding() + for name, value in self._input_ortvalues.items(): + binding.bind_ortvalue_input(name, value) + binding.bind_ortvalue_output("raw_logits", self._output_ortvalue) + self._binding = binding + + # A captured graph cannot change addresses or shapes. Give each new + # fixed-shape binding its own graph ID rather than replaying a stale graph. + self._run_options = None + if self.enable_cuda_graph: + self._run_options = ort.RunOptions() + self._run_options.add_run_config_entry("gpu_graph_id", str(self._next_graph_id)) + self._next_graph_id += 1 + + def _run_cuda_model(self, model_inputs: np.ndarray, patch_mask: np.ndarray) -> np.ndarray: + model_inputs = np.ascontiguousarray(model_inputs) + patch_mask = np.ascontiguousarray(patch_mask) + signature = (model_inputs.shape, patch_mask.shape, model_inputs.dtype.str) + if signature != self._binding_signature: + self._rebuild_cuda_binding(model_inputs, patch_mask) + self._binding_signature = signature + else: + assert self._host_model_inputs is not None + assert self._host_patch_mask is not None + np.copyto(self._host_model_inputs, model_inputs) + np.copyto(self._host_patch_mask, patch_mask) + # Both fixed-shape inputs cross the pipeline's sole host-to-device boundary. + self._input_ortvalues["model_inputs"].update_inplace(self._host_model_inputs) + self._input_ortvalues["patch_mask"].update_inplace(self._host_patch_mask) + + assert self._binding is not None + self._binding.synchronize_inputs() + self.sessions["model"].run_with_iobinding(self._binding, self._run_options) + self._binding.synchronize_outputs() + # This is the pipeline's sole device-to-host boundary. + return self._binding.copy_outputs_to_cpu()[0] + + def run_model( + self, model_inputs: np.ndarray, patch_mask: np.ndarray + ) -> tuple[np.ndarray, float]: + start = time.perf_counter() + if self.device == "cuda": + output = self._run_cuda_model(model_inputs, patch_mask) + else: + output = _run_session( + self.sessions["model"], + {"model_inputs": model_inputs, "patch_mask": patch_mask}, + )["raw_logits"] + return output, (time.perf_counter() - start) * 1_000 + + def benchmark_model( + self, + model_inputs: np.ndarray, + patch_mask: np.ndarray, + *, + warmups: int, + iterations: int, + ) -> list[float]: + """Time only learned-core execution, excluding host/device copies.""" + samples: list[float] = [] + if self.device == "cuda": + self._run_cuda_model(model_inputs, patch_mask) + assert self._binding is not None + for _ in range(warmups): + self.sessions["model"].run_with_iobinding(self._binding, self._run_options) + self._binding.synchronize_outputs() + for _ in range(iterations): + self._binding.synchronize_outputs() + start = time.perf_counter() + self.sessions["model"].run_with_iobinding(self._binding, self._run_options) + self._binding.synchronize_outputs() + samples.append((time.perf_counter() - start) * 1_000) + else: + feeds = {"model_inputs": model_inputs, "patch_mask": patch_mask} + for _ in range(warmups): + _run_session(self.sessions["model"], feeds) + for _ in range(iterations): + start = time.perf_counter() + _run_session(self.sessions["model"], feeds) + samples.append((time.perf_counter() - start) * 1_000) + return samples + + def forecast( + self, + context: np.ndarray, + horizon: int, + *, + make_positive: bool = False, + ) -> ForecastResult: + """Forecast one or more variates shaped ``[variates, context]``.""" + context = np.asarray(context) + if context.ndim == 1: + context = context[None, :] + if context.ndim != 2 or context.shape[-1] == 0: + raise ValueError("context must have shape [context] or [variates, context].") + if horizon <= 0: + raise ValueError("horizon must be positive.") + + context = np.ascontiguousarray(context[None, ...]) + batch, variates, context_length = context.shape + raw_feeds = { + "context_values": np.nan_to_num(context, nan=0.0), + "context_observed": np.isfinite(context), + "future_values": np.zeros((batch, variates, horizon), dtype=context.dtype), + "future_observed": np.zeros((batch, variates, horizon), dtype=np.bool_), + "context_lengths": np.full((batch,), context_length, dtype=np.int64), + "horizon_lengths": np.full((batch,), horizon, dtype=np.int64), + "variate_roles": np.zeros((batch, variates), dtype=np.int64), + } + + total_start = time.perf_counter() + raw, raw_ms = self._timed_cpu(self.sessions["raw_preprocessor"], raw_feeds) + pre, pre_ms = self._timed_cpu( + self.sessions["preprocessor"], + { + name: raw[name] + for name in ("values", "masks", "patch_is_target", "patch_cpm_mask") + }, + ) + raw_logits, model_ms = self.run_model(pre["model_inputs"], pre["patch_mask"]) + post, post_ms = self._timed_cpu( + self.sessions["postprocessor"], + { + "raw_logits": raw_logits, + "revin_count": pre["revin_count"], + "revin_mean": pre["revin_mean"], + "revin_std": pre["revin_std"], + "patch_cpm_mask": raw["patch_cpm_mask"], + }, + ) + stitch, stitch_ms = self._timed_cpu( + self.sessions["stitcher"], + { + "logits": post["logits"], + "make_positive": np.asarray(make_positive, dtype=np.bool_), + **{ + name: raw[name] + for name in ( + "trend_slope", + "trend_intercept", + "apply_detrend", + "target_mask", + "nonnegative_mask", + "context_lengths", + "horizon_lengths", + "context_patch_count", + "forecast_patch_counts", + ) + }, + }, + ) + end_to_end_ms = (time.perf_counter() - total_start) * 1_000 + return ForecastResult( + point=stitch["point_forecast"], + quantiles=stitch["quantile_forecasts"], + validity=stitch["validity"], + model_inputs=pre["model_inputs"], + patch_mask=pre["patch_mask"], + raw_logits=raw_logits, + component_ms={ + "raw_preprocessor": raw_ms, + "preprocessor": pre_ms, + "model": model_ms, + "postprocessor": post_ms, + "stitcher": stitch_ms, + }, + end_to_end_ms=end_to_end_ms, + ) + + +def _demo_signal(context_length: int, dtype: np.dtype[Any]) -> np.ndarray: + index = np.arange(context_length, dtype=np.float32) + signal = 0.025 * index + 1.4 * np.sin(2 * np.pi * index / 24) + signal += 0.35 * np.cos(2 * np.pi * index / 7) + return signal.astype(dtype) + + +def _build_package(args: argparse.Namespace) -> Path: + if args.model_dir is not None: + root = Path(args.model_dir) + for component in COMPONENTS: + _model_path(root, component) + print(f"Loading saved ModelPackage from {root}") + return root + + from mobius import build + + root = Path(args.output_dir) + print(f"Building {args.model_id!r} at revision {args.revision} ...") + package = build( + args.model_id, + revision=args.revision, + dtype=args.dtype, + execution_provider=args.device, + load_weights=True, + ) + if set(package) != set(COMPONENTS): + raise RuntimeError( + f"Expected TimesFM components {list(COMPONENTS)}, got {list(package)}." + ) + package.save(str(root), external_data="onnx") + print(f"Saved ModelPackage to {root}") + return root + + +def _benchmark( + driver: TimesFM3OrtDriver, + signal: np.ndarray, + horizon: int, + warmups: int, + iterations: int, +) -> ForecastResult: + for _ in range(warmups): + driver.forecast(signal, horizon) + + end_to_end: list[float] = [] + by_component = {name: [] for name in COMPONENTS} + result: ForecastResult | None = None + for _ in range(iterations): + result = driver.forecast(signal, horizon) + end_to_end.append(result.end_to_end_ms) + for name, latency in result.component_ms.items(): + by_component[name].append(latency) + + assert result is not None + core_samples = driver.benchmark_model( + result.model_inputs, + result.patch_mask, + warmups=warmups, + iterations=iterations, + ) + print("\nONNX Runtime steady-state latency") + print(f" end-to-end: {_latency_summary(end_to_end)}") + for name in COMPONENTS: + label = "model + copies" if name == "model" and driver.device == "cuda" else name + print(f" {label + ':':18}{_latency_summary(by_component[name])}") + print(f" {'model core-only:':18}{_latency_summary(core_samples)}") + return result + + +def _compare_pytorch( + result: ForecastResult, + signal: np.ndarray, + args: argparse.Namespace, +) -> None: + try: + import timesfm3 + import torch + except ImportError as error: + raise RuntimeError( + "--compare-pytorch requires the official TimesFM package. " + 'Install it with: pip install "timesfm[torch]"' + ) from error + + # The pinned upstream implementation's RoPE path promotes Q/K to float32 while V + # remains float16, which PyTorch SDPA rejects. Use its supported float32 path when + # comparing an fp16 ONNX package and report the difference explicitly. + torch_dtype = torch.float32 + if result.model_inputs.dtype == np.float16: + print("PyTorch comparison uses float32 (upstream fp16 RoPE/SDPA is unsupported).") + print("\nLoading official PyTorch checkpoint for comparison ...") + forecaster = timesfm3.TimesFM3Forecaster.from_pretrained( + args.model_id, + device=args.device, + revision=args.revision, + per_core_batch_size=1, + ) + model = forecaster.model.to(device=args.device, dtype=torch_dtype).eval() + + model_inputs = torch.from_numpy(result.model_inputs).to( + device=args.device, dtype=torch_dtype + ) + patch_mask = torch.from_numpy(result.patch_mask).to(device=args.device) + + def synchronize() -> None: + if args.device == "cuda": + torch.cuda.synchronize() + + @torch.inference_mode() + def run_core() -> Any: + hidden = model.pre_transformer_resblock(model_inputs) + hidden, _, _ = model.transformer_stack(hidden, patch_mask) + return model.output_head(hidden) + + for _ in range(args.warmups): + core_output = run_core() + synchronize() + core_samples: list[float] = [] + for _ in range(args.iterations): + synchronize() + start = time.perf_counter() + core_output = run_core() + synchronize() + core_samples.append((time.perf_counter() - start) * 1_000) + core_numpy = core_output.float().cpu().numpy() + core_error = float(np.max(np.abs(result.raw_logits.astype(np.float32) - core_numpy))) + print(f"PyTorch learned-core max absolute error: {core_error:.6g}") + print(f"PyTorch learned-core latency: {_latency_summary(core_samples)}") + + target = torch.from_numpy(signal[None, None, :]).to(device=args.device, dtype=torch_dtype) + + @torch.inference_mode() + def run_forecast() -> Any: + return model.decode(target=target, horizon=args.horizon) + + try: + for _ in range(args.warmups): + torch_forecast = run_forecast() + synchronize() + forecast_samples: list[float] = [] + for _ in range(args.iterations): + synchronize() + start = time.perf_counter() + torch_forecast = run_forecast() + synchronize() + forecast_samples.append((time.perf_counter() - start) * 1_000) + torch_quantiles = np.sort(torch_forecast.float().cpu().numpy(), axis=-1) + onnx_quantiles = result.quantiles[:, :, : args.horizon].astype(np.float32) + forecast_error = float(np.max(np.abs(onnx_quantiles - torch_quantiles))) + print(f"PyTorch user-forecast max absolute error: {forecast_error:.6g}") + print( + f"PyTorch device-resident forecast latency: {_latency_summary(forecast_samples)}" + ) + except (RuntimeError, TypeError) as error: + print(f"PyTorch user-forecast comparison unavailable: {error}") + + +def _print_forecast(result: ForecastResult, horizon: int) -> None: + point = result.point[0, 0, :horizon] + quantiles = result.quantiles[0, 0, :horizon] + middle = quantiles.shape[-1] // 2 + print("\nForecast (first target)") + labels = ( + ("q10", "q50", "q90") + if quantiles.shape[-1] == 9 + else ( + "first", + "middle", + "last", + ) + ) + print(f" step point {labels[0]:>12} {labels[1]:>12} {labels[2]:>12}") + for step in range(horizon): + print( + f"{step + 1:5d} {point[step]:12.5f} {quantiles[step, 0]:12.5f} " + f"{quantiles[step, middle]:12.5f} {quantiles[step, -1]:12.5f}" + ) + + +def _write_csv(path: str | os.PathLike[str], result: ForecastResult, horizon: int) -> None: + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + count = result.quantiles.shape[-1] + levels = QUANTILES if count == len(QUANTILES) else tuple(np.linspace(0.1, 0.9, count)) + with output.open("w", newline="", encoding="utf-8") as file: + writer = csv.writer(file) + writer.writerow(["step", "point", *(f"q{quantile:g}" for quantile in levels)]) + for step in range(horizon): + writer.writerow( + [ + step + 1, + float(result.point[0, 0, step]), + *(float(value) for value in result.quantiles[0, 0, step]), + ] + ) + print(f"Wrote forecast CSV to {output}") + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Build or load the five-component TimesFM 3 ONNX package and forecast a " + "deterministic seasonal signal. Official weights use the non-commercial " + "TimesFM license." + ), + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--model-id", default=MODEL_ID, help="Hugging Face model ID.") + parser.add_argument( + "--revision", default=REVISION, help="Pinned Hugging Face checkpoint revision." + ) + parser.add_argument( + "--model-dir", + help="Previously saved ModelPackage directory; skips build when supplied.", + ) + parser.add_argument( + "--output-dir", + default="output/timesfm3", + help="Directory for a newly built ModelPackage.", + ) + parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu") + parser.add_argument("--dtype", choices=["f32", "f16"], default="f32") + parser.add_argument("--context-length", type=int, default=512) + parser.add_argument("--horizon", type=int, default=64) + parser.add_argument("--warmups", type=int, default=3) + parser.add_argument("--iterations", type=int, default=10) + parser.add_argument( + "--cuda-graph", + action="store_true", + help="Capture the fixed-shape CUDA model core (CUDA only).", + ) + parser.add_argument( + "--benchmark", action="store_true", help="Report steady-state ORT latency." + ) + parser.add_argument( + "--compare-pytorch", + action="store_true", + help="Compare the same pinned weights, core tensors, and forecast with PyTorch.", + ) + parser.add_argument("--csv-output", help="Optional path for forecast CSV output.") + return parser + + +def main() -> None: + parser = _parser() + args = parser.parse_args() + if args.context_length <= 0: + parser.error("--context-length must be positive.") + if args.horizon <= 0: + parser.error("--horizon must be positive.") + if args.warmups < 0: + parser.error("--warmups must be non-negative.") + if args.iterations <= 0: + parser.error("--iterations must be positive.") + if args.cuda_graph and args.device != "cuda": + parser.error("--cuda-graph requires --device cuda.") + + model_dir = _build_package(args) + driver = TimesFM3OrtDriver( + model_dir, + device=args.device, + enable_cuda_graph=args.cuda_graph, + ) + signal = _demo_signal(args.context_length, driver.input_dtype) + + if args.benchmark: + result = _benchmark(driver, signal, args.horizon, args.warmups, args.iterations) + else: + result = driver.forecast(signal, args.horizon) + print(f"End-to-end latency: {result.end_to_end_ms:.3f} ms") + + _print_forecast(result, args.horizon) + if args.csv_output: + _write_csv(args.csv_output, result, args.horizon) + if args.compare_pytorch: + _compare_pytorch(result, signal, args) + + +if __name__ == "__main__": + main() diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index 51e0e5d4e..2cc9ace5b 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -156,6 +156,8 @@ SmallThinkerGGUFCausalLMModel, SmolLM3CausalLMModel, SortformerDiarizationModel, + TimesFM3Config, + TimesFM3Model, WhisperForConditionalGeneration, XverseCausalLMModel, ) @@ -1068,6 +1070,13 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: ), "fastconformer_rnnt": ModelRegistration(EncDecRNNTModel, task="fastconformer-rnnt"), "sortformer": ModelRegistration(SortformerDiarizationModel, task="diarization"), + "timesfm3": ModelRegistration( + TimesFM3Model, + task="time-series-forecasting", + config_class=TimesFM3Config, + family="timesfm", + variant="3.0", + ), "reuse": ModelRegistration( SEMambaSpeechEnhancementModel, task="speech-enhancement", diff --git a/src/mobius/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index a3c1740c3..731805dba 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -438,7 +438,7 @@ def build_transformers_model( model.metadata_props["mobius.source_revision"] = revision or "unpinned" if load_weights: - if config.block_quant_scheme is not None and hasattr( + if getattr(config, "block_quant_scheme", None) is not None and hasattr( model_module, "build_fp8_streaming_plan" ): if len(package) != 1: diff --git a/src/mobius/integrations/transformers/_config_resolver.py b/src/mobius/integrations/transformers/_config_resolver.py index c16ea0351..227d32df0 100644 --- a/src/mobius/integrations/transformers/_config_resolver.py +++ b/src/mobius/integrations/transformers/_config_resolver.py @@ -122,11 +122,14 @@ def _try_load_config_json(model_id: str, revision: str | None = None): model_type = config_dict.get("model_type") if not model_type: - model_type = _model_type_from_architectures(config_dict.get("architectures")) + model_type = _model_type_from_architectures( + config_dict.get("architectures") + ) or _model_type_from_config_signature(config_dict) if not model_type: return None logger.info( - "config.json for %s declares no model_type; inferred '%s' from architectures=%s", + "config.json for %s declares no model_type; inferred '%s' " + "from architectures=%s or its config schema", model_id, model_type, config_dict.get("architectures"), @@ -136,6 +139,24 @@ def _try_load_config_json(model_id: str, revision: str | None = None): return _dict_to_pretrained_config(config_dict) +def _model_type_from_config_signature(config: dict) -> str | None: + """Identify non-Transformers checkpoints with an unambiguous config schema.""" + transformer = config.get("transformer_config") + residual = config.get("residual_block_config") + if ( + config.get("input_patch_len") is not None + and config.get("output_patch_len") is not None + and config.get("quantiles") is not None + and config.get("use_iterative_cpm_revin") is not None + and config.get("use_variate_attention") is not None + and isinstance(transformer, dict) + and isinstance(transformer.get("transformer"), dict) + and isinstance(residual, dict) + ): + return "timesfm3" + return None + + def _model_type_from_architectures(architectures) -> str | None: """Recover a HuggingFace ``model_type`` from a config's ``architectures``. diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 40dfb2cda..413b93d48 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -180,6 +180,8 @@ "SEMambaSpeechEnhancementModel", "SortformerConfig", "SortformerDiarizationModel", + "TimesFM3Config", + "TimesFM3Model", "Qwen3TTSCodePredictorModel", "Qwen3TTSCodecDecoderModel", "Qwen3TTSCodecEncoderModel", @@ -442,6 +444,7 @@ from mobius.models.starcoder2 import StarCoder2CausalLMModel from mobius.models.t5 import T5EncoderModel, T5ForConditionalGeneration from mobius.models.talkie import TalkieForCausalLM +from mobius.models.timesfm3 import TimesFM3Config, TimesFM3Model from mobius.models.unet import ( UNet2DConditionModel, load_unet_lora_safetensors, diff --git a/src/mobius/models/timesfm3.py b/src/mobius/models/timesfm3.py new file mode 100644 index 000000000..ddb64ad98 --- /dev/null +++ b/src/mobius/models/timesfm3.py @@ -0,0 +1,1581 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""TimesFM 3 multivariate time-series forecasting model. + +Replicates Google Research's ``TimesFM3Torch`` as a five-stage ONNX pipeline: +raw-series preparation, patched feature construction, transformer inference, +CPM/RevIN postprocessing, and forecast stitching. +""" + +from __future__ import annotations + +import dataclasses +import math +from collections.abc import Mapping + +import onnx_ir as ir +from onnxscript import OpBuilder, nn + +from mobius._configs import BaseModelConfig +from mobius.components import Linear +from mobius.components._scan_utils import create_body_graph, rename_subgraph_values + +_INT64_MAX = 9223372036854775807 +_PER_DIM_SCALE = 1.442695041 +_REVIN_TOLERANCE = 1e-6 +_FLOAT32_EPS = 1.1920928955078125e-7 +_TARGET_ROLE = 0 +_PAST_ONLY_ROLE = 1 +_PAST_FUTURE_ROLE = 2 + + +def _field(value, name: str, default): + if isinstance(value, Mapping): + return value.get(name, default) + return getattr(value, name, default) + + +@dataclasses.dataclass +class TimesFM3Config(BaseModelConfig): + """Configuration parsed from a TimesFM 3 ``config.json``.""" + + input_patch_len: int = 32 + output_patch_len: int = 64 + quantiles: tuple[float, ...] = tuple(i / 10 for i in range(1, 10)) + num_layers: int = 20 + model_dims: int = 1280 + transformer_hidden_dims: int = 1280 + num_heads: int = 16 + max_variates: int = 32 + rms_norm_eps: float = _FLOAT32_EPS + use_variate_attention: bool = True + use_rope_seq: bool = True + use_rope_var: bool = False + value_clip: float = 1e20 + use_iterative_cpm_revin: bool = True + use_linear_detrending: bool = True + linear_detrending_threshold: float = 0.5 + model_type: str | None = "timesfm3" + + def __post_init__(self) -> None: + self.hidden_size = self.model_dims + self.intermediate_size = self.transformer_hidden_dims + self.num_hidden_layers = self.num_layers + self.num_attention_heads = self.num_heads + self.num_key_value_heads = self.num_heads + self.head_dim = self.model_dims // self.num_heads + + def validate(self) -> None: + if ( + not isinstance(self.input_patch_len, int) + or isinstance(self.input_patch_len, bool) + or self.input_patch_len <= 0 + or not isinstance(self.output_patch_len, int) + or isinstance(self.output_patch_len, bool) + or self.output_patch_len <= 0 + ): + raise ValueError("input_patch_len and output_patch_len must be positive integers") + if self.output_patch_len % self.input_patch_len: + raise ValueError("output_patch_len must be a multiple of input_patch_len") + if self.model_dims % self.num_heads: + raise ValueError("model_dims must be divisible by num_heads") + if self.head_dim % 2: + raise ValueError("TimesFM 3 RoPE requires an even head_dim") + if not self.quantiles: + raise ValueError("quantiles must not be empty") + if self.output_patch_len <= self.input_patch_len: + raise ValueError("TimesFM 3 stitching requires output_patch_len > input_patch_len") + + @classmethod + def from_transformers(cls, config) -> TimesFM3Config: + residual = _field(config, "residual_block_config", {}) + stack = _field(config, "transformer_config", {}) + transformer = _field(stack, "transformer", {}) + model_dims = int( + _field(transformer, "model_dims", _field(residual, "output_dims", 1280)) + ) + return cls( + input_patch_len=int(_field(config, "input_patch_len", 32)), + output_patch_len=int(_field(config, "output_patch_len", 64)), + quantiles=tuple(float(q) for q in _field(config, "quantiles", (0.5,))), + num_layers=int(_field(stack, "num_layers", 20)), + model_dims=model_dims, + transformer_hidden_dims=int(_field(transformer, "hidden_dims", model_dims)), + num_heads=int(_field(transformer, "num_heads", 16)), + max_variates=int(_field(transformer, "max_variates", 32)), + rms_norm_eps=_FLOAT32_EPS, + use_variate_attention=bool(_field(config, "use_variate_attention", True)), + use_rope_seq=bool(_field(transformer, "use_rope_seq", True)), + use_rope_var=bool(_field(transformer, "use_rope_var", False)), + value_clip=float(_field(config, "value_clip", 1e20)), + use_iterative_cpm_revin=bool(_field(config, "use_iterative_cpm_revin", True)), + use_linear_detrending=bool(_field(config, "use_linear_detrending", True)), + linear_detrending_threshold=float( + _field(config, "linear_detrending_threshold", 0.5) + ), + model_type="timesfm3", + ) + + +def _safe_divisor(op: OpBuilder, value: ir.Value) -> ir.Value: + return op.Where( + op.Less(value, op.CastLike(op.Constant(value_float=_REVIN_TOLERANCE), value)), + op.CastLike(op.Constant(value_float=1.0), value), + value, + ) + + +def _update_running_stats( + op: OpBuilder, + count: ir.Value, + mean: ir.Value, + std: ir.Value, + values: ir.Value, + masks: ir.Value, +) -> tuple[ir.Value, ir.Value, ir.Value]: + valid = op.Not(masks) + valid_f = op.Cast(valid, to=ir.DataType.FLOAT) + increment_count = op.ReduceSum(valid_f, axes=[-1], keepdims=False) + safe_increment_count = op.Max(increment_count, op.Constant(value_float=1.0)) + + valid_values = op.Where(valid, values, op.CastLike(op.Constant(value_float=0.0), values)) + increment_sum = op.ReduceSum(valid_values, axes=[-1], keepdims=False) + increment_mean = op.Where( + op.Equal(increment_count, op.Constant(value_float=0.0)), + op.Mul(increment_sum, op.Constant(value_float=0.0)), + op.Div(increment_sum, safe_increment_count), + ) + + centered = op.Sub(values, op.Unsqueeze(increment_mean, axes=[-1])) + centered_sq = op.Where( + valid, + op.Mul(centered, centered), + op.CastLike(op.Constant(value_float=0.0), values), + ) + increment_variance = op.Div( + op.ReduceSum(centered_sq, axes=[-1], keepdims=False), + safe_increment_count, + ) + new_count = op.Add(count, increment_count) + safe_new_count = op.Max(new_count, op.Constant(value_float=1.0)) + new_mean = op.Div( + op.Add(op.Mul(count, mean), op.Mul(increment_count, increment_mean)), + safe_new_count, + ) + merged_variance = op.Div( + op.Add( + op.Add( + op.Mul(count, op.Mul(std, std)), op.Mul(increment_count, increment_variance) + ), + op.Add( + op.Mul(count, op.Mul(op.Sub(mean, new_mean), op.Sub(mean, new_mean))), + op.Mul( + increment_count, + op.Mul( + op.Sub(increment_mean, new_mean), + op.Sub(increment_mean, new_mean), + ), + ), + ), + ), + safe_new_count, + ) + return new_count, new_mean, op.Sqrt(op.Max(merged_variance, op.Constant(value_float=0.0))) + + +def _running_stats_body() -> ir.Graph: + float_type = ir.TensorType(ir.DataType.FLOAT) + bool_type = ir.TensorType(ir.DataType.BOOL) + count = ir.Value(name="count", type=float_type) + mean = ir.Value(name="mean", type=float_type) + std = ir.Value(name="std", type=float_type) + values = ir.Value(name="values", type=float_type) + masks = ir.Value(name="masks", type=bool_type) + graph, builder = create_body_graph( + state_inputs=[count, mean, std], + scan_inputs=[values, masks], + name="timesfm3_running_stats", + ) + new_count, new_mean, new_std = _update_running_stats( + builder.op, count, mean, std, values, masks + ) + # Scan carry and per-step outputs must be distinct graph values. Use a + # data-dependent copy because no-op Identities are removed during optimization. + nonnegative_count = builder.op.GreaterOrEqual( + new_count, builder.op.Constant(value_float=0.0) + ) + scan_count = builder.op.Where( + nonnegative_count, new_count, builder.op.Constant(value_float=0.0) + ) + scan_mean = builder.op.Where( + nonnegative_count, new_mean, builder.op.Constant(value_float=0.0) + ) + scan_std = builder.op.Where( + nonnegative_count, new_std, builder.op.Constant(value_float=0.0) + ) + graph.outputs.extend([new_count, new_mean, new_std, scan_count, scan_mean, scan_std]) + rename_subgraph_values(graph, "timesfm3_stats_") + return graph + + +def _get_running_stats( + op: OpBuilder, values: ir.Value, masks: ir.Value +) -> tuple[ir.Value, ir.Value, ir.Value]: + # Scan patch-major inputs so the recurrence exactly matches upstream's + # population-variance merge, including its numerical evaluation order. + patch_values = op.Transpose(op.Cast(values, to=ir.DataType.FLOAT), perm=[2, 0, 1, 3]) + patch_masks = op.Transpose(masks, perm=[2, 0, 1, 3]) + bv_shape = op.Shape(values, start=0, end=2) + zeros = op.ConstantOfShape(bv_shape, value=ir.tensor([0.0], dtype=ir.DataType.FLOAT)) + _, _, _, count, mean, std = op.Scan( + zeros, + zeros, + zeros, + patch_values, + patch_masks, + body=_running_stats_body(), + num_scan_inputs=2, + _outputs=6, + ) + return ( + op.Transpose(count, perm=[1, 2, 0]), + op.Transpose(mean, perm=[1, 2, 0]), + op.Transpose(std, perm=[1, 2, 0]), + ) + + +def _nearest_observed_body() -> ir.Graph: + """Scan body carrying the nearest valid index seen so far.""" + int_type = ir.TensorType(ir.DataType.INT64) + bool_type = ir.TensorType(ir.DataType.BOOL) + previous = ir.Value(name="previous", type=int_type) + valid = ir.Value(name="valid", type=bool_type) + index = ir.Value(name="index", type=int_type) + graph, builder = create_body_graph( + state_inputs=[previous], + scan_inputs=[valid, index], + name="timesfm3_nearest_observed", + ) + current = builder.op.Where(valid, index, previous) + # Scan carry and scan output must be distinct values. + scanned = builder.op.Where( + builder.op.GreaterOrEqual(current, builder.op.Constant(value_int=-1)), + current, + builder.op.Constant(value_int=-1), + ) + graph.outputs.extend([current, scanned]) + rename_subgraph_values(graph, "timesfm3_nearest_") + return graph + + +def _interpolate_missing( + op: OpBuilder, + values: ir.Value, + observed: ir.Value, +) -> ir.Value: + """Match ``numpy.interp`` independently over the last axis.""" + values_f32 = op.Cast(values, to=ir.DataType.FLOAT) + length = op.Squeeze(op.Shape(values, start=2, end=3)) + indices = op.Range( + op.Constant(value_int=0), + length, + op.Constant(value_int=1), + ) + valid_t = op.Transpose(observed, perm=[2, 0, 1]) + leading_shape = op.Shape(values, start=0, end=2) + previous_init = op.ConstantOfShape( + leading_shape, value=ir.tensor([-1], dtype=ir.DataType.INT64) + ) + _, previous_t = op.Scan( + previous_init, + valid_t, + indices, + body=_nearest_observed_body(), + num_scan_inputs=2, + _outputs=2, + ) + + reverse_indices = op.Range( + op.Sub(length, op.Constant(value_int=1)), + op.Constant(value_int=-1), + op.Constant(value_int=-1), + ) + next_init = op.ConstantOfShape( + leading_shape, value=ir.tensor([0], dtype=ir.DataType.INT64) + ) + next_init = op.Add(next_init, length) + _, next_reversed_t = op.Scan( + next_init, + op.Gather(valid_t, reverse_indices, axis=0), + reverse_indices, + body=_nearest_observed_body(), + num_scan_inputs=2, + _outputs=2, + ) + next_t = op.Gather(next_reversed_t, reverse_indices, axis=0) + + previous = op.Transpose(previous_t, perm=[1, 2, 0]) + following = op.Transpose(next_t, perm=[1, 2, 0]) + zero = op.Constant(value_int=0) + last = op.Sub(length, op.Constant(value_int=1)) + previous_values = op.GatherElements( + values_f32, op.Min(op.Max(previous, zero), last), axis=2 + ) + following_values = op.GatherElements( + values_f32, op.Min(op.Max(following, zero), last), axis=2 + ) + + has_previous = op.GreaterOrEqual(previous, zero) + has_following = op.Less(following, length) + both = op.And(has_previous, has_following) + span = op.Cast(op.Sub(following, previous), to=ir.DataType.FLOAT) + safe_span = op.Where( + op.Equal(span, op.CastLike(op.Constant(value_float=0.0), span)), + op.CastLike(op.Constant(value_float=1.0), span), + span, + ) + position = op.Cast( + op.Sub(op.Unsqueeze(indices, axes=[0, 1]), previous), + to=ir.DataType.FLOAT, + ) + interpolated = op.Add( + previous_values, + op.Mul( + op.Div(position, safe_span), + op.Sub(following_values, previous_values), + ), + ) + filled = op.Where( + both, + interpolated, + op.Where( + has_previous, + previous_values, + op.Where( + has_following, + following_values, + op.Constant(value_float=0.0), + ), + ), + ) + return op.CastLike(op.Where(observed, values_f32, filled), values) + + +def _roll_patches( + op: OpBuilder, values: ir.Value, rolls: int, patch_len: int +) -> tuple[ir.Value, ir.Value]: + shifted = values + outputs = [] + for _ in range(rolls): + shifted = op.Concat( + op.Slice(shifted, starts=[1], ends=[_INT64_MAX], axes=[2]), + op.Slice(shifted, starts=[0], ends=[1], axes=[2]), + axis=2, + ) + outputs.append(shifted) + rolled = op.Concat(*outputs, axis=-1) + + num_patches = op.Shape(values, start=2, end=3) + patch_index = op.Range( + op.Constant(value_int=0), op.Squeeze(num_patches), op.Constant(value_int=1) + ) + output_index = op.Constant(value_ints=list(range(rolls * patch_len))) + source_patch = op.Add( + op.Add( + op.Unsqueeze(patch_index, axes=[1]), + op.Constant(value_int=1), + ), + op.Div( + op.Unsqueeze(output_index, axes=[0]), + op.Constant(value_int=patch_len), + ), + ) + wrap_mask = op.GreaterOrEqual(source_patch, num_patches) + return rolled, op.Unsqueeze(wrap_mask, axes=[0, 1]) + + +class _ResidualBlock(nn.Module): + def __init__(self, input_dims: int, hidden_dims: int, output_dims: int): + super().__init__() + self.hidden_layer = Linear(input_dims, hidden_dims, bias=False) + self.output_layer = Linear(hidden_dims, output_dims, bias=False) + self.residual_layer = Linear(input_dims, output_dims, bias=False) + + def forward(self, op: OpBuilder, values: ir.Value) -> ir.Value: + hidden = op.Relu(self.hidden_layer(op, values)) + return op.Add(self.output_layer(op, hidden), self.residual_layer(op, values)) + + +class _PerDimScale(nn.Module): + def __init__(self, head_dim: int): + super().__init__() + self.per_dim_scale = nn.Parameter([head_dim]) + self._factor = _PER_DIM_SCALE / math.sqrt(head_dim) + + def forward(self, op: OpBuilder, values: ir.Value) -> ir.Value: + scale = op.Mul(op.Softplus(self.per_dim_scale), self._factor) + return op.Mul(values, op.CastLike(scale, values)) + + +class _TimesFMRMSNorm(nn.Module): + def __init__(self, dimensions: int, eps: float): + super().__init__() + self.weight = nn.Parameter([dimensions]) + self._eps = eps + + def forward(self, op: OpBuilder, values: ir.Value) -> ir.Value: + values_f32 = op.Cast(values, to=ir.DataType.FLOAT) + variance = op.ReduceMean( + op.Mul(values_f32, values_f32), + axes=[-1], + keepdims=True, + ) + normalized = op.Mul( + values_f32, + op.Reciprocal( + op.Sqrt( + op.Add( + variance, + op.CastLike(op.Constant(value_float=self._eps), variance), + ) + ) + ), + ) + output = op.Mul(normalized, op.CastLike(self.weight, normalized)) + return op.CastLike(output, values) + + +class _TimesFMAttention(nn.Module): + def __init__(self, config: TimesFM3Config, *, causal: bool, use_rope: bool): + super().__init__() + self.query_proj = Linear(config.model_dims, config.model_dims, bias=False) + self.key_proj = Linear(config.model_dims, config.model_dims, bias=False) + self.value_proj = Linear(config.model_dims, config.model_dims, bias=False) + self.out_proj = Linear(config.model_dims, config.model_dims, bias=False) + self.query_ln = _TimesFMRMSNorm(config.head_dim, config.rms_norm_eps) + self.key_ln = _TimesFMRMSNorm(config.head_dim, config.rms_norm_eps) + self.per_dim_scale = _PerDimScale(config.head_dim) + self._num_heads = config.num_heads + self._head_dim = config.head_dim + self._model_dims = config.model_dims + self._causal = causal + self._use_rope = use_rope + self._mask_value = -65504.0 if config.dtype == ir.DataType.FLOAT16 else -1e9 + + def _rope(self, op: OpBuilder, values: ir.Value) -> ir.Value: + half_dim = self._head_dim // 2 + timescale = [10000.0 ** (2.0 * i / self._head_dim) for i in range(half_dim)] + seq_len = op.Squeeze(op.Shape(values, start=1, end=2)) + positions = op.Cast( + op.Range( + op.Constant(value_int=0), + seq_len, + op.Constant(value_int=1), + ), + to=ir.DataType.FLOAT, + ) + angles = op.Div( + op.Unsqueeze(positions, axes=[0, 2, 3]), + op.Constant(value_floats=timescale), + ) + sin = op.Sin(angles) + cos = op.Cos(angles) + first = op.Slice(values, starts=[0], ends=[half_dim], axes=[-1]) + second = op.Slice(values, starts=[half_dim], ends=[self._head_dim], axes=[-1]) + sin = op.CastLike(sin, values) + cos = op.CastLike(cos, values) + return op.Concat( + op.Sub(op.Mul(first, cos), op.Mul(second, sin)), + op.Add(op.Mul(second, cos), op.Mul(first, sin)), + axis=-1, + ) + + def forward(self, op: OpBuilder, values: ir.Value, patch_mask: ir.Value) -> ir.Value: + shape = op.Shape(values) + batch_like = op.Slice(shape, starts=[0], ends=[1]) + seq_len = op.Slice(shape, starts=[1], ends=[2]) + qkv_shape = op.Concat( + batch_like, + seq_len, + op.Constant(value_ints=[self._num_heads, self._head_dim]), + axis=0, + ) + query = op.Reshape(self.query_proj(op, values), qkv_shape) + key = op.Reshape(self.key_proj(op, values), qkv_shape) + value = op.Reshape(self.value_proj(op, values), qkv_shape) + + if self._use_rope: + query = self._rope(op, query) + key = self._rope(op, key) + query = self.per_dim_scale(op, self.query_ln(op, query)) + key = self.key_ln(op, key) + + query = op.Transpose(query, perm=[0, 2, 1, 3]) + key = op.Transpose(key, perm=[0, 2, 1, 3]) + value = op.Transpose(value, perm=[0, 2, 1, 3]) + scores = op.Mul( + op.MatMul(query, op.Transpose(key, perm=[0, 1, 3, 2])), + math.sqrt(self._head_dim), + ) + + key_valid = op.Unsqueeze(op.Not(patch_mask), axes=[1, 2]) + if self._causal: + length = op.Squeeze(seq_len) + indices = op.Range(op.Constant(value_int=0), length, op.Constant(value_int=1)) + causal = op.GreaterOrEqual( + op.Unsqueeze(indices, axes=[1]), op.Unsqueeze(indices, axes=[0]) + ) + allowed = op.And(key_valid, op.Unsqueeze(causal, axes=[0, 1])) + else: + allowed = key_valid + + masked_scores = op.Where( + allowed, + scores, + op.CastLike(op.Constant(value_float=self._mask_value), scores), + ) + probabilities = op.Softmax(masked_scores, axis=-1) + row_valid = op.Cast( + op.ReduceMax(op.Cast(allowed, to=ir.DataType.INT64), axes=[-1], keepdims=True), + to=ir.DataType.BOOL, + ) + attended = op.Where( + row_valid, + op.MatMul(probabilities, value), + op.CastLike(op.Constant(value_float=0.0), value), + ) + attended = op.Transpose(attended, perm=[0, 2, 1, 3]) + output_shape = op.Concat( + batch_like, + seq_len, + op.Constant(value_ints=[self._model_dims]), + axis=0, + ) + return self.out_proj(op, op.Reshape(attended, output_shape)) + + +class _MixingTransformer(nn.Module): + def __init__(self, config: TimesFM3Config): + super().__init__() + self.pre_seq_attn_ln = _TimesFMRMSNorm(config.model_dims, config.rms_norm_eps) + self.post_seq_attn_ln = _TimesFMRMSNorm(config.model_dims, config.rms_norm_eps) + self.seq_attn = _TimesFMAttention(config, causal=True, use_rope=config.use_rope_seq) + if config.use_variate_attention: + self.pre_var_attn_ln = _TimesFMRMSNorm(config.model_dims, config.rms_norm_eps) + self.post_var_attn_ln = _TimesFMRMSNorm(config.model_dims, config.rms_norm_eps) + self.var_attn = _TimesFMAttention( + config, causal=False, use_rope=config.use_rope_var + ) + else: + self.pre_var_attn_ln = None + self.post_var_attn_ln = None + self.var_attn = None + self.pre_ff_ln = _TimesFMRMSNorm(config.model_dims, config.rms_norm_eps) + self.post_ff_ln = _TimesFMRMSNorm(config.model_dims, config.rms_norm_eps) + self.ff0 = Linear(config.model_dims, config.transformer_hidden_dims, bias=False) + self.ff1 = Linear(config.transformer_hidden_dims, config.model_dims, bias=False) + self._model_dims = config.model_dims + + def forward(self, op: OpBuilder, values: ir.Value, patch_mask: ir.Value) -> ir.Value: + shape = op.Shape(values) + batch = op.Slice(shape, starts=[0], ends=[1]) + variates = op.Slice(shape, starts=[1], ends=[2]) + patches = op.Slice(shape, starts=[2], ends=[3]) + + seq_shape = op.Concat( + op.Mul(batch, variates), + patches, + op.Constant(value_ints=[self._model_dims]), + axis=0, + ) + seq_input = op.Reshape(self.pre_seq_attn_ln(op, values), seq_shape) + seq_mask = op.Reshape(patch_mask, op.Concat(op.Mul(batch, variates), patches, axis=0)) + seq_output = self.seq_attn(op, seq_input, seq_mask) + seq_output = op.Reshape( + seq_output, + op.Concat( + batch, + variates, + patches, + op.Constant(value_ints=[self._model_dims]), + axis=0, + ), + ) + hidden = op.Add(self.post_seq_attn_ln(op, seq_output), values) + + if self.var_attn is not None: + var_input = op.Transpose(self.pre_var_attn_ln(op, hidden), perm=[0, 2, 1, 3]) + var_shape = op.Concat( + op.Mul(batch, patches), + variates, + op.Constant(value_ints=[self._model_dims]), + axis=0, + ) + var_input = op.Reshape(var_input, var_shape) + var_mask = op.Reshape( + op.Transpose(patch_mask, perm=[0, 2, 1]), + op.Concat(op.Mul(batch, patches), variates, axis=0), + ) + var_output = self.var_attn(op, var_input, var_mask) + var_output = op.Transpose( + op.Reshape( + var_output, + op.Concat( + batch, + patches, + variates, + op.Constant(value_ints=[self._model_dims]), + axis=0, + ), + ), + perm=[0, 2, 1, 3], + ) + hidden = op.Add(self.post_var_attn_ln(op, var_output), hidden) + + ff = self.ff0(op, self.pre_ff_ln(op, hidden)) + ff = self.ff1(op, op.Relu(ff)) + return op.Add(self.post_ff_ln(op, ff), hidden) + + +class _StackedMixingTransformer(nn.Module): + def __init__(self, config: TimesFM3Config): + super().__init__() + self.layers = nn.ModuleList( + [_MixingTransformer(config) for _ in range(config.num_layers)] + ) + + def forward(self, op: OpBuilder, values: ir.Value, patch_mask: ir.Value) -> ir.Value: + for layer in self.layers: + values = layer(op, values, patch_mask) + return values + + +def _cpm_refinement_body(rolls: int, patch_len: int, value_clip: float) -> ir.Graph: + float_type = ir.TensorType(ir.DataType.FLOAT) + int_type = ir.TensorType(ir.DataType.INT64) + bool_type = ir.TensorType(ir.DataType.BOOL) + count = ir.Value(name="count", type=float_type) + mean = ir.Value(name="mean", type=float_type) + std = ir.Value(name="std", type=float_type) + anchor = ir.Value(name="anchor", type=float_type) + offset = ir.Value(name="offset", type=int_type) + actual_count = ir.Value(name="actual_count", type=float_type) + actual_mean = ir.Value(name="actual_mean", type=float_type) + actual_std = ir.Value(name="actual_std", type=float_type) + current_median = ir.Value(name="current_median", type=float_type) + is_cpm = ir.Value(name="is_cpm", type=bool_type) + graph, builder = create_body_graph( + state_inputs=[count, mean, std, anchor, offset], + scan_inputs=[actual_count, actual_mean, actual_std, current_median, is_cpm], + name="timesfm3_cpm_refinement", + ) + op = builder.op + selector = op.Cast( + op.Equal( + op.Unsqueeze(offset, axes=[1]), + op.Constant(value_ints=list(range(rolls))), + ), + to=ir.DataType.FLOAT, + ) + predicted = op.ReduceSum( + op.Mul(anchor, op.Unsqueeze(selector, axes=[1, 3])), + axes=[2], + keepdims=False, + ) + prediction_mask = op.ConstantOfShape( + op.Shape(predicted), + value=ir.tensor([False], dtype=ir.DataType.BOOL), + ) + candidate_count, candidate_mean, candidate_std = _update_running_stats( + op, count, mean, std, predicted, prediction_mask + ) + cpm_bv = op.Unsqueeze(is_cpm, axes=[1]) + output_count = op.Where(cpm_bv, candidate_count, actual_count) + output_mean = op.Where(cpm_bv, candidate_mean, actual_mean) + output_std = op.Where(cpm_bv, candidate_std, actual_std) + new_offset = op.Where( + is_cpm, + op.Mod(op.Add(offset, op.Constant(value_int=1)), op.Constant(value_int=rolls)), + op.Mul(offset, op.Constant(value_int=0)), + ) + new_anchor = op.Clip( + op.Add( + op.Mul( + current_median, + op.Unsqueeze(output_std, axes=[-1, -2]), + ), + op.Unsqueeze(output_mean, axes=[-1, -2]), + ), + -value_clip, + value_clip, + ) + updated_anchor = op.Where( + op.Unsqueeze(op.Equal(new_offset, op.Constant(value_int=0)), axes=[1, 2, 3]), + new_anchor, + anchor, + ) + nonnegative_count = op.GreaterOrEqual(actual_count, op.Constant(value_float=0.0)) + scan_mean = op.Where(nonnegative_count, output_mean, op.Constant(value_float=0.0)) + scan_std = op.Where(nonnegative_count, output_std, op.Constant(value_float=0.0)) + graph.outputs.extend( + [ + output_count, + output_mean, + output_std, + updated_anchor, + new_offset, + scan_mean, + scan_std, + ] + ) + rename_subgraph_values(graph, "timesfm3_cpm_") + return graph + + +class TimesFM3Model(nn.Module): + """TimesFM 3 padded-batch multivariate forecasting network.""" + + default_task: str = "time-series-forecasting" + category: str = "Time Series" + config_class = TimesFM3Config + + def __init__(self, config: TimesFM3Config): + super().__init__() + config.validate() + self.config = config + feature_dims = 2 * (config.input_patch_len + config.output_patch_len) + self.pre_transformer_resblock = _ResidualBlock( + feature_dims, config.model_dims, config.model_dims + ) + self.transformer_stack = _StackedMixingTransformer(config) + self.output_head = Linear( + config.model_dims, + config.output_patch_len * len(config.quantiles), + bias=True, + ) + self._rolls = config.output_patch_len // config.input_patch_len + self._value_clip = ( + min(config.value_clip, 65504.0) + if config.dtype == ir.DataType.FLOAT16 + else config.value_clip + ) + + def prepare_raw_series( + self, + op: OpBuilder, + context_values: ir.Value, + context_observed: ir.Value, + future_values: ir.Value, + future_observed: ir.Value, + context_lengths: ir.Value, + horizon_lengths: ir.Value, + variate_roles: ir.Value, + ) -> tuple[ir.Value, ...]: + """Convert right-aligned raw series to the upstream patched contract. + + ``context_values`` is right-aligned in ``[B, V, C]`` and + ``future_values`` is left-aligned in ``[B, V, H]``. The corresponding + observed tensors use True for supplied, finite observations. Roles are + 0=target, 1=past-only covariate, and 2=past-future covariate. + """ + config = self.config + patch_len = config.input_patch_len + rolls = self._rolls + extract_len = min(2 * patch_len, config.output_patch_len) + overlap = extract_len - patch_len + + batch = op.Shape(context_values, start=0, end=1) + variates = op.Shape(context_values, start=1, end=2) + context_width = op.Squeeze(op.Shape(context_values, start=2, end=3)) + future_width = op.Squeeze(op.Shape(future_values, start=2, end=3)) + max_context = op.ReduceMax(context_lengths, keepdims=False) + padded_context = op.Mul( + op.Div( + op.Add(max_context, op.Constant(value_int=patch_len - 1)), + op.Constant(value_int=patch_len), + ), + op.Constant(value_int=patch_len), + ) + padded_context = op.Max(padded_context, op.Constant(value_int=patch_len)) + context_patches = op.Div(padded_context, op.Constant(value_int=patch_len)) + + forecast_numer = op.Max( + op.Sub(horizon_lengths, op.Constant(value_int=overlap)), + op.Constant(value_int=0), + ) + forecast_patches = op.Div( + op.Add(forecast_numer, op.Constant(value_int=patch_len - 1)), + op.Constant(value_int=patch_len), + ) + forecast_patches = op.Max(forecast_patches, op.Constant(value_int=1)) + max_forecast_patches = op.ReduceMax(forecast_patches, keepdims=False) + horizon_patches = op.Add(max_forecast_patches, op.Constant(value_int=rolls - 1)) + padded_horizon = op.Mul(horizon_patches, op.Constant(value_int=patch_len)) + + is_target = op.Equal(variate_roles, op.Constant(value_int=_TARGET_ROLE)) + is_past_only = op.Equal(variate_roles, op.Constant(value_int=_PAST_ONLY_ROLE)) + is_past_future = op.Equal(variate_roles, op.Constant(value_int=_PAST_FUTURE_ROLE)) + valid_role = op.Or(op.Or(is_target, is_past_only), is_past_future) + + # Right-align every row in the common dynamic context patch grid. + context_position = op.Range( + op.Constant(value_int=0), + padded_context, + op.Constant(value_int=1), + ) + context_source = op.Add(context_position, op.Sub(context_width, padded_context)) + context_source = op.Min( + op.Max(context_source, op.Constant(value_int=0)), + op.Sub(context_width, op.Constant(value_int=1)), + ) + context_shape = op.Concat( + batch, variates, op.Unsqueeze(padded_context, axes=[0]), axis=0 + ) + context_indices = op.Expand(op.Unsqueeze(context_source, axes=[0, 1]), context_shape) + aligned_context = op.GatherElements(context_values, context_indices, axis=2) + aligned_context_observed = op.GatherElements(context_observed, context_indices, axis=2) + context_active = op.GreaterOrEqual( + op.Unsqueeze(context_position, axes=[0, 1]), + op.Sub( + padded_context, + op.Unsqueeze(context_lengths, axes=[1, 2]), + ), + ) + context_active = op.And(context_active, op.Unsqueeze(valid_role, axes=[2])) + + # Future observations are left-aligned and only past-future variates + # participate. Extra common horizon patches remain masked. + horizon_position = op.Range( + op.Constant(value_int=0), + padded_horizon, + op.Constant(value_int=1), + ) + future_source = op.Min( + horizon_position, + op.Sub(future_width, op.Constant(value_int=1)), + ) + horizon_shape = op.Concat( + batch, variates, op.Unsqueeze(padded_horizon, axes=[0]), axis=0 + ) + future_indices = op.Expand(op.Unsqueeze(future_source, axes=[0, 1]), horizon_shape) + aligned_future = op.GatherElements(future_values, future_indices, axis=2) + aligned_future_observed = op.GatherElements(future_observed, future_indices, axis=2) + future_active = op.And( + op.Less( + op.Unsqueeze(horizon_position, axes=[0, 1]), + op.Unsqueeze(horizon_lengths, axes=[1, 2]), + ), + op.Unsqueeze(is_past_future, axes=[2]), + ) + + raw_values = op.Concat(aligned_context, aligned_future, axis=2) + finite = op.Not(op.Or(op.IsNaN(raw_values), op.IsInf(raw_values))) + raw_values = op.Where( + finite, + raw_values, + op.CastLike(op.Constant(value_float=0.0), raw_values), + ) + supplied = op.Concat( + op.And(context_active, aligned_context_observed), + op.And(future_active, aligned_future_observed), + axis=2, + ) + supplied = op.And(supplied, finite) + interpolated = _interpolate_missing(op, raw_values, supplied) + + context = op.Slice( + interpolated, + starts=op.Constant(value_ints=[0]), + ends=op.Unsqueeze(padded_context, axes=[0]), + axes=op.Constant(value_ints=[2]), + ) + future = op.Slice( + interpolated, + starts=op.Unsqueeze(padded_context, axes=[0]), + ends=op.Unsqueeze(op.Add(padded_context, padded_horizon), axes=[0]), + axes=op.Constant(value_ints=[2]), + ) + + # Fit y = m*t + c using each row's unpadded context length. The + # right-aligned time grid is -(length-1)..0 for every row. + context_f32 = op.Cast(context, to=ir.DataType.FLOAT) + context_valid = context_active + valid_f32 = op.Cast(context_valid, to=ir.DataType.FLOAT) + length_f32 = op.Cast(op.Unsqueeze(context_lengths, axes=[1, 2]), to=ir.DataType.FLOAT) + time = op.Sub( + op.Cast( + op.Unsqueeze(context_position, axes=[0, 1]), + to=ir.DataType.FLOAT, + ), + op.Cast( + op.Sub(padded_context, op.Constant(value_int=1)), + to=ir.DataType.FLOAT, + ), + ) + normalized_time = op.Div(time, op.Max(length_f32, op.Constant(value_float=1.0))) + zeros = op.ConstantOfShape( + op.Shape(context_f32), value=ir.tensor([0.0], dtype=ir.DataType.FLOAT) + ) + valid_values = op.Where(context_valid, context_f32, zeros) + valid_time = op.Where(context_valid, normalized_time, zeros) + count = op.ReduceSum(valid_f32, axes=[-1], keepdims=True) + safe_count = op.Max(count, op.Constant(value_float=1.0)) + sum_time = op.ReduceSum(valid_time, axes=[-1], keepdims=True) + sum_time2 = op.ReduceSum( + op.Where(context_valid, op.Mul(normalized_time, normalized_time), zeros), + axes=[-1], + keepdims=True, + ) + sum_values = op.ReduceSum(valid_values, axes=[-1], keepdims=True) + sum_time_values = op.ReduceSum( + op.Where( + context_valid, + op.Mul(normalized_time, context_f32), + zeros, + ), + axes=[-1], + keepdims=True, + ) + determinant = op.Sub(op.Mul(count, sum_time2), op.Mul(sum_time, sum_time)) + determinant_zero = op.Equal(determinant, op.Constant(value_float=0.0)) + safe_determinant = op.Where( + determinant_zero, op.Constant(value_float=1.0), determinant + ) + trend_slope = op.Where( + determinant_zero, + op.Constant(value_float=0.0), + op.Div( + op.Sub( + op.Mul(count, sum_time_values), + op.Mul(sum_time, sum_values), + ), + safe_determinant, + ), + ) + trend_intercept = op.Where( + determinant_zero, + op.Where( + op.Greater(count, op.Constant(value_float=0.0)), + op.Div(sum_values, safe_count), + op.Constant(value_float=0.0), + ), + op.Div( + op.Sub(sum_values, op.Mul(trend_slope, sum_time)), + safe_count, + ), + ) + detrended_context = op.Sub( + context_f32, + op.Add(op.Mul(trend_slope, normalized_time), trend_intercept), + ) + mean = op.Div(sum_values, safe_count) + sum_values2 = op.ReduceSum( + op.Where(context_valid, op.Mul(context_f32, context_f32), zeros), + axes=[-1], + keepdims=True, + ) + original_variance = op.Max( + op.Sub(op.Div(sum_values2, safe_count), op.Mul(mean, mean)), + op.Constant(value_float=0.0), + ) + detrended_values = op.Where(context_valid, detrended_context, zeros) + detrended_sum = op.ReduceSum(detrended_values, axes=[-1], keepdims=True) + detrended_mean = op.Div(detrended_sum, safe_count) + detrended_sum2 = op.ReduceSum( + op.Mul(detrended_values, detrended_values), + axes=[-1], + keepdims=True, + ) + detrended_variance = op.Max( + op.Sub( + op.Div(detrended_sum2, safe_count), + op.Mul(detrended_mean, detrended_mean), + ), + op.Constant(value_float=0.0), + ) + apply_detrend = op.Less( + op.Sqrt(detrended_variance), + op.Mul( + op.Constant(value_float=config.linear_detrending_threshold), + op.Sqrt(original_variance), + ), + ) + if not config.use_linear_detrending: + apply_detrend = op.And( + apply_detrend, + op.ConstantOfShape( + op.Shape(apply_detrend), + value=ir.tensor([False], dtype=ir.DataType.BOOL), + ), + ) + context = op.CastLike( + op.Where(apply_detrend, detrended_context, context_f32), + context_values, + ) + + horizon_step = op.Cast( + op.Add(horizon_position, op.Constant(value_int=1)), + to=ir.DataType.FLOAT, + ) + future_trend = op.Add( + op.Mul( + trend_slope, + op.Div( + op.Unsqueeze(horizon_step, axes=[0, 1]), + op.Max(length_f32, op.Constant(value_float=1.0)), + ), + ), + trend_intercept, + ) + future_f32 = op.Cast(future, to=ir.DataType.FLOAT) + future = op.CastLike( + op.Where( + apply_detrend, + op.Sub(future_f32, future_trend), + future_f32, + ), + future_values, + ) + + context_masks = op.Not(context_active) + future_masks = op.Not(future_active) + all_values = op.Concat( + op.Where( + context_masks, + op.CastLike(op.Constant(value_float=0.0), context), + context, + ), + op.Where( + future_masks, + op.CastLike(op.Constant(value_float=0.0), future), + future, + ), + axis=2, + ) + all_masks = op.Concat(context_masks, future_masks, axis=2) + values = op.Reshape( + all_values, + op.Concat( + batch, + variates, + op.Constant(value_ints=[-1, patch_len]), + axis=0, + ), + ) + masks = op.Reshape( + all_masks, + op.Concat( + batch, + variates, + op.Constant(value_ints=[-1, patch_len]), + axis=0, + ), + ) + total_patches = op.Squeeze(op.Shape(values, start=2, end=3)) + patch_index = op.Range( + op.Constant(value_int=0), + total_patches, + op.Constant(value_int=1), + ) + patch_shape = op.Concat(batch, variates, op.Unsqueeze(total_patches, axes=[0]), axis=0) + patch_is_target = op.Expand( + op.Unsqueeze(op.Or(is_target, is_past_only), axes=[2]), + patch_shape, + ) + patch_cpm_mask = op.Expand( + op.Unsqueeze(op.GreaterOrEqual(patch_index, context_patches), axes=[0]), + op.Concat(batch, op.Unsqueeze(total_patches, axes=[0]), axis=0), + ) + + # Positivity policy is based on the original finite observations, + # before interpolation or detrending. + original_finite = op.Not(op.Or(op.IsNaN(aligned_context), op.IsInf(aligned_context))) + original_valid = op.And( + context_active, + op.And(aligned_context_observed, original_finite), + ) + original_count = op.ReduceSum( + op.Cast(original_valid, to=ir.DataType.INT64), + axes=[-1], + keepdims=False, + ) + original_nonnegative = op.ReduceMin( + op.Cast( + op.Or( + op.Not(original_valid), + op.GreaterOrEqual( + aligned_context, + op.CastLike(op.Constant(value_float=0.0), aligned_context), + ), + ), + to=ir.DataType.INT64, + ), + axes=[-1], + keepdims=False, + ) + nonnegative = op.And( + op.Greater(original_count, op.Constant(value_int=0)), + op.Cast(original_nonnegative, to=ir.DataType.BOOL), + ) + nonnegative = op.And(nonnegative, is_target) + + return ( + values, + masks, + patch_is_target, + patch_cpm_mask, + op.Squeeze(trend_slope, axes=[-1]), + op.Squeeze(trend_intercept, axes=[-1]), + op.Squeeze(apply_detrend, axes=[-1]), + is_target, + nonnegative, + context_lengths, + horizon_lengths, + context_patches, + forecast_patches, + ) + + def _refine_cpm_stats( + self, + op: OpBuilder, + raw_logits: ir.Value, + running_count: ir.Value, + running_mean: ir.Value, + running_std: ir.Value, + patch_cpm_mask: ir.Value, + ) -> tuple[ir.Value, ir.Value]: + config = self.config + structured = op.Reshape( + raw_logits, + [0, 0, 0, self._rolls, config.input_patch_len, len(config.quantiles)], + ) + median = op.Gather(structured, len(config.quantiles) // 2, axis=-1) + median_t = op.Transpose(median, perm=[2, 0, 1, 3, 4]) + count_t = op.Transpose(running_count, perm=[2, 0, 1]) + mean_t = op.Transpose(running_mean, perm=[2, 0, 1]) + std_t = op.Transpose(running_std, perm=[2, 0, 1]) + cpm_t = op.Transpose(patch_cpm_mask, perm=[1, 0]) + + bv_shape = op.Shape(running_count, start=0, end=2) + zeros = op.ConstantOfShape(bv_shape, value=ir.tensor([0.0], dtype=ir.DataType.FLOAT)) + anchor_shape = op.Concat( + bv_shape, + op.Constant(value_ints=[self._rolls, config.input_patch_len]), + axis=0, + ) + anchor = op.ConstantOfShape( + anchor_shape, value=ir.tensor([0.0], dtype=ir.DataType.FLOAT) + ) + batch_shape = op.Shape(running_count, start=0, end=1) + offset = op.ConstantOfShape(batch_shape, value=ir.tensor([0], dtype=ir.DataType.INT64)) + _, _, _, _, _, refined_mean, refined_std = op.Scan( + zeros, + zeros, + zeros, + anchor, + offset, + count_t, + mean_t, + std_t, + median_t, + cpm_t, + body=_cpm_refinement_body(self._rolls, config.input_patch_len, self._value_clip), + num_scan_inputs=5, + _outputs=7, + ) + return ( + op.Transpose(refined_mean, perm=[1, 2, 0]), + op.Transpose(refined_std, perm=[1, 2, 0]), + ) + + def preprocess( + self, + op: OpBuilder, + values: ir.Value, + masks: ir.Value, + patch_is_target: ir.Value, + patch_cpm_mask: ir.Value, + ) -> tuple[ir.Value, ir.Value, ir.Value, ir.Value, ir.Value]: + """Prepare patched inputs and RevIN statistics for the transformer.""" + config = self.config + values = op.Where( + op.IsNaN(values), + op.CastLike(op.Constant(value_float=0.0), values), + values, + ) + values = op.Clip(values, -self._value_clip, self._value_clip) + running_count, running_mean, running_std = _get_running_stats(op, values, masks) + + cpm_target = op.And( + op.Unsqueeze(patch_cpm_mask, axes=[1, 3]), + op.Unsqueeze(patch_is_target, axes=[3]), + ) + effective_masks = op.Or(masks, cpm_target) + model_mean = op.CastLike(running_mean, values) + divisor = op.CastLike(_safe_divisor(op, running_std), values) + current = op.Div( + op.Sub(values, op.Unsqueeze(model_mean, axes=[-1])), + op.Unsqueeze(divisor, axes=[-1]), + ) + current = op.Where( + effective_masks, + op.CastLike(op.Constant(value_float=0.0), current), + current, + ) + + future, wrap_mask = _roll_patches(op, values, self._rolls, config.input_patch_len) + future = op.Div( + op.Sub(future, op.Unsqueeze(model_mean, axes=[-1])), + op.Unsqueeze(divisor, axes=[-1]), + ) + rolled_masks, _ = _roll_patches( + op, + op.Cast(effective_masks, to=ir.DataType.FLOAT), + self._rolls, + config.input_patch_len, + ) + future_masks = op.Or( + op.Or( + op.Cast(rolled_masks, to=ir.DataType.BOOL), op.Unsqueeze(patch_is_target, [3]) + ), + wrap_mask, + ) + future = op.Where( + future_masks, + op.CastLike(op.Constant(value_float=0.0), future), + future, + ) + + values_cat = op.Concat(current, future, axis=-1) + masks_cat = op.Concat(effective_masks, future_masks, axis=-1) + residual_input = op.Concat(values_cat, op.CastLike(masks_cat, values_cat), axis=-1) + patch_mask = op.Cast( + op.ReduceMin(op.Cast(masks_cat, to=ir.DataType.INT64), axes=[-1], keepdims=False), + to=ir.DataType.BOOL, + ) + + prefix_count = op.CumSum( + op.Cast(patch_mask, to=ir.DataType.INT64), op.Constant(value_int=2) + ) + num_patches = op.Squeeze(op.Shape(patch_mask, start=2, end=3)) + patch_ordinals = op.Add( + op.Range(op.Constant(value_int=0), num_patches, op.Constant(value_int=1)), + op.Constant(value_int=1), + ) + effective_patch_mask = op.Equal( + prefix_count, op.Unsqueeze(patch_ordinals, axes=[0, 1]) + ) + return ( + residual_input, + effective_patch_mask, + running_count, + running_mean, + running_std, + ) + + def forecast( + self, + op: OpBuilder, + model_inputs: ir.Value, + patch_mask: ir.Value, + ) -> ir.Value: + """Run every learned layer in the capture-friendly model component.""" + hidden = self.pre_transformer_resblock(op, model_inputs) + hidden = self.transformer_stack(op, hidden, patch_mask) + return self.output_head(op, hidden) + + def postprocess( + self, + op: OpBuilder, + raw_logits: ir.Value, + running_count: ir.Value, + running_mean: ir.Value, + running_std: ir.Value, + patch_cpm_mask: ir.Value, + ) -> tuple[ir.Value, ir.Value, ir.Value]: + """Apply CPM refinement, reverse RevIN, clipping, and output shaping.""" + config = self.config + + output_mean, output_std = running_mean, running_std + if config.use_iterative_cpm_revin: + refined_mean, refined_std = self._refine_cpm_stats( + op, + op.Cast(raw_logits, to=ir.DataType.FLOAT), + running_count, + running_mean, + running_std, + patch_cpm_mask, + ) + cpm = op.Unsqueeze(patch_cpm_mask, axes=[1]) + output_mean = op.Where(cpm, refined_mean, running_mean) + output_std = op.Where(cpm, refined_std, running_std) + + logits = op.Add( + op.Mul( + raw_logits, + op.Unsqueeze(op.CastLike(output_std, raw_logits), axes=[-1]), + ), + op.Unsqueeze(op.CastLike(output_mean, raw_logits), axes=[-1]), + ) + logits = op.Clip(logits, -self._value_clip, self._value_clip) + logits = op.Reshape( + logits, + [0, 0, 0, config.output_patch_len, len(config.quantiles)], + ) + return logits, running_mean, running_std + + def stitch_forecast( + self, + op: OpBuilder, + logits: ir.Value, + trend_slope: ir.Value, + trend_intercept: ir.Value, + apply_detrend: ir.Value, + target_mask: ir.Value, + nonnegative_mask: ir.Value, + make_positive: ir.Value, + context_lengths: ir.Value, + horizon_lengths: ir.Value, + context_patch_count: ir.Value, + forecast_patch_counts: ir.Value, + ) -> tuple[ir.Value, ir.Value, ir.Value]: + """Stitch overlapping patches and apply the public forecast policies.""" + config = self.config + patch_len = config.input_patch_len + extract_len = min(2 * patch_len, config.output_patch_len) + overlap = extract_len - patch_len + num_quantiles = len(config.quantiles) + + max_forecast_patches = op.ReduceMax(forecast_patch_counts, keepdims=False) + forecast_indices = op.Add( + op.Range( + op.Constant(value_int=0), + max_forecast_patches, + op.Constant(value_int=1), + ), + op.Sub(context_patch_count, op.Constant(value_int=1)), + ) + patch_predictions = op.Gather(logits, forecast_indices, axis=2) + patch_predictions = op.Slice( + patch_predictions, + starts=[0], + ends=[extract_len], + axes=[3], + ) + + # Express stitch_patches as indexed gathers so the single-patch case + # needs no control flow. Output positions comprise P points per patch + # plus the final overlap tail. + stitched_length = op.Add( + op.Mul(max_forecast_patches, op.Constant(value_int=patch_len)), + op.Constant(value_int=overlap), + ) + position = op.Range( + op.Constant(value_int=0), + stitched_length, + op.Constant(value_int=1), + ) + base_patch = op.Div(position, op.Constant(value_int=patch_len)) + offset = op.Mod(position, op.Constant(value_int=patch_len)) + patch_end = op.Mul( + op.Unsqueeze(forecast_patch_counts, axes=[1]), + op.Constant(value_int=patch_len), + ) + position_b = op.Unsqueeze(position, axes=[0]) + base_patch_b = op.Unsqueeze(base_patch, axes=[0]) + offset_b = op.Unsqueeze(offset, axes=[0]) + is_tail = op.GreaterOrEqual(position_b, patch_end) + current_patch = op.Min( + base_patch_b, + op.Unsqueeze( + op.Sub(forecast_patch_counts, op.Constant(value_int=1)), + axes=[1], + ), + ) + current_offset = op.Where( + is_tail, + op.Add( + op.Constant(value_int=patch_len), + op.Sub(position_b, patch_end), + ), + offset_b, + ) + current_linear_index = op.Add( + op.Mul(current_patch, op.Constant(value_int=extract_len)), + current_offset, + ) + previous_patch = op.Max( + op.Sub(base_patch_b, op.Constant(value_int=1)), + op.Constant(value_int=0), + ) + previous_linear_index = op.Add( + op.Mul(previous_patch, op.Constant(value_int=extract_len)), + op.Add(op.Constant(value_int=patch_len), offset_b), + ) + + batch = op.Shape(logits, start=0, end=1) + variates = op.Shape(logits, start=1, end=2) + gather_shape = op.Concat( + batch, + variates, + op.Unsqueeze(stitched_length, axes=[0]), + op.Constant(value_ints=[num_quantiles]), + axis=0, + ) + flattened_predictions = op.Reshape( + patch_predictions, + op.Concat( + batch, + variates, + op.Constant(value_ints=[-1, num_quantiles]), + axis=0, + ), + ) + current = op.GatherElements( + flattened_predictions, + op.Expand( + op.Unsqueeze(current_linear_index, axes=[1, 3]), + gather_shape, + ), + axis=2, + ) + previous = op.GatherElements( + flattened_predictions, + op.Expand( + op.Unsqueeze(previous_linear_index, axes=[1, 3]), + gather_shape, + ), + axis=2, + ) + blend_position = op.Cast(offset, to=ir.DataType.FLOAT) + blend_weight = op.Sub( + op.Constant(value_float=1.0), + op.Div( + blend_position, + op.Constant(value_float=float(max(overlap - 1, 1))), + ), + ) + blend_weight = op.CastLike(op.Unsqueeze(blend_weight, axes=[0, 1, 3]), current) + blended = op.Add( + op.Mul(blend_weight, previous), + op.Mul( + op.Sub( + op.CastLike(op.Constant(value_float=1.0), blend_weight), + blend_weight, + ), + current, + ), + ) + should_blend = op.And( + op.And( + op.GreaterOrEqual(position_b, op.Constant(value_int=patch_len)), + op.Not(is_tail), + ), + op.Less(offset_b, op.Constant(value_int=overlap)), + ) + stitched = op.Where( + op.Unsqueeze(should_blend, axes=[1, 3]), + blended, + current, + ) + + max_horizon = op.ReduceMax(horizon_lengths, keepdims=False) + stitched = op.Slice( + stitched, + starts=op.Constant(value_ints=[0]), + ends=op.Unsqueeze(max_horizon, axes=[0]), + axes=op.Constant(value_ints=[2]), + ) + horizon_position = op.Range( + op.Constant(value_int=0), + max_horizon, + op.Constant(value_int=1), + ) + trend = op.Add( + op.Mul( + op.Unsqueeze(trend_slope, axes=[2]), + op.Div( + op.Cast( + op.Unsqueeze( + op.Add(horizon_position, op.Constant(value_int=1)), + axes=[0, 1], + ), + to=ir.DataType.FLOAT, + ), + op.Cast( + op.Unsqueeze(context_lengths, axes=[1, 2]), + to=ir.DataType.FLOAT, + ), + ), + ), + op.Unsqueeze(trend_intercept, axes=[2]), + ) + trend = op.Where( + op.Unsqueeze(apply_detrend, axes=[2]), + trend, + op.Constant(value_float=0.0), + ) + forecasts = op.Add( + stitched, + op.Unsqueeze(op.CastLike(trend, stitched), axes=[3]), + ) + + # Upstream sorts quantiles before selecting the configured median. + quantile_forecasts, _ = op.TopK( + forecasts, + op.Constant(value_ints=[num_quantiles]), + axis=-1, + largest=0, + sorted=1, + _outputs=2, + ) + apply_positive = op.And(nonnegative_mask, make_positive) + quantile_forecasts = op.Where( + op.Unsqueeze(apply_positive, axes=[2, 3]), + op.Max( + quantile_forecasts, + op.CastLike(op.Constant(value_float=0.0), quantile_forecasts), + ), + quantile_forecasts, + ) + validity = op.And( + op.Unsqueeze(target_mask, axes=[2]), + op.Less( + op.Unsqueeze(horizon_position, axes=[0, 1]), + op.Unsqueeze(horizon_lengths, axes=[1, 2]), + ), + ) + quantile_forecasts = op.Where( + op.Unsqueeze(validity, axes=[3]), + quantile_forecasts, + op.CastLike(op.Constant(value_float=0.0), quantile_forecasts), + ) + point_forecast = op.Gather(quantile_forecasts, num_quantiles // 2, axis=-1) + return point_forecast, quantile_forecasts, validity + + def forward( + self, + op: OpBuilder, + values: ir.Value, + masks: ir.Value, + patch_is_target: ir.Value, + patch_cpm_mask: ir.Value, + ) -> tuple[ir.Value, ir.Value, ir.Value]: + """Compose the three export stages for direct module use.""" + model_inputs, patch_mask, running_count, running_mean, running_std = self.preprocess( + op, + values, + masks, + patch_is_target, + patch_cpm_mask, + ) + raw_logits = self.forecast(op, model_inputs, patch_mask) + return self.postprocess( + op, + raw_logits, + running_count, + running_mean, + running_std, + patch_cpm_mask, + ) diff --git a/src/mobius/models/timesfm3_test.py b/src/mobius/models/timesfm3_test.py new file mode 100644 index 000000000..8fced699c --- /dev/null +++ b/src/mobius/models/timesfm3_test.py @@ -0,0 +1,670 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import json + +import numpy as np +import onnx_ir as ir +import pytest + +from mobius import build_from_module +from mobius._configs import BaseModelConfig +from mobius._registry import registry +from mobius.integrations.transformers._config_resolver import _try_load_config_json +from mobius.models.timesfm3 import TimesFM3Config, TimesFM3Model +from mobius.tasks import TimeSeriesForecastingTask, get_task + + +def _tiny_config() -> TimesFM3Config: + return TimesFM3Config( + input_patch_len=4, + output_patch_len=8, + quantiles=(0.1, 0.5, 0.9), + num_layers=1, + model_dims=16, + transformer_hidden_dims=24, + num_heads=4, + max_variates=4, + ) + + +def _build_tiny(): + config = _tiny_config() + module = TimesFM3Model(config) + package = build_from_module( + module, + config, + task=TimeSeriesForecastingTask(), + ) + return config, module, package + + +@pytest.mark.parametrize( + ("input_patch_len", "output_patch_len"), + [(0, 8), (-1, 8), (4, 0), (4, -1)], +) +def test_config_rejects_nonpositive_patch_lengths( + input_patch_len: int, output_patch_len: int +) -> None: + config = _tiny_config() + config.input_patch_len = input_patch_len + config.output_patch_len = output_patch_len + + with pytest.raises(ValueError, match="must be positive"): + config.validate() + + +def test_forecasting_task_rejects_config_without_patch_length() -> None: + with pytest.raises(TypeError, match="integer input_patch_len"): + TimeSeriesForecastingTask().build(TimesFM3Model(_tiny_config()), BaseModelConfig()) + + +@pytest.fixture(scope="module") +def tiny_package(): + return _build_tiny() + + +def _run_pipeline(package, feeds: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + from mobius._testing.ort_inference import OnnxModelSession + + preprocessor = OnnxModelSession(package["preprocessor"]) + preprocessed = preprocessor.run(feeds) + preprocessor.close() + + model = OnnxModelSession(package["model"]) + raw_logits = model.run( + { + "model_inputs": preprocessed["model_inputs"], + "patch_mask": preprocessed["patch_mask"], + } + ) + model.close() + + postprocessor = OnnxModelSession(package["postprocessor"]) + outputs = postprocessor.run( + { + "raw_logits": raw_logits["raw_logits"], + "revin_count": preprocessed["revin_count"], + "revin_mean": preprocessed["revin_mean"], + "revin_std": preprocessed["revin_std"], + "patch_cpm_mask": feeds["patch_cpm_mask"], + } + ) + postprocessor.close() + return outputs + + +def _run_full_pipeline( + package, + feeds: dict[str, np.ndarray], + *, + make_positive: bool = False, +) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: + from mobius._testing.ort_inference import OnnxModelSession + + raw_session = OnnxModelSession(package["raw_preprocessor"]) + raw = raw_session.run(feeds) + raw_session.close() + patched = _run_pipeline( + package, + {name: raw[name] for name in ("values", "masks", "patch_is_target", "patch_cpm_mask")}, + ) + stitch_session = OnnxModelSession(package["stitcher"]) + outputs = stitch_session.run( + { + "logits": patched["logits"], + "make_positive": np.array(make_positive), + **{ + name: raw[name] + for name in ( + "trend_slope", + "trend_intercept", + "apply_detrend", + "target_mask", + "nonnegative_mask", + "context_lengths", + "horizon_lengths", + "context_patch_count", + "forecast_patch_counts", + ) + }, + } + ) + stitch_session.close() + return raw, outputs + + +def test_config_parses_official_nested_schema() -> None: + config = TimesFM3Config.from_transformers( + { + "input_patch_len": 32, + "output_patch_len": 64, + "quantiles": [0.1, 0.5, 0.9], + "use_iterative_cpm_revin": True, + "use_variate_attention": True, + "value_clip": 1e20, + "residual_block_config": {"output_dims": 1280}, + "transformer_config": { + "num_layers": 20, + "transformer": { + "model_dims": 1280, + "hidden_dims": 1280, + "num_heads": 16, + "max_variates": 32, + "use_rope_seq": True, + "use_rope_var": False, + }, + }, + } + ) + + assert config.input_patch_len == 32 + assert config.output_patch_len == 64 + assert config.num_layers == 20 + assert config.head_dim == 80 + assert config.quantiles == (0.1, 0.5, 0.9) + assert config.rms_norm_eps == np.finfo(np.float32).eps + assert config.use_linear_detrending + assert config.linear_detrending_threshold == pytest.approx(0.5) + + +def test_raw_config_detection(tmp_path) -> None: + (tmp_path / "config.json").write_text( + json.dumps( + { + "input_patch_len": 32, + "output_patch_len": 64, + "quantiles": [0.1, 0.5, 0.9], + "use_iterative_cpm_revin": True, + "use_variate_attention": True, + "residual_block_config": {"output_dims": 1280}, + "transformer_config": { + "num_layers": 20, + "transformer": {"model_dims": 1280, "max_variates": 32}, + }, + } + ) + ) + + config = _try_load_config_json(str(tmp_path)) + + assert config is not None + assert config.model_type == "timesfm3" + + +def test_graph_contract_and_checkpoint_weight_names(tiny_package) -> None: + config, module, package = tiny_package + + assert set(package) == { + "raw_preprocessor", + "preprocessor", + "model", + "postprocessor", + "stitcher", + } + assert {value.name for value in package["raw_preprocessor"].graph.inputs} == { + "context_values", + "context_observed", + "future_values", + "future_observed", + "context_lengths", + "horizon_lengths", + "variate_roles", + } + assert {value.name for value in package["raw_preprocessor"].graph.outputs} == { + "values", + "masks", + "patch_is_target", + "patch_cpm_mask", + "trend_slope", + "trend_intercept", + "apply_detrend", + "target_mask", + "nonnegative_mask", + "context_lengths", + "horizon_lengths", + "context_patch_count", + "forecast_patch_counts", + } + assert {value.name for value in package["preprocessor"].graph.inputs} == { + "values", + "masks", + "patch_is_target", + "patch_cpm_mask", + } + assert {value.name for value in package["preprocessor"].graph.outputs} == { + "model_inputs", + "patch_mask", + "revin_count", + "revin_mean", + "revin_std", + } + assert {value.name for value in package["model"].graph.inputs} == { + "model_inputs", + "patch_mask", + } + assert {value.name for value in package["model"].graph.outputs} == {"raw_logits"} + assert {value.name for value in package["postprocessor"].graph.inputs} == { + "raw_logits", + "revin_count", + "revin_mean", + "revin_std", + "patch_cpm_mask", + } + assert {value.name for value in package["postprocessor"].graph.outputs} == { + "logits", + "revin_mean", + "revin_std", + } + assert {value.name for value in package["stitcher"].graph.outputs} == { + "point_forecast", + "quantile_forecasts", + "validity", + } + assert package["preprocessor"].graph.inputs[0].shape[-1] == config.input_patch_len + assert not {node.op_type for node in package["model"].graph}.intersection( + {"Scan", "Loop", "If"} + ) + + initializers = { + name for component in package.values() for name in component.graph.initializers + } + assert "pre_transformer_resblock.hidden_layer.weight" in initializers + assert "transformer_stack.layers.0.seq_attn.query_proj.weight" in initializers + assert "transformer_stack.layers.0.var_attn.per_dim_scale.per_dim_scale" in initializers + assert "output_head.weight" in initializers + assert "output_head.bias" in initializers + parameter_names = {name for name, _ in module.named_parameters()} + assert parameter_names.issubset(initializers) + assert all( + sum(name in component.graph.initializers for component in package.values()) == 1 + for name in parameter_names + ) + + import torch + + state_dict = { + name: torch.zeros(tuple(parameter.shape), dtype=torch.float32) + for name, parameter in module.named_parameters() + } + package.apply_weights(state_dict, fold_constants=False) + applied_names = { + name + for component in package.values() + for name, initializer in component.graph.initializers.items() + if initializer.const_value is not None + } + assert parameter_names.issubset(applied_names) + assert len(list(TimesFM3Model(TimesFM3Config()).named_parameters())) == 445 + + +def test_raw_preprocessor_matches_interpolation_detrending_and_padding( + tiny_package, +) -> None: + from mobius._testing.ort_inference import OnnxModelSession + + _, _, package = tiny_package + context = np.zeros((2, 3, 6), dtype=np.float32) + context[0, 0, -5:] = [0.0, np.nan, 2.0, 3.0, 4.0] + context[0, 1, -5:] = [0.0, np.nan, 4.0, 9.0, 16.0] + context[0, 2, -5:] = [10.0, np.nan, 14.0, 16.0, 18.0] + context[1, 0, -3:] = [2.0, 4.0, 6.0] + context[1, 1, -3:] = [1.0, 2.0, 5.0] + context[1, 2, -3:] = [3.0, 6.0, 9.0] + context_observed = np.isfinite(context) + + future = np.zeros((2, 3, 10), dtype=np.float32) + future[0, 2] = np.arange(20.0, 40.0, 2.0) + future[0, 2, 1] = np.nan + future[1, 2, :3] = [12.0, np.nan, 18.0] + future_observed = np.isfinite(future) + + session = OnnxModelSession(package["raw_preprocessor"]) + outputs = session.run( + { + "context_values": context, + "context_observed": context_observed, + "future_values": future, + "future_observed": future_observed, + "context_lengths": np.array([5, 3], dtype=np.int64), + "horizon_lengths": np.array([10, 3], dtype=np.int64), + "variate_roles": np.array([[0, 1, 2], [0, 1, 2]], dtype=np.int64), + } + ) + session.close() + + assert outputs["values"].shape == (2, 3, 5, 4) + np.testing.assert_array_equal( + outputs["patch_cpm_mask"], + np.array([[False, False, True, True, True]] * 2), + ) + np.testing.assert_array_equal( + outputs["forecast_patch_counts"], np.array([2, 1], dtype=np.int64) + ) + assert outputs["context_patch_count"] == 2 + + flat = outputs["values"].reshape(2, 3, -1) + masks = outputs["masks"].reshape(2, 3, -1) + # Recover the interpolated nonlinear row from its emitted trend metadata. + past_time = np.arange(-4, 1, dtype=np.float32) / 5.0 + restored_past = np.where( + outputs["apply_detrend"][0, 1], + flat[0, 1, 3:8] + + outputs["trend_slope"][0, 1] * past_time + + outputs["trend_intercept"][0, 1], + flat[0, 1, 3:8], + ) + np.testing.assert_allclose(restored_past, [0.0, 2.0, 4.0, 9.0, 16.0], atol=1e-5) + np.testing.assert_array_equal( + masks[0, 1, :8], [True, True, True, False, False, False, False, False] + ) + # Perfectly linear target and past-future rows become zero, including + # interpolated future covariates across the context/horizon boundary. + np.testing.assert_allclose(flat[0, 0, 3:8], 0.0, atol=1e-6) + np.testing.assert_allclose(flat[0, 2, 3:18], 0.0, atol=1e-6) + np.testing.assert_allclose(outputs["trend_slope"][0, [0, 2]], [5.0, 10.0], atol=1e-5) + np.testing.assert_allclose(outputs["trend_intercept"][0, [0, 2]], [4.0, 18.0], atol=1e-5) + np.testing.assert_array_equal(outputs["apply_detrend"][0, [0, 2]], [True, True]) + np.testing.assert_array_equal( + outputs["nonnegative_mask"], + np.array([[True, False, False], [True, False, False]]), + ) + + +def test_stitcher_matches_overlap_trend_sort_mask_and_clipping(tiny_package) -> None: + from mobius._testing.ort_inference import OnnxModelSession + + _, _, package = tiny_package + logits = np.zeros((2, 2, 5, 8, 3), dtype=np.float32) + quantile_offsets = np.array([2.0, -1.0, 1.0], dtype=np.float32) + for batch_index in range(2): + for variate_index in range(2): + for forecast_index, patch_index in enumerate((1, 2)): + center = ( + -12.0 + + 20.0 * batch_index + + 5.0 * variate_index + + 8.0 * forecast_index + + np.arange(8, dtype=np.float32) + ) + logits[batch_index, variate_index, patch_index] = ( + center[:, None] + quantile_offsets + ) + + trend_slope = np.array([[5.0, 0.0], [0.0, 0.0]], dtype=np.float32) + trend_intercept = np.array([[4.0, 0.0], [0.0, 0.0]], dtype=np.float32) + apply_detrend = np.array([[True, False], [False, False]]) + target_mask = np.array([[True, False], [True, True]]) + nonnegative = np.array([[True, False], [False, False]]) + context_lengths = np.array([5, 3], dtype=np.int64) + horizon_lengths = np.array([6, 10], dtype=np.int64) + forecast_patch_counts = np.array([1, 2], dtype=np.int64) + + session = OnnxModelSession(package["stitcher"]) + outputs = session.run( + { + "logits": logits, + "trend_slope": trend_slope, + "trend_intercept": trend_intercept, + "apply_detrend": apply_detrend, + "target_mask": target_mask, + "nonnegative_mask": nonnegative, + "make_positive": np.array(True), + "context_lengths": context_lengths, + "horizon_lengths": horizon_lengths, + "context_patch_count": np.array(2, dtype=np.int64), + "forecast_patch_counts": forecast_patch_counts, + } + ) + session.close() + + weights = np.linspace(1.0, 0.0, 4, dtype=np.float32)[None, :, None] + expected = np.zeros((2, 2, 10, 3), dtype=np.float32) + for batch_index, patch_count in enumerate(forecast_patch_counts): + selected = logits[batch_index, :, 1 : 1 + patch_count] + if patch_count == 1: + stitched = selected[:, 0] + else: + stitched = np.concatenate( + [ + selected[:, 0, :4], + ( + weights[0] * selected[:, 0, 4:8] + + (1.0 - weights[0]) * selected[:, 1, :4] + ), + selected[:, 1, 4:8], + ], + axis=1, + ) + expected[batch_index, :, : horizon_lengths[batch_index]] = stitched[ + :, : horizon_lengths[batch_index] + ] + steps = np.arange(1, 11, dtype=np.float32)[None, None, :] + trend = ( + trend_slope[:, :, None] * steps / context_lengths[:, None, None] + + trend_intercept[:, :, None] + ) + expected += np.where(apply_detrend[:, :, None], trend, 0.0)[..., None] + expected = np.sort(expected, axis=-1) + expected = np.where(nonnegative[:, :, None, None], np.maximum(expected, 0.0), expected) + validity = target_mask[:, :, None] & ( + np.arange(10)[None, None, :] < horizon_lengths[:, None, None] + ) + expected = np.where(validity[..., None], expected, 0.0) + + np.testing.assert_allclose(outputs["quantile_forecasts"], expected, atol=1e-6) + np.testing.assert_allclose(outputs["point_forecast"], expected[..., 1], atol=1e-6) + np.testing.assert_array_equal(outputs["validity"], validity) + assert np.all(np.diff(outputs["quantile_forecasts"], axis=-1) >= 0) + assert np.all(outputs["quantile_forecasts"][0, 0] >= 0) + + +def test_full_fp32_pipeline_matches_linear_extrapolation(tiny_package) -> None: + _, module, package = tiny_package + for name, parameter in module.named_parameters(): + for component in package.values(): + initializer = component.graph.initializers.get(name) + if initializer is not None: + initializer.const_value = ir.tensor( + np.zeros(list(parameter.shape), dtype=np.float32) + ) + + context = np.array( + [ + [[0.0, 1.0, 2.0, 3.0, 4.0]], + [[0.0, 0.0, 10.0, 12.0, 14.0]], + ], + dtype=np.float32, + ) + _, outputs = _run_full_pipeline( + package, + { + "context_values": context, + "context_observed": np.array([[[True] * 5], [[False, False, True, True, True]]]), + "future_values": np.zeros((2, 1, 6), dtype=np.float32), + "future_observed": np.zeros((2, 1, 6), dtype=np.bool_), + "context_lengths": np.array([5, 3], dtype=np.int64), + "horizon_lengths": np.array([6, 2], dtype=np.int64), + "variate_roles": np.zeros((2, 1), dtype=np.int64), + }, + ) + + expected = np.array( + [[[5.0, 6.0, 7.0, 8.0, 9.0, 10.0]], [[16.0, 18.0, 0.0, 0.0, 0.0, 0.0]]], + dtype=np.float32, + ) + np.testing.assert_allclose(outputs["point_forecast"], expected, atol=1e-5) + np.testing.assert_allclose( + outputs["quantile_forecasts"], + np.repeat(expected[..., None], 3, axis=-1), + atol=1e-5, + ) + np.testing.assert_array_equal( + outputs["validity"], + np.array([[[True] * 6], [[True, True, False, False, False, False]]]), + ) + + +@pytest.mark.parametrize( + ("dtype", "numpy_dtype"), + [ + (ir.DataType.FLOAT, np.float32), + (ir.DataType.FLOAT16, np.float16), + ], +) +def test_tiny_model_runs_with_random_weights(dtype, numpy_dtype) -> None: + + config = _tiny_config() + config.dtype = dtype + model = build_from_module( + TimesFM3Model(config), + config, + task=TimeSeriesForecastingTask(), + ) + rng = np.random.default_rng(0) + for component in model.values(): + for initializer in component.graph.initializers.values(): + if initializer.const_value is None: + initializer.const_value = ir.tensor( + (rng.standard_normal(list(initializer.shape)) * 0.02).astype(numpy_dtype) + ) + values = rng.standard_normal((2, 3, 5, config.input_patch_len)).astype(numpy_dtype) + masks = np.zeros_like(values, dtype=np.bool_) + masks[:, :, 0] = True + patch_is_target = np.ones((2, 3, 5), dtype=np.bool_) + patch_cpm_mask = np.zeros((2, 5), dtype=np.bool_) + patch_cpm_mask[:, -2:] = True + + outputs = _run_pipeline( + model, + { + "values": values, + "masks": masks, + "patch_is_target": patch_is_target, + "patch_cpm_mask": patch_cpm_mask, + }, + ) + + assert outputs["logits"].shape == (2, 3, 5, 8, 3) + assert outputs["revin_mean"].shape == (2, 3, 5) + assert outputs["revin_std"].shape == (2, 3, 5) + assert np.isfinite(outputs["logits"]).all() + + if dtype == ir.DataType.FLOAT: + # A row's earlier forecast is invariant to another row forcing extra + # left context patches and right horizon patches in the padded batch. + context = np.zeros((2, 3, 9), dtype=np.float32) + context[0, :, -5:] = np.array( + [[1.0, 2.0, 4.0, 3.0, 5.0], [2.0, 1.0, 3.0, 2.0, 4.0], [0.0, 2.0, 1.0, 4.0, 3.0]] + ) + context[1] = rng.standard_normal((3, 9)).astype(np.float32) + future = np.zeros((2, 3, 10), dtype=np.float32) + future[:, 2] = rng.standard_normal((2, 10)).astype(np.float32) + common = { + "context_values": context, + "context_observed": np.ones_like(context, dtype=np.bool_), + "future_values": future, + "future_observed": np.ones_like(future, dtype=np.bool_), + "variate_roles": np.array([[0, 1, 2], [0, 1, 2]], dtype=np.int64), + } + _, padded = _run_full_pipeline( + model, + { + **common, + "context_lengths": np.array([5, 9], dtype=np.int64), + "horizon_lengths": np.array([3, 10], dtype=np.int64), + }, + ) + _, alone = _run_full_pipeline( + model, + {name: value[:1] for name, value in common.items()} + | { + "context_lengths": np.array([5], dtype=np.int64), + "horizon_lengths": np.array([3], dtype=np.int64), + }, + ) + np.testing.assert_allclose( + padded["quantile_forecasts"][0, :, :3], + alone["quantile_forecasts"][0], + rtol=1e-4, + atol=1e-4, + ) + + +def test_running_revin_and_cpm_refinement_match_reference() -> None: + + config = TimesFM3Config( + input_patch_len=2, + output_patch_len=4, + quantiles=(0.1, 0.5, 0.9), + num_layers=0, + model_dims=4, + transformer_hidden_dims=4, + num_heads=1, + max_variates=1, + use_variate_attention=False, + ) + model = build_from_module( + TimesFM3Model(config), + config, + task=TimeSeriesForecastingTask(), + ) + for component in model.values(): + for initializer in component.graph.initializers.values(): + if initializer.const_value is None: + initializer.const_value = ir.tensor( + np.zeros(list(initializer.shape), dtype=np.float32) + ) + output_bias = np.zeros(12, dtype=np.float32) + output_bias[1::3] = [1.0, 2.0, 3.0, 4.0] + model["model"].graph.initializers["output_head.bias"].const_value = ir.tensor(output_bias) + + outputs = _run_pipeline( + model, + { + "values": np.array( + [[[[1.0, 3.0], [5.0, 7.0], [0.0, 0.0], [0.0, 0.0]]]], + dtype=np.float32, + ), + "masks": np.array( + [[[[False, False], [False, False], [True, True], [True, True]]]] + ), + "patch_is_target": np.ones((1, 1, 4), dtype=np.bool_), + "patch_cpm_mask": np.array([[False, False, True, True]]), + }, + ) + + np.testing.assert_allclose( + outputs["revin_mean"], + np.array([[[2.0, 4.0, 4.0, 4.0]]], dtype=np.float32), + rtol=1e-6, + atol=1e-6, + ) + np.testing.assert_allclose( + outputs["revin_std"], + np.array([[[1.0, np.sqrt(5.0), np.sqrt(5.0), np.sqrt(5.0)]]], dtype=np.float32), + rtol=1e-6, + atol=1e-6, + ) + np.testing.assert_allclose( + outputs["logits"][0, 0, :, :, 1], + np.array( + [ + [3.0, 4.0, 5.0, 6.0], + [6.236068, 8.472136, 10.708204, 12.944272], + [7.618034, 10.118034, 12.618034, 15.118034], + [10.460805, 14.126524, 17.792244, 21.457964], + ], + dtype=np.float32, + ), + rtol=1e-6, + atol=1e-6, + ) + + +def test_registry_and_task_registration() -> None: + assert registry.get("timesfm3") is TimesFM3Model + assert registry.get_config_class("timesfm3") is TimesFM3Config + assert isinstance(get_task("time-series-forecasting"), TimeSeriesForecastingTask) diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index 8c73cff49..feaa2f487 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -105,6 +105,7 @@ "TASK_REGISTRY", "TTSTask", "T5TextEncoderTask", + "TimeSeriesForecastingTask", "VAETask", "VideoDenoisingTask", "VideoVAETask", @@ -202,6 +203,7 @@ from mobius.tasks._speech_to_text import SpeechToTextTask from mobius.tasks._ssm_causal_lm import SSM2CausalLMTask, SSMCausalLMTask from mobius.tasks._t5_text_encoder import T5TextEncoderTask +from mobius.tasks._time_series_forecasting import TimeSeriesForecastingTask from mobius.tasks._tts import TTSTask from mobius.tasks._vae import VAETask from mobius.tasks._video_denoising import VideoDenoisingTask @@ -255,6 +257,7 @@ "moshi-depformer": MoshiDepformerTask, "moshi-temporal": MoshiTemporalTask, "text-generation": CausalLMTask, + "time-series-forecasting": TimeSeriesForecastingTask, "smallthinker-gguf-text-generation": SmallThinkerGGUFCausalLMTask, "t5-text-encoding": T5TextEncoderTask, "deepseek-v4": DeepSeekV4Task, diff --git a/src/mobius/tasks/_time_series_forecasting.py b/src/mobius/tasks/_time_series_forecasting.py new file mode 100644 index 000000000..14a3cda10 --- /dev/null +++ b/src/mobius/tasks/_time_series_forecasting.py @@ -0,0 +1,312 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Padded-batch raw-series time-series forecasting task.""" + +from __future__ import annotations + +from typing import ClassVar + +import onnx_ir as ir +from onnxscript import nn + +from mobius._configs import BaseModelConfig +from mobius._model_package import ModelPackage +from mobius.tasks._base import ModelTask, _make_graph, _make_model + + +class TimeSeriesForecastingTask(ModelTask): + """Build the TimesFM patched-input forecasting pipeline. + + ``raw_preprocessor`` accepts right-aligned context tensors ``[B, V, C]`` + and left-aligned future-covariate tensors ``[B, V, H]``. Boolean observed + tensors distinguish missing values from padding; ``context_lengths`` and + ``horizon_lengths`` define each row's valid extents. ``variate_roles`` uses + 0=target, 1=past-only, and 2=past-future. Lengths must be positive and no + larger than their padded axes; invalid roles fail closed as masked variates. + + The package separates data-dependent preprocessing and postprocessing from + the transformer so ``model`` contains no ONNX control-flow operators and + can be captured independently by an execution provider. Symmetric + averaging and outer z-normalization remain opt-in host wrapper policies. + """ + + model_roles: ClassVar[dict[str, str]] = { + "raw_preprocessor": "encoder", + "preprocessor": "encoder", + "model": "encoder", + "postprocessor": "encoder", + "stitcher": "encoder", + } + + def build(self, module: nn.Module, config: BaseModelConfig) -> ModelPackage: + patch_len = getattr(config, "input_patch_len", None) + if not isinstance(patch_len, int) or isinstance(patch_len, bool): + raise TypeError( + "TimeSeriesForecastingTask requires a config with an integer input_patch_len" + ) + batch = ir.SymbolicDim("batch") + variates = ir.SymbolicDim("variates") + patches = ir.SymbolicDim("patches") + + raw_graph, builder = _make_graph(name="timesfm3_raw_preprocessor") + context = ir.SymbolicDim("context") + horizon = ir.SymbolicDim("horizon") + context_values = builder.input( + "context_values", + dtype=config.dtype, + shape=[batch, variates, context], + ) + context_observed = builder.input( + "context_observed", + dtype=ir.DataType.BOOL, + shape=[batch, variates, context], + ) + future_values = builder.input( + "future_values", + dtype=config.dtype, + shape=[batch, variates, horizon], + ) + future_observed = builder.input( + "future_observed", + dtype=ir.DataType.BOOL, + shape=[batch, variates, horizon], + ) + context_lengths = builder.input( + "context_lengths", + dtype=ir.DataType.INT64, + shape=[batch], + ) + horizon_lengths = builder.input( + "horizon_lengths", + dtype=ir.DataType.INT64, + shape=[batch], + ) + variate_roles = builder.input( + "variate_roles", + dtype=ir.DataType.INT64, + shape=[batch, variates], + ) + ( + raw_values, + raw_masks, + raw_patch_is_target, + raw_patch_cpm_mask, + trend_slope, + trend_intercept, + apply_detrend, + target_mask, + nonnegative_mask, + raw_context_lengths, + raw_horizon_lengths, + context_patch_count, + forecast_patch_counts, + ) = module.prepare_raw_series( + builder.op, + context_values, + context_observed, + future_values, + future_observed, + context_lengths, + horizon_lengths, + variate_roles, + ) + builder.add_output(raw_values, "values") + builder.add_output(raw_masks, "masks") + builder.add_output(raw_patch_is_target, "patch_is_target") + builder.add_output(raw_patch_cpm_mask, "patch_cpm_mask") + builder.add_output(trend_slope, "trend_slope") + builder.add_output(trend_intercept, "trend_intercept") + builder.add_output(apply_detrend, "apply_detrend") + builder.add_output(target_mask, "target_mask") + builder.add_output(nonnegative_mask, "nonnegative_mask") + builder.add_output(raw_context_lengths, "context_lengths") + builder.add_output(raw_horizon_lengths, "horizon_lengths") + builder.add_output(context_patch_count, "context_patch_count") + builder.add_output(forecast_patch_counts, "forecast_patch_counts") + + preprocess_graph, builder = _make_graph(name="timesfm3_preprocessor") + values = builder.input( + "values", + dtype=config.dtype, + shape=[batch, variates, patches, patch_len], + ) + masks = builder.input( + "masks", + dtype=ir.DataType.BOOL, + shape=[batch, variates, patches, patch_len], + ) + patch_is_target = builder.input( + "patch_is_target", + dtype=ir.DataType.BOOL, + shape=[batch, variates, patches], + ) + patch_cpm_mask = builder.input( + "patch_cpm_mask", + dtype=ir.DataType.BOOL, + shape=[batch, patches], + ) + model_inputs, patch_mask, running_count, running_mean, running_std = module.preprocess( + builder.op, + values=values, + masks=masks, + patch_is_target=patch_is_target, + patch_cpm_mask=patch_cpm_mask, + ) + builder.add_output(model_inputs, "model_inputs") + builder.add_output(patch_mask, "patch_mask") + builder.add_output(running_count, "revin_count") + builder.add_output(running_mean, "revin_mean") + builder.add_output(running_std, "revin_std") + + model_graph, builder = _make_graph(name="timesfm3_model") + model_inputs = builder.input( + "model_inputs", + dtype=config.dtype, + shape=[ + batch, + variates, + patches, + 2 * (config.input_patch_len + config.output_patch_len), + ], + ) + patch_mask = builder.input( + "patch_mask", + dtype=ir.DataType.BOOL, + shape=[batch, variates, patches], + ) + raw_logits = module.forecast(builder.op, model_inputs, patch_mask) + builder.add_output(raw_logits, "raw_logits") + + postprocess_graph, builder = _make_graph(name="timesfm3_postprocessor") + raw_logits = builder.input( + "raw_logits", + dtype=config.dtype, + shape=[ + batch, + variates, + patches, + config.output_patch_len * len(config.quantiles), + ], + ) + running_count = builder.input( + "revin_count", + dtype=ir.DataType.FLOAT, + shape=[batch, variates, patches], + ) + running_mean = builder.input( + "revin_mean", + dtype=ir.DataType.FLOAT, + shape=[batch, variates, patches], + ) + running_std = builder.input( + "revin_std", + dtype=ir.DataType.FLOAT, + shape=[batch, variates, patches], + ) + patch_cpm_mask = builder.input( + "patch_cpm_mask", + dtype=ir.DataType.BOOL, + shape=[batch, patches], + ) + logits, running_mean, running_std = module.postprocess( + builder.op, + raw_logits, + running_count, + running_mean, + running_std, + patch_cpm_mask, + ) + builder.add_output(logits, "logits") + builder.add_output(running_mean, "revin_mean") + builder.add_output(running_std, "revin_std") + + stitch_graph, builder = _make_graph(name="timesfm3_stitcher") + logits = builder.input( + "logits", + dtype=config.dtype, + shape=[ + batch, + variates, + patches, + config.output_patch_len, + len(config.quantiles), + ], + ) + trend_slope = builder.input( + "trend_slope", + dtype=ir.DataType.FLOAT, + shape=[batch, variates], + ) + trend_intercept = builder.input( + "trend_intercept", + dtype=ir.DataType.FLOAT, + shape=[batch, variates], + ) + apply_detrend = builder.input( + "apply_detrend", + dtype=ir.DataType.BOOL, + shape=[batch, variates], + ) + target_mask = builder.input( + "target_mask", + dtype=ir.DataType.BOOL, + shape=[batch, variates], + ) + nonnegative_mask = builder.input( + "nonnegative_mask", + dtype=ir.DataType.BOOL, + shape=[batch, variates], + ) + make_positive = builder.input( + "make_positive", + dtype=ir.DataType.BOOL, + shape=[], + ) + context_lengths = builder.input( + "context_lengths", + dtype=ir.DataType.INT64, + shape=[batch], + ) + horizon_lengths = builder.input( + "horizon_lengths", + dtype=ir.DataType.INT64, + shape=[batch], + ) + context_patch_count = builder.input( + "context_patch_count", + dtype=ir.DataType.INT64, + shape=[], + ) + forecast_patch_counts = builder.input( + "forecast_patch_counts", + dtype=ir.DataType.INT64, + shape=[batch], + ) + point_forecast, quantile_forecasts, validity = module.stitch_forecast( + builder.op, + logits, + trend_slope, + trend_intercept, + apply_detrend, + target_mask, + nonnegative_mask, + make_positive, + context_lengths, + horizon_lengths, + context_patch_count, + forecast_patch_counts, + ) + builder.add_output(point_forecast, "point_forecast") + builder.add_output(quantile_forecasts, "quantile_forecasts") + builder.add_output(validity, "validity") + return ModelPackage( + { + "raw_preprocessor": _make_model(raw_graph), + "preprocessor": _make_model(preprocess_graph), + "model": _make_model(model_graph), + "postprocessor": _make_model(postprocess_graph), + "stitcher": _make_model(stitch_graph), + }, + config=config, + ) diff --git a/tests/build_graph/_support.py b/tests/build_graph/_support.py index e5c932876..ba93a717c 100644 --- a/tests/build_graph/_support.py +++ b/tests/build_graph/_support.py @@ -205,6 +205,8 @@ def _make_params(configs: list[tuple[str, dict, bool]]) -> list: "sew", "sew-d", "sortformer", + # TimesFM 3 patched-input forecasting (co-located models/timesfm3_test.py). + "timesfm3", "speecht5", "unispeech", "unispeech-sat", diff --git a/tests/model_coverage_test.py b/tests/model_coverage_test.py index a93772804..cb193e5fb 100644 --- a/tests/model_coverage_test.py +++ b/tests/model_coverage_test.py @@ -306,6 +306,8 @@ def _all_registered_with_test_id() -> dict[str, str]: "mms": "CTC ASR model — tested via TestBuildMMSGraph", "fastconformer_rnnt": "NeMo .nemo RNN-T ASR — tested via tests/nemo_rnnt_integration_test.py", "sortformer": "NeMo .nemo speaker diarization — tested via tests/sortformer_integration_test.py", + "timesfm3": "Official 3.0 weights are non-commercial and cannot be redistributed; " + "the patched core graph is covered by src/mobius/models/timesfm3_test.py.", # --- Models requiring trust_remote_code --- "chatglm": "Requires trust_remote_code (custom HF modeling code)", "dots1": "Requires trust_remote_code (custom HF modeling code)",