Add JAX implementation for InvGammaRV

[rlouf]

No description provided.

[manish-p-gupta]

Hey, I'm new here but I have some idea of working with JAX. Is this up for contributions?

[rlouf]

By all means!

[manish-p-gupta]

Great! Can you point me to any starting documentation that'll be helpful. I am aware about the contributing guidelines. Let me know if you think I should do anything else. I'll open a draft PR to discuss things in detail there too!

[rlouf]

Here's an explanation I gave in a related issue. This should be enough to get you started: #1335 (comment)

[manish-p-gupta]

Okay thanks! I'll go over it and start the work.

[PaulScemama]

@rlouf I've been taking a look at this and may have come across a small notational choice that may cause confusion with users.

In aesara/tensor/random/basic.py, there exists an Inverse Gamma random variable. In it's documentation, it says "the probability density function for invgamma in terms of its shape parameter $\alpha$ and $\text{scale}$ parameter $\beta$ is:

$$ f(x; \alpha, \beta) = \frac{\beta^\alpha}{\Gamma(\alpha)} x^{-(\alpha+1)} \exp\left(-\frac{\beta}{x}\right) $$

".

This, however, creates a confusing discrepancy between Aesara's documentation of the Gamma random variable.In it's documentation, says "the probability density function for gamma in terms of the shape parameter $\alpha$ and $\text{rate}$ parameter $\beta$ is:*

$$ f(x; \alpha, \beta) = \frac{\beta^\alpha}{\Gamma(\alpha)}x^{\alpha-1}e^{-\beta x} $$

".

I think the InvGamma documentation should follow the Gamma, and the rate parameter ($\beta$) should be passed in, and then transformed to the scale by its reciprocal.

BTW: @manish-p-gupta are you still working on this? If not, I don't mind taking a stab at it. Let me know what you both think! Thanks.

[rlouf]

I agree with the fact that it is confusing; I actually made a mistake because of this when using Aesara in a project. I believe that we were blindly following SciPy when implementing this. But now the issue is that downstream code may be relying on this API so we have to be cautious.

You can get a stab at the JAX implementation if you'd like!

[PaulScemama]

In terms of the parameter name thing: Ah I see. I can take a peek and see how much downstream code is reliant on it. I agree that being cautious is the best plan.

Also sounds good, thanks! I have to get public release permission (from work) if I were to literally contribute so it may take a week or so.

[PaulScemama]

Hi @rlouf @brandonwillard just checking in here.

To review: the InvGammaRV in aesara/tensor/random/basic.py has some confusing variable naming and documentation that doesn't align with GammaRV in the same file. You mentioned we should be cautious in fixing this in case downstream code depends on it.

Here's what I've found regarding the downstream code dependence on InvGammaRV:

• searching InvGammaRV in this repository yields this
  • we see that in aesara/tensor/random/basic.py, we instantiate InvGammaRV with invgamma.
• searching invgamma yields this
  • we see that in tests/tensor/random/test_basic.py we test invgamma.

In conclusion, the files affected by a change in InvGammaRV are...

1. the file it is defined in aesara/tensor/random/basic.py
2. the file testing it tests/tensor/random/test_basic.py

Note, invgamma shows up in doc/reference/random/index.rst, but I don't think changing anything about InvGammaRV will change this?

Let me know what you think.

[brandonwillard]

Let me know what you think.

I think the bigger issue is user-written code and other packages that depend on the current parameterization.

[PaulScemama]

@brandonwillard makes sense. So should I go ahead then and give the JAX implementation a go with how our Inverse-Gamma currently stands?

And I totally forgot about the other "children" packages of Aesara 😅, that is definitely not just 2 files then.

[brandonwillard]

So should I go ahead then and give the JAX implementation a go with how our Inverse-Gamma currently stands?

Yeah, feel free to create a PR.