mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 21:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

softmax: document NaN outputs for infinite inputs

Browse Source
This commit is contained in:
Jake VanderPlas 2024-05-29 15:00:20 -07:00
parent 851a05dc1b
commit 75f570e8b0
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
jax/_src/nn/functions.py
View File

@ -512,6 +512,10 @@ def log_softmax(x: ArrayLike,
Returns:
An array.
Note:
If any input values are ``+inf``, the result will be all ``NaN``: this reflects the
fact that ``inf / inf`` is not well-defined in the context of floating-point math.
See also:
:func:`softmax`
"""
@ -557,6 +561,10 @@ def softmax(x: ArrayLike,
Returns:
An array.
Note:
If any input values are ``+inf``, the result will be all ``NaN``: this reflects the
fact that ``inf / inf`` is not well-defined in the context of floating-point math.
See also:
:func:`log_softmax`
"""