Fix #93: Reject non-positive softmax temperatures

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 finite or is not greater than zero.
     """
+    if not np.isfinite(temperature) or temperature <= 0:
+      raise ValueError(
+          f"temperature must be finite and greater than 0. Got {temperature}."
+      )
+
     x = x / temperature
     # Subtract max for numerical stability
     x_max = np.max(x, axis=axis, keepdims=True)

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_invalid_temperature_raises(self):
+    logits = np.array([[2.0, 1.0, 0.0]], dtype=np.float64)
+
+    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 finite and 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):