From 9efe799c75082f8c715aac5ab8c28fccd44b841e Mon Sep 17 00:00:00 2001 From: Aditi Gautam Date: Mon, 26 Oct 2026 09:21:13 -0700 Subject: [PATCH 1/2] Add DARR retrieval-augmented Kumo forecasting Signed-off-by: Aditi Gautam --- docs/source/api/models.rst | 17 ++ examples/timeseries/kumo_darr.py | 33 +++ .../kumo/timeseries/forecast/__init__.py | 2 + sdm/models/kumo/timeseries/forecast/darr.py | 272 ++++++++++++++++++ sdm/models/kumo/timeseries/forecast/model.py | 14 +- sdm/models/kumo/timeseries/forecast/recipe.py | 10 +- .../kumo/timeseries/forecast/test_darr.py | 272 ++++++++++++++++++ 7 files changed, 615 insertions(+), 5 deletions(-) create mode 100644 examples/timeseries/kumo_darr.py create mode 100644 sdm/models/kumo/timeseries/forecast/darr.py create mode 100644 test/models/kumo/timeseries/forecast/test_darr.py diff --git a/docs/source/api/models.rst b/docs/source/api/models.rst index 9a7b49bc2..189584acc 100644 --- a/docs/source/api/models.rst +++ b/docs/source/api/models.rst @@ -78,3 +78,20 @@ Models {% for name in api_names("sdm.models") if name not in api_names("sdm.models.base") %} ~sdm.models.{{ name }} {% endfor %} + +Retrieval-Augmented Forecasting +------------------------------ + +Reuse a pretrained forecaster with an explicitly fitted historical memory: + +.. code-block:: python + + from sdm.models import KumoForecast + from sdm.models.kumo.timeseries.forecast import DARR + + darr = DARR(KumoForecast(device="cuda"), alpha=0.01) + darr.fit(reference, horizon=72) # Observations known at the forecast origin. + forecast = darr(context) # Same columns; at least 512 observed rows. + +.. autoclass:: sdm.models.kumo.timeseries.forecast.DARR + :members: fit, forward, clear diff --git a/examples/timeseries/kumo_darr.py b/examples/timeseries/kumo_darr.py new file mode 100644 index 000000000..150d61909 --- /dev/null +++ b/examples/timeseries/kumo_darr.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pandas as pd +import torch + +import sdm +from sdm.models.kumo.timeseries.forecast import DARR + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +model = sdm.models.KumoForecast(device=device) +darr = DARR(model, k=64, temperature=0.05, alpha=0.01) + +url = "https://autogluon.s3.amazonaws.com/datasets/timeseries/m4_hourly/train.csv" +frame = pd.read_csv(url) +series = frame[frame["item_id"] == "H170"].sort_values("timestamp") +table = sdm.TableTensor.from_tensor( + torch.tensor( + series[["target"]].to_numpy(), dtype=torch.float32, device=device + ), + columns=["target"], +) + +# Keep the final 72 observations strictly out of both memory and query history. +history, held_out = table[:-72], table[-72:] +darr.fit(history, horizon=72, stride=72) +forecast = darr(history[-512:]) +print( + "DARR MAE:", (forecast.numerical - held_out.numerical).abs().mean().item() +) + +# Memory is reused across queries. For rolling evaluation, never fit it using +# observations later than the earliest evaluated query's forecast origin. diff --git a/sdm/models/kumo/timeseries/forecast/__init__.py b/sdm/models/kumo/timeseries/forecast/__init__.py index 2a769a8ab..c89999e06 100644 --- a/sdm/models/kumo/timeseries/forecast/__init__.py +++ b/sdm/models/kumo/timeseries/forecast/__init__.py @@ -4,8 +4,10 @@ """KumoForecast.""" from sdm.models.kumo.timeseries.forecast.model import KumoForecast +from sdm.models.kumo.timeseries.forecast.darr import DARR __all__ = [ "KumoForecast", + "DARR", ] diff --git a/sdm/models/kumo/timeseries/forecast/darr.py b/sdm/models/kumo/timeseries/forecast/darr.py new file mode 100644 index 000000000..56d70c58b --- /dev/null +++ b/sdm/models/kumo/timeseries/forecast/darr.py @@ -0,0 +1,272 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import cast + +import torch +import torch.nn.functional as F +from torch import Tensor + +from sdm import Stype, TableTensor +from sdm.models.kumo.timeseries.forecast.model import KumoForecast +from sdm.models.kumo.timeseries.forecast.recipe import standardizer + + +class DARR(torch.nn.Module): + """Retrieval-augmented forecasting with a :class:`KumoForecast` encoder. + + Fit a memory of historical windows and their observed continuations, then + retrieve the nearest windows using cosine similarity of mean-pooled + encoder embeddings. The output is ``alpha * direct + (1 - alpha) * + retrieved``, in the input's units, for every numerical column. + + Reference and query series must have the same variables, units and sampling + interval. Rows must be ordered. Only supply reference observations that + would have been available at the query's forecast origin: this API does not + inspect timestamps or remove overlapping evaluation data automatically. + Batched reference series contribute to one shared memory, without windows + crossing series boundaries. Query batch dimensions remain independent. + + Uses the checkpoint's frozen normalization, not a user-fitted recipe. + Non-finite continuation values are excluded independently at each channel + and horizon step; when all selected values are missing, the retrieval + prediction is the checkpoint standardizer's mean. Missing history follows + :class:`KumoForecast` normalization. Equal-similarity neighbor ordering is + unspecified. This is an inference-only wrapper; call ``eval()`` before use. + + The wrapped model's weights are shared. Refit after changing weights. + Memory moves with ``to()`` but is not included in ``state_dict()``; refit + after loading weights. No additional checkpoint or dependencies are needed. + + Ported from the Apache-2.0 Kumo-TS implementation at + ``af74798bbf7d020d3298f545fb9be383560a18ce`` (``Kumo-Forecast/sdk``). + + Args: + model: The pretrained forecaster to reuse. + k: Maximum number of historical neighbors to retrieve. + temperature: Softmax temperature for cosine similarities. ``None`` or + a negative value uses uniform weights; zero uses ``1e-12``. + alpha: Weight of the direct forecast, between zero and one. The SDK + default is ``0.01``. One gives direct-only predictions. + """ + + _keys: Tensor | None + _values: Tensor | None + + def __init__( + self, + model: KumoForecast, + *, + k: int = 64, + temperature: float | None = 0.05, + alpha: float = 0.01, + ) -> None: + super().__init__() + if k < 1: + raise ValueError("Expected at least one neighbor") + if not 0 <= alpha <= 1: + raise ValueError("Expected 'alpha' between zero and one") + + self.model = model.model + self.standardizer = standardizer().to(self.model.lin.weight.device) + self.k = k + self.temperature = temperature + self.alpha = alpha + self.register_buffer("_keys", None, persistent=False) + self.register_buffer("_values", None, persistent=False) + self._columns: tuple[str, ...] | None = None + self.eval() + + @torch.no_grad() + def fit( + self, + reference: Tensor | TableTensor, # [..., R, C] + *, + horizon: int = 72, + stride: int = 72, + batch_size: int = 8, + ) -> None: + """Replace retrieval memory using only complete historical windows. + + Args: + reference: Ordered numerical observations of shape ``[..., R, C]``. + Leading dimensions identify independent reference series. + Every continuation must be known before predicting a query. + horizon: Number of future steps to retrieve, between 1 and 512. + Beyond 72 steps, direct forecasts are autoregressive. + Windows require ``512 + max(72, horizon)`` reference rows. + stride: Step between reference windows, defaulting to the native + forecast horizon, as in the SDK. + batch_size: Maximum number of reference windows encoded together. + """ + self._check_eval() + if not 1 <= horizon <= 512: + raise ValueError("Expected 'horizon' between 1 and 512") + table = self._as_table(reference) + length = self.model.context_size + max( + self.model.prediction_size, horizon + ) + if table.size(-2) < length: + raise ValueError(f"Expected at least {length} reference rows") + + x = self.standardizer(table).numerical + x = x.reshape(-1, x.size(-2), x.size(-1)) + num_windows = (x.size(-2) - length) // stride + 1 + keys, values = [], [] + offsets = torch.arange(length, device=x.device) + # Gather bounded batches without materializing all overlapping windows. + for start in range(0, x.size(0) * num_windows, batch_size): + index = torch.arange( + start, + min(start + batch_size, x.size(0) * num_windows), + device=x.device, + ) + rows = (index % num_windows).unsqueeze(-1) * stride + offsets + window = x[(index // num_windows).unsqueeze(-1), rows] + history = window[:, : self.model.context_size].transpose(-1, -2) + encoded, _, _ = self.model.encode(history) + keys.append(_embedding(encoded)) + values.append( + window[ + :, + self.model.context_size : self.model.context_size + + horizon, + ].float() + ) + + self._keys = torch.cat(keys) + self._values = torch.cat(values) + self._columns = table.numerical_columns + + @torch.no_grad() + def forward( + self, + context: Tensor | TableTensor, # [..., R, C] + *, + batch_size: int = 8, + ) -> TableTensor: # [..., horizon, C] + """Forecast every column using a previously fitted reference memory. + + Args: + context: Ordered observed history of shape ``[..., R, C]``. Uses + the last 512 rows. Named columns are aligned to the fitted + reference; plain tensors use positional order. Returns columns + in reference order, with all leading query dimensions retained. + batch_size: Maximum number of query windows processed together. + + Returns: + Numerical forecasts in original units, with the fitted horizon. + The output retains the query's device and floating-point dtype. + Select target columns from the output when other columns are + historical covariates; covariates are also forecast for rollouts. + """ + self._check_eval() + if self._keys is None or self._values is None: + raise RuntimeError("Call 'fit()' with reference history first") + table = self._as_table(context, columns=self._columns) + if table.size(-2) < self.model.context_size: + raise ValueError( + f"Expected at least {self.model.context_size} context rows" + ) + + x = self.standardizer(table[..., -self.model.context_size :, :]) + history = x.numerical.transpose(-1, -2) + history = history.reshape(-1, history.size(-2), history.size(-1)) + horizon = self._values.size(-2) + predictions = [] + for window in history.split(batch_size): + encoded, mean, scale = self.model.encode(window) + if self.alpha != 1: + retrieved = _knn_forecast( + query=_embedding(encoded), + keys=self._keys, + values=self._values, + k=self.k, + temperature=self.temperature, + ) + if self.alpha == 0: + predictions.append(retrieved) + continue + + direct = self.model.revin.inverse( + x=self.model.head(encoded.flatten(-2)), + mean=mean, + scale=scale, + ) + rollout = [direct] + for _ in range( + self.model.prediction_size, horizon, self.model.prediction_size + ): + window = torch.cat([window, direct], dim=-1) + window = window[..., -self.model.context_size :] + direct = self.model(window) + rollout.append(direct) + direct = torch.cat(rollout, dim=-1)[..., :horizon].transpose( + -1, -2 + ) + predictions.append( + direct + if self.alpha == 1 + else self.alpha * direct + (1 - self.alpha) * retrieved + ) + + out = torch.cat(predictions).reshape(*table.size()[:-2], horizon, -1) + out = TableTensor.from_tensor(out, columns=self._columns) + out = self.standardizer.inverse_transform(out) + return cast(TableTensor, out.to(dtype=context.dtype)) + + def clear(self) -> None: + """Discard reference memory without changing model weights.""" + self._keys = self._values = None + self._columns = None + + def _check_eval(self) -> None: + """Require deterministic inference embeddings for retrieval.""" + if self.training or self.model.training: + raise RuntimeError("DARR is inference-only; call 'eval()' first") + + def _as_table( + self, + x: Tensor | TableTensor, + columns: tuple[str, ...] | None = None, + ) -> TableTensor: + """Normalize inputs and align reference and query variables.""" + if not isinstance(x, TableTensor): + x = TableTensor.from_tensor(x, columns=columns) + if x.size(-1) == 0 or x.active_stypes != {Stype.numerical}: + raise ValueError("Expected one or more numerical columns only") + if columns is not None: + if set(x.numerical_columns) != set(columns): + raise ValueError("Context columns must match the reference") + x = x.select_columns(columns) + return cast(TableTensor, x.to(dtype=self.model.lin.weight.dtype)) + + +def _embedding(encoded: Tensor) -> Tensor: + """Match the SDK's channel-then-patch mean pooling and cosine keys.""" + pooled = encoded.mean(dim=-3).mean(dim=-2).float() + pooled = pooled.nan_to_num(nan=0, posinf=0, neginf=0) + return F.normalize(pooled, dim=-1, eps=1e-12) + + +def _knn_forecast( + query: Tensor, # [B, D], unit-normalized + keys: Tensor, # [N, D], unit-normalized + values: Tensor, # [N, H, C] + k: int, + temperature: float | None, +) -> Tensor: # [B, H, C] + """Retrieve cosine neighbors and average their observed continuations.""" + with torch.amp.autocast(query.device.type, enabled=False): + similarity, index = (query.float() @ keys.float().t()).topk( + k=min(k, keys.size(0)), dim=-1 + ) + if temperature is None or temperature < 0: + weights = torch.ones_like(similarity) / similarity.size(-1) + else: + weights = (similarity / max(temperature, 1e-12)).softmax(dim=-1) + neighbors = values[index].float() # [B, K, H, C] + valid = neighbors.isfinite() + weights = weights[..., None, None] * valid + total = (weights * neighbors.masked_fill(~valid, 0)).sum(dim=1) + return total / weights.sum(dim=1).clamp_min(1e-12) diff --git a/sdm/models/kumo/timeseries/forecast/model.py b/sdm/models/kumo/timeseries/forecast/model.py index 5ac402cfe..3af4e4a6c 100644 --- a/sdm/models/kumo/timeseries/forecast/model.py +++ b/sdm/models/kumo/timeseries/forecast/model.py @@ -40,6 +40,8 @@ class KumoForecast(ICLModel): Optional cross-channel mixing lets forecasts use dependencies between variates after temporal encoding. The forecasting head predicts the native up to 72-step horizon directly from a 512-step history. + For retrieval-augmented forecasting over historical continuations, wrap + this model with :class:`~sdm.models.kumo.timeseries.forecast.DARR`. Within the :class:`~sdm.models.ICLModel` protocol, :class:`KumoForecast` treats rows as ordered time steps. In a default forward pass, @@ -282,7 +284,16 @@ def forward( self, x: Tensor, # [..., C, R] ) -> Tensor: # [..., C, prediction_size] + """Forecast from the shared patch encoder.""" + encoded, mean, scale = self.encode(x) + out = self.head(encoded.flatten(-2)) + return self.revin.inverse(out, mean, scale) + def encode( + self, + x: Tensor, # [..., C, R] + ) -> tuple[Tensor, Tensor, Tensor]: + """Return encoded patches and the instance-normalization statistics.""" mean, scale = self.revin.compute_stats(x) x = self.revin(x, mean, scale) x = x.nan_to_num(nan=0, posinf=0, neginf=0) @@ -293,5 +304,4 @@ def forward( if self.cross_channel is not None: x = self.cross_channel(x) # [..., C, N, D] - x = self.head(x.flatten(-2)) # [..., C, prediction_size] - return self.revin.inverse(x, mean, scale) + return x, mean, scale diff --git a/sdm/models/kumo/timeseries/forecast/recipe.py b/sdm/models/kumo/timeseries/forecast/recipe.py index ec8300b07..6ba6c4376 100644 --- a/sdm/models/kumo/timeseries/forecast/recipe.py +++ b/sdm/models/kumo/timeseries/forecast/recipe.py @@ -4,13 +4,17 @@ import sdm.processing as sp -def default_recipe() -> sp.Recipe: # noqa: D103 +def standardizer() -> sp.Standardize: + """Return the frozen checkpoint normalization shared by both modes.""" # Frozen statistics from standardizer.pkl in nvidia/Kumo-Forecast: mean = -5.671202659606934 scale = 8.693312644958496 + return sp.Standardize().with_state(mean, scale) + +def default_recipe() -> sp.Recipe: # noqa: D103 return sp.Recipe( - features=sp.Standardize().with_state(mean, scale), - target=sp.Standardize().with_state(mean, scale), + features=standardizer(), + target=standardizer(), output=sp.AverageEstimators(), ) diff --git a/test/models/kumo/timeseries/forecast/test_darr.py b/test/models/kumo/timeseries/forecast/test_darr.py new file mode 100644 index 000000000..0e9238559 --- /dev/null +++ b/test/models/kumo/timeseries/forecast/test_darr.py @@ -0,0 +1,272 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from functools import partial + +import numpy as np +import pytest +import torch + +from sdm import TableTensor +from sdm.models import KumoForecast +from sdm.models.kumo.timeseries.forecast import DARR +from sdm.models.kumo.timeseries.forecast import model as forecasting +from sdm.models.kumo.timeseries.forecast.darr import _embedding, _knn_forecast +from sdm.testing import withCUDA + + +@pytest.fixture +def model(monkeypatch: pytest.MonkeyPatch) -> KumoForecast: + """Build the same forecast architecture at a small test scale.""" + monkeypatch.setattr( + forecasting, + "_KumoForecast", + partial( + forecasting._KumoForecast, + context_size=16, + prediction_size=8, + patch_size=4, + channels=16, + hidden_channels=24, + num_layers=2, + num_heads=4, + cross_channel_heads=4, + ), + ) + return KumoForecast(pretrained=False) + + +@withCUDA +@pytest.mark.parametrize("temperature", [None, -1.0, 0.0, 0.05, 2.0]) +@pytest.mark.parametrize("k", [1, 3, 20]) +def test_retrieval_matches_numpy( + temperature: float | None, + k: int, + device: torch.device, +) -> None: + """Match the upstream cosine/top-k/weighted-continuation equations.""" + # Distinct similarities avoid depending on unspecified top-k tie order. + keys = np.array([[1, 0], [0, 1], [-1, 0], [0, -1]], dtype=np.float32) + query = np.array([[0.6, 0.8], [-0.8, 0.6]], dtype=np.float32) + values = np.arange(32, dtype=np.float32).reshape(4, 4, 2) + values[0, 0, 0] = np.nan + values[1, 0, 0] = np.inf + values[:, 3, 1] = np.nan + similarity = query @ keys.T + index = np.argsort(-similarity, axis=-1)[:, : min(k, len(keys))] + scores = np.take_along_axis(similarity, index, axis=-1) + if temperature is None or temperature < 0: + weights = np.ones_like(scores) / scores.shape[-1] + else: + scores = scores / max(temperature, 1e-12) + weights = np.exp(scores - scores.max(axis=-1, keepdims=True)) + weights /= weights.sum(axis=-1, keepdims=True) + neighbors = values[index] + valid = np.isfinite(neighbors) + weights = weights[..., None, None] * valid + expected = (weights * np.where(valid, neighbors, 0)).sum(axis=1) + expected /= np.maximum(weights.sum(axis=1), 1e-12) + + out = _knn_forecast( + query=torch.from_numpy(query).to(device), + keys=torch.from_numpy(keys).to(device), + values=torch.from_numpy(values).to(device), + k=k, + temperature=temperature, + ) + torch.testing.assert_close(out, torch.from_numpy(expected).to(device)) + assert out.device == device + + +def test_embedding_pooling() -> None: + """Pool channels and then patches as the SDK does, including zero keys.""" + encoded = torch.arange(48, dtype=torch.float32).reshape(2, 2, 3, 4) + pooled = encoded.mean(dim=1).sum(dim=1) / 3 + expected = pooled / pooled.norm(dim=-1, keepdim=True) + torch.testing.assert_close(_embedding(encoded), expected) + torch.testing.assert_close( + _embedding(torch.zeros_like(encoded)), torch.zeros(2, 4) + ) + + +@withCUDA +def test_retrieval_uses_query_history( + model: KumoForecast, + device: torch.device, +) -> None: + """Retrieve the matching continuation from a shared two-series memory.""" + model = model.to(device) + time = torch.arange(16, device=device).float() + history = torch.stack([time.square(), time.flip(0).square()]) + future = torch.stack([time[:8] + 100, time[:8] - 100]) + reference = torch.cat([history, future], dim=-1).unsqueeze(-1) + darr = DARR(model, k=1, alpha=0) + darr.fit(reference, horizon=8, batch_size=1) + out = darr(reference[:, :16], batch_size=2).numerical + torch.testing.assert_close(out, reference[:, 16:]) + + +@withCUDA +def test_retrieval_autocast(device: torch.device) -> None: + """Keep similarity and weight accumulation in float32 under autocast.""" + keys = torch.tensor([[1.0, 0.0], [0.0, 1.0]], device=device) + values = torch.tensor([[[1.0]], [[3.0]]], device=device) + query = torch.tensor([[0.50001, 0.5]], device=device) + expected = _knn_forecast(query, keys, values, k=2, temperature=1e-5) + with torch.amp.autocast(device.type, dtype=torch.bfloat16): + out = _knn_forecast(query, keys, values, k=2, temperature=1e-5) + torch.testing.assert_close(out, expected, rtol=0, atol=0) + assert out.dtype == torch.float32 + + +@withCUDA +@pytest.mark.parametrize("cross_channel", [False, True]) +@pytest.mark.parametrize("horizon", [3, 8, 19]) +def test_blending_and_rollout( + model: KumoForecast, + cross_channel: bool, + horizon: int, + device: torch.device, +) -> None: + """Match both blending endpoints and SDK-style native-block rollouts.""" + model = model.to(device) + if not cross_channel: + model.model.cross_channel = None + reference = torch.randn(35, 2, device=device) + context = torch.randn(2, 3, 20, 2, device=device) + darr = DARR(model, temperature=None, alpha=0) + darr.fit(reference, horizon=horizon, stride=4, batch_size=2) + retrieved = darr(context, batch_size=2).numerical + starts = range(0, 35 - 16 - max(8, horizon) + 1, 4) + expected = torch.stack( + [reference[s + 16 : s + 16 + horizon] for s in starts] + ).mean(0) + torch.testing.assert_close(retrieved, expected.expand(2, 3, -1, -1)) + + # Independent use of the existing public forecaster for each native block. + current = context.reshape(-1, 20, 2) + rollout = [] + for _ in range(0, horizon, 8): + prediction = model( + x_context=current[..., :0], + y_context=current, + x_query=current.new_empty(current.size(0), 8, 0), + num_estimators=1, + ).numerical + rollout.append(prediction) + current = torch.cat([current, prediction], dim=-2)[..., -16:, :] + expected_direct = torch.cat(rollout, dim=-2)[..., :horizon, :].reshape( + 2, 3, horizon, 2 + ) + darr.alpha = 1 + direct = darr(context).numerical + torch.testing.assert_close(direct, expected_direct, rtol=2e-5, atol=2e-5) + darr.alpha = 0.25 + mixed = darr(context).numerical + torch.testing.assert_close(mixed, 0.25 * direct + 0.75 * retrieved) + assert not mixed.requires_grad + assert mixed.device == device + + +@withCUDA +def test_reference_boundaries( + model: KumoForecast, device: torch.device +) -> None: + """Never cross independent series or borrow incomplete future windows.""" + model = model.to(device) + reference = ( + torch.stack( + [ + torch.arange(28, device=device), + torch.arange(28, device=device) + 100, + ] + ) + .float() + .unsqueeze(-1) + ) + darr = DARR(model, temperature=None, alpha=0) + darr.fit(reference, horizon=3, stride=8, batch_size=1) + out = darr(reference[:, :16]).numerical + # Even a short horizon requires a complete native 8-step continuation. + expected = reference[:, 16:19].mean(dim=0).expand(2, -1, -1) + torch.testing.assert_close(out, expected) + # A query's extra leading history must not augment or alter fitted memory. + longer = torch.cat([reference * 1000, reference[:, :16]], dim=-2) + torch.testing.assert_close(darr(longer).numerical, out) + darr.fit(reference, horizon=3, stride=8, batch_size=8) + torch.testing.assert_close(darr(longer).numerical, out) + + +def test_column_alignment_and_refit(model: KumoForecast) -> None: + """Align named channels, preserve units, and replace memory on refit.""" + reference = TableTensor.from_tensor(torch.randn(24, 2), columns=["a", "b"]) + darr = DARR(model, k=1, alpha=0) + darr.fit(reference, horizon=8) + context = reference[:16].select_columns(["b", "a"]) + out = darr(context) + assert out.numerical_columns == ("a", "b") + torch.testing.assert_close(out.numerical, reference[16:].numerical) + torch.testing.assert_close( + darr(reference[:16].numerical).numerical, out.numerical + ) + with pytest.raises(ValueError, match="columns must match"): + darr(TableTensor.from_tensor(torch.randn(16, 2), columns=["a", "c"])) + reference = reference.replace_blocks(numerical=reference.numerical + 50) + darr.fit(reference, horizon=8) + torch.testing.assert_close( + darr(context).numerical, reference[16:].numerical + ) + darr.clear() + with pytest.raises(RuntimeError, match="fit"): + darr(context) + + +def test_missing_continuations(model: KumoForecast) -> None: + """Renormalize finite neighbors; all-missing values use the frozen mean.""" + reference = torch.randn(2, 24, 2) + reference[0, 16:, 0] = torch.nan + reference[:, 16:, 1] = torch.inf + darr = DARR(model, temperature=None, alpha=0) + darr.fit(reference, horizon=8) + out = darr(reference[1, :16]).numerical + torch.testing.assert_close(out[:, 0], reference[1, 16:, 0]) + torch.testing.assert_close(out[:, 1], torch.full((8,), -5.671202659606934)) + + +@withCUDA +def test_memory_conversion(model: KumoForecast, device: torch.device) -> None: + """Move memory with weights; exclude observations from state dicts.""" + reference = torch.randn(24, 1) + darr = DARR(model, alpha=0) + before = set(darr.state_dict()) + darr.fit(reference, horizon=8) + assert set(darr.state_dict()) == before + darr = darr.to(device=device, dtype=torch.float64) + out = darr(reference[:16].to(device)).numerical + assert out.dtype == reference.dtype + torch.testing.assert_close( + out, + reference[16:].to(device=device, dtype=out.dtype), + rtol=2e-6, + atol=2e-6, + ) + assert out.device == device + + +def test_input_contract(model: KumoForecast) -> None: + """Reject absent memory, insufficient history, and incompatible data.""" + darr = DARR(model) + with pytest.raises(RuntimeError, match="fit"): + darr(torch.zeros(16, 1)) + with pytest.raises(ValueError, match="24 reference rows"): + darr.fit(torch.zeros(23, 1), horizon=3) + with pytest.raises(ValueError, match="horizon"): + darr.fit(torch.zeros(24, 1), horizon=0) + with pytest.raises(ValueError, match="numerical columns"): + darr.fit(TableTensor.from_tensor(torch.ones(24, 1, dtype=torch.long))) + darr.fit(torch.zeros(24, 1), horizon=8) + with pytest.raises(ValueError, match="16 context rows"): + darr(torch.zeros(15, 1)) + darr.train() + with pytest.raises(RuntimeError, match="inference-only"): + darr(torch.zeros(16, 1)) From b4504681f2424930f5bf72c3cd391a315965427f Mon Sep 17 00:00:00 2001 From: Aditi Gautam Date: Mon, 26 Oct 2026 10:41:00 -0700 Subject: [PATCH 2/2] Fix DARR documentation heading underline Signed-off-by: Aditi Gautam --- docs/source/api/models.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/api/models.rst b/docs/source/api/models.rst index 189584acc..4d60c5dfb 100644 --- a/docs/source/api/models.rst +++ b/docs/source/api/models.rst @@ -80,7 +80,7 @@ Models {% endfor %} Retrieval-Augmented Forecasting ------------------------------- +------------------------------- Reuse a pretrained forecaster with an explicitly fitted historical memory: