1import numpy as np2import torch34from chronos_x.models.model import MRPT, MRPTConfig5from chronos_x.data.windows import F_BAR, F_GLOBAL6from chronos_x.models.student_t import masked_student_t_nll, student_t_logpdf789def _batch(B=3, H=60, LA=390, LB=390, LC=260):10 g = torch.Generator().manual_seed(0)11 return {"xA": torch.randn(B, LA, F_BAR, generator=g), "mA": torch.ones(B, LA, dtype=torch.uint8),12 "xB": torch.randn(B, LB, F_BAR, generator=g), "mB": torch.ones(B, LB, dtype=torch.uint8),13 "xC": torch.randn(B, LC, F_BAR, generator=g), "mC": torch.ones(B, LC, dtype=torch.uint8),14 "glob": torch.randn(B, F_GLOBAL, generator=g), "fut_cal": torch.rand(B, H, 4, generator=g),15 "sigma": torch.full((B,), 3e-4), "y": torch.randn(B, H, generator=g) * 3e-4,16 "y_mask": torch.ones(B, H, dtype=torch.uint8), "sym": torch.zeros(B, dtype=torch.long)}171819def test_forward_shapes_and_param_counts():20 for preset, lo, hi in (("tiny", 1e4, 5e5), ("small", 2e6, 6e6), ("spec", 5e6, 12e6), ("base", 15e6, 30e6)):21 m = MRPT(MRPTConfig.preset(preset))22 n = m.n_params()23 assert lo <= n <= hi, (preset, n)24 m = MRPT(MRPTConfig.preset("tiny"))25 out = m(_batch())26 assert out["mu"].shape == (3, 60) and out["scale"].shape == (3, 60) and out["nu"].shape == (3, 60)27 assert (out["scale"] > 0).all() and (out["nu"] > 2).all()28 loss = masked_student_t_nll(_batch()["y"], _batch()["y_mask"], out)29 assert torch.isfinite(loss)30 loss.backward()313233def test_variants_run():34 b = _batch()35 for kw in (dict(branches=["A"]), dict(branches=["A", "B"]), dict(fusion="cross"), dict(head="gaussian"),36 dict(head="quantile"), dict(regime="none"), dict(regime="learned"), dict(n_symbols=4),37 dict(copula_rank=3)):38 m = MRPT(MRPTConfig.preset("tiny", **kw))39 out = m(b)40 key = "quantiles" if kw.get("head") == "quantile" else "mu"41 assert torch.isfinite(out[key]).all()424344def test_student_t_logpdf_matches_scipy():45 from scipy import stats46 y = torch.tensor([[-0.002, 0.0, 0.0031]]); mu = torch.tensor([[0.0001, 0.0, -0.0002]])47 sc = torch.tensor([[0.001, 0.002, 0.0005]]); nu = torch.tensor([[3.0, 5.5, 20.0]])48 ours = student_t_logpdf(y, mu, sc, nu).numpy()49 ref = stats.t.logpdf(y.numpy(), nu.numpy(), loc=mu.numpy(), scale=sc.numpy())50 np.testing.assert_allclose(ours, ref, rtol=1e-5)51