Skip to content

support null shape in jax.random.poisson #8390

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
Merged
merged 1 commit into from
Dec 30, 2021
Merged

support null shape in jax.random.poisson #8390

merged 1 commit into from
Dec 30, 2021

Conversation

rakki-18
Copy link
Contributor

Fixes #7861

Issue

jax.random.poisson(key,lam)

was giving errors, because default shape was not being broadcasted to lam

Changes introduced

  • If shape argument is not given to jax.random.poisson,
    it takes the shape of lam.shape
  • Added tests to cover the case of default shape
jax.random.poisson(key,lam)

@google-cla google-cla bot added the cla: yes label Oct 28, 2021
@rakki-18 rakki-18 marked this pull request as draft October 29, 2021 11:05
@rakki-18 rakki-18 marked this pull request as ready for review October 29, 2021 11:11
@rakki-18
Copy link
Contributor Author

rakki-18 commented Oct 29, 2021 •
edited
Loading

@jakevdp Could you please review this?

@jakevdp jakevdp self-assigned this Nov 30, 2021
jakevdp
jakevdp requested changes Nov 30, 2021
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.

Hi - sorry this fell off my radar. Thanks for the contribution! A few small changes and then we can get this merged.

@rakki-18
Copy link
Contributor Author

@jakevdp Really sorry for the delay. I have made the changes requested. Could you please review it now?

jakevdp
jakevdp approved these changes Dec 28, 2021
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, thanks!

@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Dec 28, 2021
@jakevdp
Copy link
Collaborator

jakevdp commented Dec 28, 2021

Looks like there are a couple flake8 issues – once these are addressed we can get this merged. Thanks!

@rakki-18
Copy link
Contributor Author

Looks like there are a couple flake8 issues – once these are addressed we can get this merged. Thanks!

@jakevdp I have fixed them now

@jakevdp
Copy link
Collaborator

jakevdp commented Dec 28, 2021

Thanks - can you squash all the changes into a single commit please?

@rakki-18
Copy link
Contributor Author

Thanks - can you squash all the changes into a single commit please?

@jakevdp I have squashed changes into a single commit now.

@copybara-service copybara-service bot merged commit 73de2a7 into jax-ml:main Dec 30, 2021
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@jakevdp jakevdp jakevdp approved these changes

Assignees

@jakevdp jakevdp

Labels

cla: yes pull ready Ready for copybara import and testing

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

jax.random.poisson does not broadcast to shape of rate param

3 participants

@rakki-18 @jakevdp @kokoro-team