Fix softmax: improve type hints, exception handling, lint compliance

Author: shreyasNaik0101

maths/softmax.py

@@ -1,10 +1,10 @@
 import numpy as np
-from typing import Optional, Union
+from typing import Optional
 
 
 def softmax(
-    vector: Union[np.ndarray, list, tuple],
-    axis: Optional[int] = -1
+    vector: np.ndarray | list | tuple,
+    axis: int | None = -1
 ) -> np.ndarray:
     """
     Compute the softmax of `vector` along `axis` in a numerically-stable way.
@@ -20,8 +20,8 @@ def softmax(
     Returns
     -------
     np.ndarray
-        Same shape as `vector`, with softmax applied along `axis`. Probabilities sum to 1
-        along `axis` (or to 1 overall if axis is None).
+        Same shape as `vector`, with softmax applied along `axis`. Probabilities
+        sum to 1 along `axis` (or to 1 overall if axis is None).
 
     Raises
     ------
@@ -30,8 +30,9 @@ def softmax(
     """
     try:
         vector = np.asarray(vector, dtype=float)
-    except Exception as e:
-        raise ValueError(f"Could not convert input to float ndarray: {e}")
+    except TypeError as e:
+        msg = f"Could not convert input to float ndarray: {e}"
+        raise ValueError(msg)
 
     if vector.size == 0:
         raise ValueError("softmax input must be non-empty")
@@ -49,42 +50,41 @@ def softmax(
     return e_vector / denom
 
 
-# Example unit tests
 def _test_softmax():
     import numpy.testing as npt
-    
+
     # Typical 1D input
     result = softmax([1, 2, 3])
     npt.assert_almost_equal(result.sum(), 1)
-    
+
     # Typical 2D, axis=-1
     result = softmax([[1, 2, 3], [4, 5, 6]])
     npt.assert_almost_equal(result.sum(axis=-1).tolist(), [1, 1])
-    
+
     # Scalar input
     result = softmax([0])
     npt.assert_almost_equal(result, [1.0])
-    
+
     # Identical values
     result = softmax([5, 5])
     npt.assert_almost_equal(result, [0.5, 0.5])
-    
+
     # Large values for numeric stability
     result = softmax([1000, 1001])
     npt.assert_almost_equal(result.sum(), 1)
-    
+
     # axis=None flatten
     data = np.array([[1, 2], [3, 4]])
     flat_result = softmax(data, axis=None)
     npt.assert_almost_equal(flat_result.sum(), 1)
-    
+
     # Empty input error
     try:
         softmax([])
-        assert False, "Expected ValueError for empty input"
+        raise AssertionError("Expected ValueError for empty input")
     except ValueError:
         pass
-    
+
     print("All tests passed.")