Skip to content

Add axis argument to smooth_labels function. #1492

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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add axis argument to smooth_labels function. #1492

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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: 6 additions & 1 deletion optax/losses/_smoothing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
# ==============================================================================
"""Smoothing functions."""

from typing import Union

from jax import typing
import jax.numpy as jnp
from optax._src import utils
Expand All @@ -22,6 +24,8 @@
def smooth_labels(
labels: typing.ArrayLike,
alpha: float,
*,
axis: Union[int, tuple[int, ...], None] = -1,
) -> jnp.ndarray:
"""Apply label smoothing.

Expand All @@ -32,6 +36,7 @@ def smooth_labels(
Args:
labels: One hot labels to be smoothed.
alpha: The smoothing factor.
axis: Axis or axes along which to compute.

Returns:
a smoothed version of the one hot input labels.
Expand All @@ -41,5 +46,5 @@ def smooth_labels(
<https://arxiv.org/abs/1906.02629>`_, 2019
"""
utils.check_subdtype(labels, jnp.floating)
num_categories = labels.shape[-1]
num_categories = jnp.size(labels, axis)
return (1.0 - alpha) * labels + alpha / num_categories
Loading