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

Jump to bottom

Improve Gamma sampling #12943

Closed
wants to merge 5 commits into from
Closed

Improve Gamma sampling #12943

wants to merge 5 commits into from

Conversation

adamgayoso
Copy link

@adamgayoso adamgayoso commented Oct 24, 2022 •
edited
Loading

This fixes:

  • a todo to use lax cond in order to speed up a computation during gamma sampling

@google-cla
Copy link

google-cla bot commented Oct 24, 2022

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@adamgayoso adamgayoso mentioned this pull request Oct 24, 2022
Open
@adamgayoso adamgayoso changed the title Use lax.cond in Gamma sampling instead of lax.bitwise_and Improve Gamma sampling Oct 24, 2022
@adamgayoso adamgayoso mentioned this pull request Oct 24, 2022
Closed
jakevdp
jakevdp reviewed Oct 24, 2022
View reviewed changes
@jakevdp jakevdp self-assigned this Oct 24, 2022
jakevdp
jakevdp approved these changes Oct 24, 2022
View reviewed changes
Copy link
Collaborator

@jakevdp jakevdp 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, though I'm a bit hesitant to merge this without understanding why lax.map was used in the first place, and confirming that that reason no longer is important. Have you been able to discern that?

@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Oct 24, 2022
@adamgayoso
Copy link
Author

adamgayoso commented Oct 24, 2022 •
edited
Loading

cc @fehiepsi -- looks like it was added when the grad was still implemented in this file.

I can try a quick benchmark for now?

@adamgayoso
Copy link
Author

@jakevdp actually it does indeed faster to map on cpu, I had a bug in my benchmark. Also would probably good to leave these potential changes in separate PRs. Therefore I changed this PR to only include the lax cond change

fehiepsi
fehiepsi reviewed Oct 24, 2022
View reviewed changes
jax/_src/random.py
lax.log(V))))))
lax.log(V)))))
false_fun = lambda: np.bool_(False)
cond = lax.cond(pred, true_fun, false_fun)
Copy link
Contributor

Choose a reason for hiding this comment

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

Does this change make things faster? Could you post some benchmark numbers?

@jakevdp
Copy link
Collaborator

jakevdp commented Nov 3, 2023

It seems like this has become stale. I'm going to close for now – feel free to re-open if you'd like to continue working on this!

@jakevdp jakevdp closed this Nov 3, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@fehiepsi fehiepsi fehiepsi left review comments

@jakevdp jakevdp jakevdp approved these changes

Assignees

@jakevdp jakevdp

Labels
pull ready Ready for copybara import and testing
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@adamgayoso @jakevdp @fehiepsi @kokoro-team