Support smaller context in KumoForecast

sdm/models/kumo/timeseries/forecast/model.py

@@ -40,8 +40,7 @@ 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.
-    Forecast horizons beyond 72 steps are extended autoregressively, up to
-    512 steps.
+    Forecast horizons beyond 72 steps are extended autoregressively.
 
     For longer observation histories, Domain-Aware Representation and
     Retrieval (DARR) can augment the direct forecast with patterns retrieved
@@ -181,19 +180,7 @@ def _forward(
                 f"Expected 'alpha' between zero and one (got {alpha})"
             )
 
-        if y_context.size(-2) < self.model.context_size:
-            raise ValueError(
-                f"Expected at least {self.model.context_size} context rows "
-                f"(got {y_context.size(-2)})"
-            )
-
         horizon = x_query.size(-2)
-        if not 1 <= horizon <= 512:
-            raise ValueError(
-                "Number of query rows must be between 1 and 512 steps "
-                f"(got {horizon})"
-            )
-
         context_size = self.model.context_size
         prediction_size = self.model.prediction_size
         memory_horizon = max(prediction_size, horizon)
@@ -218,6 +205,12 @@ def _forward(
 
         # Forecast:
         last_x = x[..., -self.model.context_size :]
+        if last_x.size(-1) < context_size:
+            last_x = torch.nn.functional.pad(
+                last_x,
+                pad=(context_size - last_x.size(-1), 0),
+                value=float("nan"),
+            )
         encoded, mean, scale = self.model.encode(last_x)  # [..., C, P, D]
         out = self.model.decode(encoded, mean, scale)
 

test/models/kumo/timeseries/forecast/test_model.py

@@ -14,10 +14,16 @@
 
 
 @withCUDA
+@pytest.mark.parametrize(
+    ("context_size", "horizon"),
+    [(1, 120), (37, 600), (1200, 600)],
+)
 @pytest.mark.parametrize("cross_channel", [False, True])
 @pytest.mark.parametrize("alpha", [0.0, 0.5, 1.0])
 def test_forward(
     monkeypatch: pytest.MonkeyPatch,
+    context_size: int,
+    horizon: int,
     cross_channel: bool,
     alpha: float,
     device: torch.device,
@@ -38,24 +44,24 @@ def test_forward(
 
     # Past and past-and-future covariates:
     x_context = TableTensor.from_tensor(
-        torch.randn(1024, 2, device=device),
+        torch.randn(context_size, 2, device=device),
         columns=["x1", "x2"],
     )
-    x_query = TableTensor.from_tensor(torch.randn(120, 0, device=device))
+    x_query = TableTensor.from_tensor(torch.randn(horizon, 0, device=device))
 
     # Target variates:
     y_context = TableTensor.from_tensor(
-        torch.randn(1024, 3, device=device),
+        torch.randn(context_size, 3, device=device),
         columns=["y1", "y2", "y3"],
     )
 
     with (
         nullcontext()
-        if cross_channel or alpha < 1
+        if cross_channel or (alpha < 1 and context_size >= 512 + horizon)
         else pytest.warns(RuntimeWarning, match="Historical covariates")
     ):
         out = model(x_context, y_context, x_query, alpha=alpha)
-    assert out.size() == (120, 3)
+    assert out.size() == (horizon, 3)
     assert out.device == device
     assert out.numerical.isfinite().all()
     assert "y1" in out.numerical_columns