From 375723c96cb318ed6b9b9a5ae01911a0288dee96 Mon Sep 17 00:00:00 2001
From: Sylvester Kaczmarek
 <16242628+sylvesterkaczmarek@users.noreply.github.com>
Date: Tue, 8 Sep 2026 22:54:50 +0100
Subject: [PATCH] Validate softmax temperature

---
 tabfm/src/classifier_and_regressor.py      |  3 +++
 tabfm/src/classifier_and_regressor_test.py | 18 ++++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/tabfm/src/classifier_and_regressor.py b/tabfm/src/classifier_and_regressor.py
index b18f929..c981076 100644
--- a/tabfm/src/classifier_and_regressor.py
+++ b/tabfm/src/classifier_and_regressor.py
@@ -3370,6 +3370,9 @@ def softmax(
     Returns:
       Softmax probabilities with the same shape as ``x``.
     """
+    if temperature <= 0:
+      raise ValueError('temperature must be greater than 0')
+
     x = x / temperature
     # Subtract max for numerical stability
     x_max = np.max(x, axis=axis, keepdims=True)
diff --git a/tabfm/src/classifier_and_regressor_test.py b/tabfm/src/classifier_and_regressor_test.py
index 419c983..6d0502f 100644
--- a/tabfm/src/classifier_and_regressor_test.py
+++ b/tabfm/src/classifier_and_regressor_test.py
@@ -35,6 +35,24 @@
 
 # pylint: disable=invalid-name
 
+
+class SoftmaxTest(absltest.TestCase):
+
+  def test_rejects_non_positive_temperature(self):
+    logits = np.array([[1.0, 2.0, 3.0]])
+
+    for temperature in (0.0, -1.0):
+      with self.assertRaisesRegex(ValueError, 'temperature must be greater than 0'):
+        TabFMClassifier.softmax(logits, temperature=temperature)
+
+  def test_positive_temperature_returns_probabilities(self):
+    logits = np.array([[1.0, 2.0, 3.0]])
+
+    probabilities = TabFMClassifier.softmax(logits, temperature=0.9)
+
+    np.testing.assert_allclose(probabilities.sum(axis=-1), 1.0)
+
+
 class EnsembleGeneratorTest(absltest.TestCase):
 
   def test_permute_categorical_structure(self):