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.

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

Extend log_softmax stabilization rewrite to graphs with indexing and expand_dims #352

Merged
merged 2 commits into from
Jun 29, 2023

Conversation

ricardoV94
Copy link
Member

@ricardoV94 ricardoV94 commented Jun 20, 2023

This makes the rewrite applicable to cases like the PMF of a Categorical, when probabilities are provided as a softmax and there is an expand_dims or indexing in between the log and softmax

Also removes the axis warning in softmax functions

@ricardoV94 ricardoV94 changed the title Implement rewrite to stabilize log(softmax[idx]) Implement rewrite to stabilize log(softmax[idx]) Jun 20, 2023
@ricardoV94 ricardoV94 force-pushed the stabilize_log_subtensor_softmax branch from 225add5 to 7c54695 Compare June 20, 2023 12:46
@ricardoV94 ricardoV94 force-pushed the stabilize_log_subtensor_softmax branch from 7c54695 to 9aa9d06 Compare June 20, 2023 14:24
@ricardoV94 ricardoV94 changed the title Implement rewrite to stabilize log(softmax[idx]) Extend log_softmax stabilization rewrite Jun 20, 2023
@ricardoV94 ricardoV94 added enhancement New feature or request graph rewriting labels Jun 20, 2023
@ricardoV94 ricardoV94 force-pushed the stabilize_log_subtensor_softmax branch from 9aa9d06 to a31e827 Compare June 20, 2023 14:41
@ricardoV94 ricardoV94 marked this pull request as ready for review June 20, 2023 14:55
@ricardoV94 ricardoV94 force-pushed the stabilize_log_subtensor_softmax branch 2 times, most recently from e3e2611 to 74c2882 Compare June 20, 2023 15:27
@ricardoV94 ricardoV94 force-pushed the stabilize_log_subtensor_softmax branch from 74c2882 to 29dd1d1 Compare June 21, 2023 07:24
@ricardoV94 ricardoV94 changed the title Extend log_softmax stabilization rewrite Extend log_softmax stabilization rewrite to graphs with indexing and expand_dims Jun 21, 2023
@ricardoV94 ricardoV94 changed the title Extend log_softmax stabilization rewrite to graphs with indexing and expand_dims Extend log_softmax stabilization rewrite to graphs with indexing and expand_dims Jun 21, 2023
Copy link
Member

@aseyboldt aseyboldt left a comment

Choose a reason for hiding this comment

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

Looks good to me :-)

@ricardoV94 ricardoV94 merged commit ce0b503 into pymc-devs:main Jun 29, 2023
@ricardoV94 ricardoV94 deleted the stabilize_log_subtensor_softmax branch July 13, 2023 12:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request graph rewriting
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants