CI deploy 0e23c622

src/adapters.py

@@ -192,13 +192,26 @@ class ChronosAdapter(ForecastAdapter):
     BOLT_CHUNK = 256
     T5_CHUNK = 16
     T5_NUM_SAMPLES = 20
+    CHRONOS2_CHUNK = 128
 
     @property
     def _is_bolt(self) -> bool:
         return "bolt" in (self.model_id or "").lower()
 
+    @property
+    def _is_chronos2(self) -> bool:
+        """Chronos-2 (`amazon/chronos-2`).
+
+        Detected from the id rather than from the loaded pipeline class so the
+        chunk size is known before the weights are fetched.
+        """
+        name = (self.model_id or "").lower()
+        return "chronos-2" in name or "chronos2" in name
+
     @property
     def chunk_size(self) -> int:
+        if self._is_chronos2:
+            return self.CHRONOS2_CHUNK
         return self.BOLT_CHUNK if self._is_bolt else self.T5_CHUNK
 
     def predict(self, context_windows: np.ndarray) -> Forecast:
@@ -211,7 +224,12 @@ class ChronosAdapter(ForecastAdapter):
             ctx = ctx[None, :]
 
         q_levels = list(DEFAULT_QUANTILES)
-        extra = {} if self._is_bolt else {"num_samples": self.T5_NUM_SAMPLES}
+        # Only the original T5 Chronos samples paths; Bolt and Chronos-2 emit
+        # quantiles directly.
+        if self._is_bolt or self._is_chronos2:
+            extra = {}
+        else:
+            extra = {"num_samples": self.T5_NUM_SAMPLES}
         step = max(1, self.chunk_size)
         parts = []
 
@@ -222,7 +240,7 @@ class ChronosAdapter(ForecastAdapter):
                 quantiles, _mean = self._model.predict_quantiles(
                     tensors, prediction_length=1, quantile_levels=q_levels, **extra,
                 )
-            parts.append(quantiles.float().cpu().numpy()[:, 0, :])
+            parts.append(self._to_array(quantiles))
             del quantiles, tensors
             _release(self.device)
 
@@ -230,6 +248,48 @@ class ChronosAdapter(ForecastAdapter):
         return Forecast(q10=arr[:, 0], q50=arr[:, 1], q90=arr[:, 2],
                         context_len=ctx.shape[1])
 
+    @staticmethod
+    def _to_array(quantiles) -> np.ndarray:
+        """Normalise a `predict_quantiles` result to `(batch, n_quantiles)`.
+
+        The two shapes differ and the difference is not cosmetic:
+
+          Bolt / T5      one stacked tensor, `(batch, horizon, quantiles)`
+          Chronos-2      a *list* of per-item tensors, each
+                         `(n_variates, horizon, quantiles)`
+
+        Calling `.float()` on the list raises `AttributeError`, which is
+        exactly what `amazon/chronos-2` did before this existed. Horizon is
+        always 1 here, and these are univariate price series, so the leading
+        variate axis is taken at index 0.
+        """
+        if isinstance(quantiles, (list, tuple)):
+            rows = []
+            for item in quantiles:
+                array = ChronosAdapter._as_numpy(item)
+                # (n_variates, horizon, quantiles) -> (quantiles,)
+                while array.ndim > 1:
+                    array = array[0]
+                rows.append(array)
+            return np.vstack(rows)
+
+        return ChronosAdapter._as_numpy(quantiles)[:, 0, :]
+
+    @staticmethod
+    def _as_numpy(value) -> np.ndarray:
+        """Tensor -> ndarray, by duck typing rather than an isinstance check.
+
+        Deliberately does not import torch: the test suite runs offline with no
+        GPU stack installed, and a shape-normalising helper that cannot be
+        tested without a 2GB dependency would not get tested.
+        """
+        if hasattr(value, "float") and hasattr(value, "cpu"):
+            return np.asarray(value.float().cpu().numpy())
+        if isinstance(value, np.ndarray):
+            return value
+        raise AdapterError(
+            f"unexpected predict_quantiles result: {type(value).__name__}")
+
 
 # --------------------------------------------------------------------------
 # TimesFM

src/config.py

@@ -172,6 +172,14 @@ SEED_MODELS: dict[str, ModelSpec] = {
         ModelSpec("chronos-t5-small", "amazon/chronos-t5-small", "chronos",
                   "Chronos T5 Small", context_len=512),
 
+        # Chronos-2. Loads through the same adapter and the already-pinned
+        # chronos-forecasting 2.3.1, but `predict_quantiles` returns a list of
+        # per-item tensors rather than one stacked tensor -- see
+        # `ChronosAdapter._to_array`. Seedable, so it costs GPU quota on the
+        # next seed run; drop it from SEEDABLE_MODELS if that is not wanted yet.
+        ModelSpec("chronos-2", "amazon/chronos-2", "chronos",
+                  "Chronos-2", context_len=512),
+
         # Naive baselines, deliberately first-class. A forecasting model that
         # cannot beat "tomorrow looks like today" is not worth deploying, and
         # the leaderboard should make that impossible to miss.

tests/test_adapters.py

@@ -343,3 +343,73 @@ def test_chronos_100_step_run_is_schema_valid():
     assert out["inference_version"].nunique() == 1
     assert out["inference_version"].iloc[0] != config.PLACEHOLDER_VERSION
     assert np.isfinite(out[["q10", "q50", "q90"]].to_numpy()).all()
+
+
+# --------------------------------------------------------------------------
+# Chronos-2
+# --------------------------------------------------------------------------
+#
+# `predict_quantiles` returns a different shape per pipeline, and the
+# difference is not cosmetic:
+#
+#   Bolt / T5   one stacked tensor, (batch, horizon, quantiles)
+#   Chronos-2   a LIST of per-item tensors, (n_variates, horizon, quantiles)
+#
+# Calling .float() on the list raises AttributeError, which is what
+# amazon/chronos-2 did before `_to_array` existed.
+
+
+def test_bolt_style_tensor_is_normalised():
+    import numpy as np
+
+    from src.adapters import ChronosAdapter
+
+    stacked = np.zeros((4, 1, 3), dtype="float32")
+    stacked[:, 0, 1] = 5.0                     # q50 for every row
+    out = ChronosAdapter._to_array(stacked)
+    assert out.shape == (4, 3)
+    assert np.allclose(out[:, 1], 5.0)
+
+
+def test_chronos2_list_of_tensors_is_normalised():
+    import numpy as np
+
+    from src.adapters import ChronosAdapter
+
+    # Four items, each (n_variates=1, horizon=1, quantiles=3).
+    listed = [np.array([[[1.0, 2.0, 3.0]]], dtype="float32") for _ in range(4)]
+    out = ChronosAdapter._to_array(listed)
+    assert out.shape == (4, 3), "Chronos-2's list shape was not handled"
+    assert np.allclose(out[0], [1.0, 2.0, 3.0])
+
+
+def test_an_unexpected_result_type_raises_a_named_error():
+    from src.adapters import AdapterError, ChronosAdapter
+
+    with pytest.raises(AdapterError):
+        ChronosAdapter._to_array({"not": "a tensor"})
+
+
+@pytest.mark.parametrize("model_id,is_c2,chunk_attr", [
+    ("amazon/chronos-2", True, "CHRONOS2_CHUNK"),
+    ("amazon/chronos-bolt-base", False, "BOLT_CHUNK"),
+    ("amazon/chronos-t5-small", False, "T5_CHUNK"),
+])
+def test_chunk_size_per_variant(model_id, is_c2, chunk_attr):
+    from src.adapters import ChronosAdapter
+
+    chunk = getattr(ChronosAdapter, chunk_attr)
+    adapter = ChronosAdapter.__new__(ChronosAdapter)
+    adapter.model_id = model_id
+    assert adapter._is_chronos2 is is_c2
+    assert adapter.chunk_size == chunk
+
+
+def test_chronos2_does_not_ask_for_samples():
+    """Only the original T5 Chronos samples paths; asking Chronos-2 for
+    `num_samples` is an unexpected keyword."""
+    from src.adapters import ChronosAdapter
+
+    adapter = ChronosAdapter.__new__(ChronosAdapter)
+    adapter.model_id = "amazon/chronos-2"
+    assert adapter._is_bolt is False and adapter._is_chronos2 is True