Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Simplify implementation of nn.relu. #25331

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Simplify implementation of nn.relu. #25331

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
879a599
Simplify implementation of nn.relu.
carlosgmartin Dec 11, 2024
777e577
Testing.
carlosgmartin Dec 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 2 additions & 5 deletions jax/_src/nn/functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@

import jax
import jax.numpy as jnp
from jax import custom_jvp
from jax import lax
from jax._src import config
from jax._src import core
Expand All @@ -50,7 +49,6 @@ def __repr__(self):

# activations

@custom_jvp
@jax.jit
def relu(x: ArrayLike) -> Array:
r"""Rectified linear unit activation function.
Expand Down Expand Up @@ -83,9 +81,8 @@ def relu(x: ArrayLike) -> Array:
:func:`relu6`

"""
return jnp.maximum(x, 0)
# For behavior at 0, see https://dl.acm.org/doi/10.5555/3540261.3540297
relu.defjvps(lambda g, ans, x: lax.select(x > 0, g, lax.full_like(g, 0)))
z = lax.zeros_like_array(x)
return lax.select(lax.ge(x, z), x, z)

@jax.jit
def squareplus(x: ArrayLike, b: ArrayLike = 4) -> Array:
Expand Down