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

Fix documentation for nn.elu, nn.celu, and lax.expm1. (#3116)

Browse Source
This commit is contained in:
Ed Schmerling 2020-05-15 20:51:53 -07:00 committed by GitHub
parent e675f804ff
commit 510af1de64
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
jax/lax/lax.py
View File

@ -145,7 +145,7 @@ def exp(x: Array) -> Array:
return exp_p.bind(x)
def expm1(x: Array) -> Array:
r"""Elementwise :math:`e^{x - 1}`."""
r"""Elementwise :math:`e^{x} - 1`."""
return expm1_p.bind(x)
def log(x: Array) -> Array:
@ -5257,4 +5257,4 @@ def _canonicalize_axis(axis, num_dims):
raise ValueError(
"axis {} is out of bounds for array of dimension {}".format(
axis, num_dims))
return axis
return axis

4
jax/nn/functions.py
View File

@ -95,7 +95,7 @@ def elu(x, alpha=1.0):
.. math::
\mathrm{elu}(x) = \begin{cases}
x, & x > 0\\
\alpha \exp(x - 1), & x \le 0
\alpha \left(\exp(x) - 1\right), & x \le 0
\end{cases}
"""
safe_x = jnp.where(x > 0, 0., x)
@ -138,7 +138,7 @@ def celu(x, alpha=1.0):
.. math::
\mathrm{celu}(x) = \begin{cases}
x, & x > 0\\
\alpha \exp(\frac{x}{\alpha} - 1), & x \le 0
\alpha \left(\exp(\frac{x}{\alpha}) - 1\right), & x \le 0
\end{cases}
For more information, see