Skip to content

Fix softmax_cross_entropy to handle -inf logits correctly when corresponding label is 0. - #898

Merged
copybara-service[bot] merged 1 commit into
google-deepmind:mainfrom
carlosgmartin:softmax_cross_entropy_inf
Apr 7, 2024
Merged

Fix softmax_cross_entropy to handle -inf logits correctly when corresponding label is 0.#898
copybara-service[bot] merged 1 commit into
google-deepmind:mainfrom
carlosgmartin:softmax_cross_entropy_inf

Conversation

@carlosgmartin

Copy link
Copy Markdown
Contributor

@fabianp

fabianp commented Apr 3, 2024

Copy link
Copy Markdown
Member

Would perhaps using logsumexp (https://jax.readthedocs.io/en/latest/_autosummary/jax.scipy.special.logsumexp.html) be a solution here?

Also, please add tests :-)

@carlosgmartin

Copy link
Copy Markdown
Contributor Author

@fabianp Added tests.

log_softmax uses logsumexp internally: #896 (comment).

@vroulet vroulet left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks again!
I understood also my confusion, you are modifying cross_entropy and not cross_entropy_with_integer_labels so it makes a lot of sense. Sorry for the confusion and thanks again!

@copybara-service
copybara-service Bot merged commit b0c04dc into google-deepmind:main Apr 7, 2024
@carlosgmartin
carlosgmartin deleted the softmax_cross_entropy_inf branch April 7, 2024 22:48
@vroulet vroulet mentioned this pull request Apr 8, 2024
copybara-service Bot pushed a commit that referenced this pull request Apr 8, 2024
Added tests for correct gradient evaluations.

PiperOrigin-RevId: 622956379
copybara-service Bot pushed a commit that referenced this pull request Apr 9, 2024
Added tests for correct gradient evaluations.

PiperOrigin-RevId: 622956379
copybara-service Bot pushed a commit that referenced this pull request Apr 9, 2024
Added tests for correct gradient evaluations.

PiperOrigin-RevId: 623180578
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants