From ad3ec93617fdcc8d2e27d1129b2f5f0f46a0e8c7 Mon Sep 17 00:00:00 2001
From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com>
Date: Mon, 27 Jul 2026 22:03:31 +0200
Subject: [PATCH] fix(architectures): initialize Trompt's init_rec prompt
 parameter

nn.Parameter(torch.empty(P, d_model)) was never initialized, so the
initial prompt representation kept whatever happened to be in the
allocation. When that memory held garbage the model produced NaN logits:
this is the cause of the intermittent CI failures in
test_experimental_classifier_fit_predict_evaluate[TromptClassifier] and
test_experimental_lss_fit_predict_evaluate[TromptLSS], where the LSS run
died with 'Expected parameter loc ... to satisfy the constraint Real(),
but found invalid values: tensor([nan, ...])'.

The sibling prompt embeddings in ImportanceGetter use exactly this
pattern with a following torch.nn.init.normal_(std=0.01); init_rec just
missed the call. Matching that convention also makes the model
seed-reproducible, which it could not be before.

Part of #423

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 deeptab/architectures/experimental/trompt.py |  3 ++
 tests/test_trompt_init.py                    | 44 ++++++++++++++++++++
 2 files changed, 47 insertions(+)
 create mode 100644 tests/test_trompt_init.py

diff --git a/deeptab/architectures/experimental/trompt.py b/deeptab/architectures/experimental/trompt.py
index ec2c480..0df214f 100644
--- a/deeptab/architectures/experimental/trompt.py
+++ b/deeptab/architectures/experimental/trompt.py
@@ -61,6 +61,9 @@ def __init__(
         self.cells = nn.ModuleList(TromptCell(feature_information, config) for _ in range(config.n_cycles))
         self.decoder = TromptDecoder(config.d_model, num_classes)
         self.init_rec = nn.Parameter(torch.empty(config.P, config.d_model))
+        # Matches the prompt-embedding initialisation in ImportanceGetter; without
+        # it the parameter keeps whatever torch.empty returned (NaNs in CI).
+        torch.nn.init.normal_(self.init_rec, std=0.01)
         self.n_cycles = config.n_cycles
 
     def forward(self, *data):
diff --git a/tests/test_trompt_init.py b/tests/test_trompt_init.py
new file mode 100644
index 0000000..1138eac
--- /dev/null
+++ b/tests/test_trompt_init.py
@@ -0,0 +1,44 @@
+"""Regression test: Trompt's init_rec prompt parameter must be initialized.
+
+``nn.Parameter(torch.empty(...))`` keeps whatever was in the allocation. When
+that memory held garbage the model produced NaN logits -- observed as
+intermittent CI failures of the TromptClassifier/TromptLSS model tests.
+"""
+
+import numpy as np
+import pandas as pd
+import torch
+
+from deeptab.models.experimental import TromptRegressor
+
+
+def _data(n=40, seed=0):
+    rng = np.random.RandomState(seed)
+    return pd.DataFrame({"a": rng.randn(n), "b": rng.randn(n)}), rng.randn(n)
+
+
+def test_init_rec_is_initialized_not_left_as_raw_allocation(monkeypatch):
+    """Simulate an allocation full of garbage; the constructor must overwrite it."""
+    real_empty = torch.empty
+
+    def poisoned_empty(*args, **kwargs):
+        tensor = real_empty(*args, **kwargs)
+        if tensor.is_floating_point():
+            tensor.fill_(float("nan"))
+        return tensor
+
+    monkeypatch.setattr(torch, "empty", poisoned_empty)
+
+    X, y = _data()
+    model = TromptRegressor()
+    model.build_model(X, y)
+
+    estimator = model._task_model.estimator  # type: ignore[union-attr]
+    assert torch.isfinite(estimator.init_rec).all()
+
+
+def test_trompt_predictions_are_finite():
+    X, y = _data()
+    model = TromptRegressor()
+    model.fit(X, y, max_epochs=1, batch_size=16, accelerator="cpu")
+    assert np.isfinite(model.predict(X)).all()