From d983e58d4374ecebe8f9a53638adf035bf9bed19 Mon Sep 17 00:00:00 2001
From: Echo-Nie <nyxchaoji123@163.com>
Date: Mon, 17 Aug 2026 11:11:15 +0800
Subject: [PATCH 1/2] Validate softmax temperature

---
 tabfm/src/classifier_and_regressor.py      | 12 +++++++--
 tabfm/src/classifier_and_regressor_test.py | 29 ++++++++++++++++++++++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/tabfm/src/classifier_and_regressor.py b/tabfm/src/classifier_and_regressor.py
index e5ddc29..611564c 100644
--- a/tabfm/src/classifier_and_regressor.py
+++ b/tabfm/src/classifier_and_regressor.py
@@ -3363,12 +3363,20 @@ def softmax(
     Args:
       x: Input logit array of any shape.
       axis: Axis along which to compute softmax.
-      temperature: Scaling factor applied before the softmax; values < 1
-        produce a sharper distribution.
+      temperature: Strictly positive scaling factor applied before the softmax;
+        values < 1 produce a sharper distribution.
 
     Returns:
       Softmax probabilities with the same shape as ``x``.
+
+    Raises:
+      ValueError: If ``temperature`` is not greater than zero.
     """
+    if temperature <= 0:
+      raise ValueError(
+          f"temperature must be greater than 0. Got {temperature}."
+      )
+
     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..54ba4d2 100644
--- a/tabfm/src/classifier_and_regressor_test.py
+++ b/tabfm/src/classifier_and_regressor_test.py
@@ -704,6 +704,35 @@ def test_classifier_predict_oof_proba_raises_on_regression_model(self):
       classifier._predict_oof_proba(cv=2)
 
 
+class SoftmaxTemperatureTest(absltest.TestCase):
+
+  def test_non_positive_temperature_raises(self):
+    logits = np.array([[2.0, 1.0, 0.0]], dtype=np.float64)
+
+    for temperature in (0.0, -1.0):
+      with self.subTest(temperature=temperature):
+        with self.assertRaisesRegex(
+            ValueError, "temperature must be greater than 0"
+        ):
+          TabFMClassifier.softmax(
+              logits, axis=-1, temperature=temperature
+          )
+
+  def test_positive_temperature_returns_valid_probabilities(self):
+    logits = np.array([[2.0, 1.0, 0.0]], dtype=np.float64)
+
+    probabilities = TabFMClassifier.softmax(
+        logits, axis=-1, temperature=0.9
+    )
+
+    self.assertTrue(np.all(np.isfinite(probabilities)))
+    np.testing.assert_allclose(
+        probabilities.sum(axis=-1),
+        np.ones(logits.shape[0]),
+    )
+    self.assertEqual(np.argmax(probabilities, axis=-1).item(), 0)
+
+
 @unittest.skipUnless(HAS_JAX, "JAX is required")
 class CalibrationTest(absltest.TestCase):
 

From 9f120e0b2b21f1920498829476658fcca6b3b298 Mon Sep 17 00:00:00 2001
From: Echo-Nie <nyxchaoji123@163.com>
Date: Fri, 4 Sep 2026 15:32:49 +0800
Subject: [PATCH 2/2] Reject non-finite softmax temperatures

---
 tabfm/src/classifier_and_regressor.py      | 6 +++---
 tabfm/src/classifier_and_regressor_test.py | 6 +++---
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tabfm/src/classifier_and_regressor.py b/tabfm/src/classifier_and_regressor.py
index 611564c..e3190cd 100644
--- a/tabfm/src/classifier_and_regressor.py
+++ b/tabfm/src/classifier_and_regressor.py
@@ -3370,11 +3370,11 @@ def softmax(
       Softmax probabilities with the same shape as ``x``.
 
     Raises:
-      ValueError: If ``temperature`` is not greater than zero.
+      ValueError: If ``temperature`` is not finite or is not greater than zero.
     """
-    if temperature <= 0:
+    if not np.isfinite(temperature) or temperature <= 0:
       raise ValueError(
-          f"temperature must be greater than 0. Got {temperature}."
+          f"temperature must be finite and greater than 0. Got {temperature}."
       )
 
     x = x / temperature
diff --git a/tabfm/src/classifier_and_regressor_test.py b/tabfm/src/classifier_and_regressor_test.py
index 54ba4d2..81938f6 100644
--- a/tabfm/src/classifier_and_regressor_test.py
+++ b/tabfm/src/classifier_and_regressor_test.py
@@ -706,13 +706,13 @@ def test_classifier_predict_oof_proba_raises_on_regression_model(self):
 
 class SoftmaxTemperatureTest(absltest.TestCase):
 
-  def test_non_positive_temperature_raises(self):
+  def test_invalid_temperature_raises(self):
     logits = np.array([[2.0, 1.0, 0.0]], dtype=np.float64)
 
-    for temperature in (0.0, -1.0):
+    for temperature in (0.0, -1.0, np.nan, np.inf, -np.inf):
       with self.subTest(temperature=temperature):
         with self.assertRaisesRegex(
-            ValueError, "temperature must be greater than 0"
+            ValueError, "temperature must be finite and greater than 0"
         ):
           TabFMClassifier.softmax(
               logits, axis=-1, temperature=temperature