Support smaller context in KumoForecast

docs/source/nim.md

@@ -9,10 +9,10 @@ Thus, an application can call the model without loading its weights or managing
 NIMs for structured data models follow SDM's unified prediction interface.
 We have published the following NIMs:
 
-| Model ID          | Model                                  | Deployment Instructions                                                            |
-| ----------------- | -------------------------------------- | ---------------------------------------------------------------------------------- |
-| `kumo-tabular`    | {py:class}`~sdm.models.KumoTabular`    | [Kumo Tabular on NVIDIA Build](https://build.nvidia.com/nvidia/kumo-tabular)       |
-| `kumo-relational` | {py:class}`~sdm.models.KumoRelational` | [Kumo Relational on NVIDIA Build](https://build.nvidia.com/nvidia/kumo-relational) |
+| Model ID          | Model                                  | NVIDIA Build                                              | NGC Catalog                                                                              |
+| ----------------- | -------------------------------------- | --------------------------------------------------------- | ---------------------------------------------------------------------------------------- |
+| `kumo-tabular`    | {py:class}`~sdm.models.KumoTabular`    | [Deploy](https://build.nvidia.com/nvidia/kumo-tabular)    | [Container](https://catalog.ngc.nvidia.com/orgs/nim/nvidia/containers/kumo-tabular/-)    |
+| `kumo-relational` | {py:class}`~sdm.models.KumoRelational` | [Deploy](https://build.nvidia.com/nvidia/kumo-relational) | [Container](https://catalog.ngc.nvidia.com/orgs/nim/nvidia/containers/kumo-relational/-) |
 
 ## The NIM API
 

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.
+    Shorter histories are left-padded with missing values, which are excluded
+    from normalization statistics and replaced with zeros after normalization.
     Forecast horizons beyond 72 steps are extended autoregressively, up to
     512 steps.
 
@@ -181,12 +183,6 @@ 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(
@@ -218,6 +214,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

@@ -16,10 +16,12 @@
 @withCUDA
 @pytest.mark.parametrize("cross_channel", [False, True])
 @pytest.mark.parametrize("alpha", [0.0, 0.5, 1.0])
+@pytest.mark.parametrize("context_size", [1, 37, 1024])
 def test_forward(
     monkeypatch: pytest.MonkeyPatch,
     cross_channel: bool,
     alpha: float,
+    context_size: int,
     device: torch.device,
 ) -> None:
     monkeypatch.setattr(
@@ -38,20 +40,20 @@ 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))
 
     # 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 >= 632)
         else pytest.warns(RuntimeWarning, match="Historical covariates")
     ):
         out = model(x_context, y_context, x_query, alpha=alpha)
@@ -61,3 +63,33 @@ def test_forward(
     assert "y1" in out.numerical_columns
     assert "y2" in out.numerical_columns
     assert "y3" in out.numerical_columns
+
+    if context_size < 512:
+        padded_x_context = TableTensor.from_tensor(
+            torch.nn.functional.pad(
+                x_context.numerical,
+                pad=(0, 0, 512 - context_size, 0),
+                value=float("nan"),
+            ),
+            columns=x_context.numerical_columns,
+        )
+        padded_y_context = TableTensor.from_tensor(
+            torch.nn.functional.pad(
+                y_context.numerical,
+                pad=(0, 0, 512 - context_size, 0),
+                value=float("nan"),
+            ),
+            columns=y_context.numerical_columns,
+        )
+        with (
+            nullcontext()
+            if cross_channel
+            else pytest.warns(RuntimeWarning, match="Historical covariates")
+        ):
+            expected = model(
+                padded_x_context,
+                padded_y_context,
+                x_query,
+                alpha=alpha,
+            )
+        torch.testing.assert_close(out.numerical, expected.numerical)