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

Add ChiSquareRV JAX implementation #1363

Merged
merged 2 commits into from
Mar 10, 2023

Conversation

larryshamalama
Copy link
Contributor

@larryshamalama larryshamalama commented Dec 12, 2022

Closes #1322

Given the shape/rate parametrization of GammaRV in Aesara, this addition uses the fact that $$\chi^2_{\text{df}} \equiv \Gamma\left(\frac{\text{df}}{2}, \frac{1}{2} \right).$$

Wiki reference here

@brandonwillard brandonwillard added enhancement New feature or request JAX Involves JAX transpilation random variables Involves random variables and/or sampling labels Dec 12, 2022
@rlouf
Copy link
Member

rlouf commented Dec 13, 2022

Looks good, waiting for #1345 to be merged to update the key splitting mechanism.

rlouf
rlouf previously approved these changes Feb 21, 2023
@rlouf
Copy link
Member

rlouf commented Feb 21, 2023

Thanks for your contribution @larryshamalama!

@rlouf
Copy link
Member

rlouf commented Feb 21, 2023

I am unable to reproduce the test failure locally (it even works with smaller a smaller number of samples). We need to understand what causes this.

@rlouf rlouf force-pushed the chisq-jax branch 2 times, most recently from a4356c1 to 27f09c4 Compare February 23, 2023 14:40
@larryshamalama
Copy link
Contributor Author

I can have a look later, but I'm currently on vacation until early March so I will only be able to do so then. Admittedly, I have not looked much into this PR after its inception...

@rlouf
Copy link
Member

rlouf commented Feb 24, 2023

I don't think the failures have anything to do with your implementation; the tests pass locally for me.

@brandonwillard brandonwillard merged commit 984ee55 into aesara-devs:main Mar 10, 2023
@brandonwillard
Copy link
Member

Thanks, @larryshamalama!

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 ChiSquareRV
3 participants