Transform result of JAX sampling

[martiningram]

Hi all,

I'm assuming this is known but since no issue seems to have been raised about it (let me know if I just missed it, sorry!), I wanted to raise one.

When sampling using JAX, it appears that the result is given on the unconstrained representation of the variables. For example, in the Radon example, the parameter sigma_a yields samples called sigma_a_log__, which are in fact samples from log(sigma_a). This is easy enough to change as a user in this case (just transform with np.exp e.g.), but it gets a bit more fiddly when sampling e.g. covariance matrices, where two transformations have to be composed (exp for the diagonal elements followed by L L^T to get the covariance). I'm assuming that pymc3's original NUTS sampler required these transformations also so hopefully the same logic can be used here, too.

Best,
Martin

[twiecki]

Hi Martin, thanks - that'd be definitely be useful. Is that something you'd be interested in working on? We're still looking for help on the JAX backend.

[martiningram]

I'm happy to give it a go too @twiecki ! @kc611 not sure if we could work together on this somehow? I'm planning to take a first look today to see if I can work out the original logic, can keep you posted.

[martiningram]

Hi all, I believe the pull request addresses this issue. Sorry @kc611 , I just gave it a shot, but if you like, if you see anything that could be improved please let me know!

[martiningram]

Fixed by #4427

[PedroSebe]

Is the fix on #4427 exclusive to the Numpyro sampler? I used the JAX-TFP sampler today and I got neither my transformed variables nor my deterministic variables. If that is the case, we probably should reopen this issue.

(My install is pymc3 3.11.2, jax 0.2.11, tensorflow 2.5.0, tensorflow-probability 0.13.0. running on Google Colab with GPU)

[twiecki]

Any reason you don't use the numpyro sampler? I thought we might want to remove the TFP one.

…

On Wed, Jul 14, 2021, 07:02 Pedro Sebe ***@***.***> wrote: Is the fix on #4427 <#4427> exclusive to the Numpyro sampler? I used the JAX-TFP sampler today and I got neither my transformed variables nor my deterministic variables. If that is the case, we probably should reopen this issue. (My install is pymc3 3.11.2, jax 0.2.11, tensorflow 2.5.0, tensorflow-probability 0.13.0. running on Google Colab with GPU) — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#4415 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAFETGBGWDKBMP4TSAMB2VLTXUK5ZANCNFSM4V44RMZQ> .

[PedroSebe]

I was running on a Google Colab instance, and using Numpyro caused the problem reported at #4645. I was not aware that TFP sampler is going to be removed, will it happen on the next release?

[twiecki]

This is fixed in a recent aesara version. Can you try installing aesara and pymc3 main?