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

Fixes the JAX implementation of CELU returning NaN gradients for input

Browse Source 
values >= 88.7229.

When a JAX where() op is used to avoid a NaN or undefined value, reverse
differentiation can still return NaN even though the NaN input is not selected
by the conditional:

https://jax.readthedocs.io/en/latest/faq.html#gradients-contain-nan-where-using-where

This change uses jnp.maximum and jnp.minimum to compute CELU without producing an undefined value.

PiperOrigin-RevId: 461678140
This commit is contained in:
jax authors 2022-07-18 11:57:29 -07:00
parent d98d5ddce5
commit b0805a8a31
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -200,7 +200,7 @@ def celu(x: Array, alpha: Array = 1.0) -> Array:
x : input array
alpha : array or scalar (default: 1.0)
"""
return jnp.where(x > 0, x, alpha * jnp.expm1(x / alpha))
return jnp.maximum(x, 0.0) + alpha * jnp.expm1(jnp.minimum(x, 0.0) / alpha)
@jax.jit
def selu(x: Array) -> Array: