Uh oh!
There was an error while loading.Please reload this page.
- Notifications
You must be signed in to change notification settings - Fork48.5k
Open
Description
Feature description
The softmax function does not have specification for the input type required to compute softmax. At the least, the function argument should have strict type checking, in place, to prevent crashing on non-supported inputs such as string or dict.
Other recommendations are:
- Numerical stability: subtract the max on the reduction axis before exponentiating to avoid overflow.
- Make the function N‑D and axis-aware: add an axis parameter (default -1) so it works for 1D/2D/3D tensors and is explicit about where probabilities sum to 1.
- Input normalization and dtype: convert inputs with np.asarray(..., dtype=float) so lists/tuples/ints behave correctly.
- Use keepdims=True on reductions so broadcasting keeps shapes correct.
- Add simple type hints and a clearer docstring describing inputs/outputs and axis semantics.
- Add unit tests (happy path + edge cases: scalar, single element, identical logits, large values, empty) and a smallmain demo.
- Defensive checks: raise meaningful error for empty input and document behavior for NaNs/Infs.
A sample improved version is as follows:
import numpy as npfrom typing import Optionaldef softmax(vector, axis: Optional[int] = -1) -> np.ndarray: """ Compute the softmax of `vector` along `axis` in a numerically-stable way. Parameters ---------- vector : array_like Input data (vector, matrix, tensor). Will be converted to float ndarray. axis : int or None, optional Axis along which to compute softmax. If None, compute softmax over the flattened array (single distribution). Default is -1 (last axis). 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). Raises ------ ValueError If input is empty. """ try: x = np.asarray(x, dtype=float) except Exception as e: raise ValueError(f"Could not convert input to float ndarray: {e}") if vector.size == 0: raise ValueError("softmax input must be non-empty") if axis is None: # flatten to single distribution vector_max = np.max(vector) e_vector = np.exp(vector - vector_max) return e_vector / e_vector.sum() # subtract max along axis with keepdims for numerical stability/broadcasting vector_max = np.max(vector, axis=axis, keepdims=True) e_vector = np.exp(vector - vector_max) denom = e_vector.sum(axis=axis, keepdims=True) return e_vector / denomif __name__ == "__main__": print(softmax((0,)))credit to Github Copilot