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

JAX: Implement geometric RV sampling #1444

Merged
merged 2 commits into from
Feb 21, 2023
Merged

Conversation

Smit-create
Copy link
Member

Fixes #1323
Here are a few important guidelines and requirements to check before your PR can be merged:

  • There is an informative high-level description of the changes.
  • The description and/or commit message(s) references the relevant GitHub issue(s).
  • pre-commit is installed and set up.
  • The commit messages follow these guidelines.
  • The commits correspond to relevant logical changes, and there are no commits that fix changes introduced by other commits in the same branch/BR.
  • There are tests covering the changes introduced in the PR.

Comment on lines +359 to +404
sample_num = jax.numpy.log(jax.random.uniform(sampling_key, size))
sample = sample_num / jax.numpy.log1p(-p)
sample_ceil = jax.numpy.ceil(sample)
Copy link
Member Author

Choose a reason for hiding this comment

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

@codecov
Copy link

codecov bot commented Feb 20, 2023

Codecov Report

Merging #1444 (552346f) into main (4c74650) will increase coverage by 0.00%.
The diff coverage is 100.00%.

Additional details and impacted files

Impacted file tree graph

@@           Coverage Diff           @@
##             main    #1444   +/-   ##
=======================================
  Coverage   74.81%   74.81%           
=======================================
  Files         194      194           
  Lines       49995    50007   +12     
  Branches    10572    10572           
=======================================
+ Hits        37402    37414   +12     
  Misses      10267    10267           
  Partials     2326     2326           
Impacted Files Coverage Δ
aesara/link/jax/dispatch/random.py 100.00% <100.00%> (ø)

rlouf
rlouf previously approved these changes Feb 20, 2023
brandonwillard
brandonwillard previously approved these changes Feb 20, 2023
Copy link
Member

@brandonwillard brandonwillard 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, but appears to need some conflicts resolved.

@brandonwillard brandonwillard added enhancement New feature or request JAX Involves JAX transpilation random variables Involves random variables and/or sampling labels Feb 21, 2023
@Smit-create Smit-create marked this pull request as ready for review February 21, 2023 05:05
@rlouf rlouf merged commit 311c901 into aesara-devs:main Feb 21, 2023
@rlouf
Copy link
Member

rlouf commented Feb 21, 2023

Thank you @Smit-create!

@Smit-create Smit-create deleted the i-1323 branch February 21, 2023 06:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request JAX Involves JAX transpilation random variables Involves random variables and/or sampling
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Add JAX implementation for GeometricRV
3 participants