From 3207df5d032c97dd384df8af64ab8749b6fe68c5 Mon Sep 17 00:00:00 2001 From: Nishant Date: Mon, 19 Jan 2026 14:48:35 +0530 Subject: [PATCH] refactor: add type hints to softmax function --- maths/softmax.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/maths/softmax.py b/maths/softmax.py index 95c95e66f59e..23c4bf3c0239 100644 --- a/maths/softmax.py +++ b/maths/softmax.py @@ -13,25 +13,22 @@ import numpy as np -def softmax(vector): +def softmax(vector: np.ndarray | list | tuple) -> np.ndarray: """ Implements the softmax function Parameters: - vector (np.array,list,tuple): A numpy array of shape (1,n) - consisting of real values or a similar list,tuple - + vector (np.array | list | tuple): A numpy array of shape (1,n) + consisting of real values or a similar list, tuple Returns: - softmax_vec (np.array): The input numpy array after applying - softmax. + np.array: The input numpy array after applying softmax. - The softmax vector adds up to one. We need to ceil to mitigate for - precision - >>> float(np.ceil(np.sum(softmax([1,2,3,4])))) + The softmax vector adds up to one. We need to ceil to mitigate for precision + >>> float(np.ceil(np.sum(softmax([1, 2, 3, 4])))) 1.0 - >>> vec = np.array([5,5]) + >>> vec = np.array([5, 5]) >>> softmax(vec) array([0.5, 0.5])