Numpyro: NUTS with GPU (doesn't work for bnn.py)

[ziatdinovmax]

Hello,

I've noticed that choosing a GPU option results in memory problems when "training" even a very simple Bayesian neural network with NUTS (to be more specific it crashes with a "segmentation fault" message). I first thought I'm doing something wrong but when I ran the original bnn.py script from NumPyro examples it also crashed. Interestingly, the gp.py example works fine with a GPU option. I was running these scripts using Colab VM and GCP VM with NVIDIA's P100 GPU. Here's a link to Colab notebook:
https://colab.research.google.com/drive/1TlRo8mr7JTM6fUF6IUIe01yt8U8_Z73Q

Thanks for your help in advance!

[fehiepsi]

Thanks! I can confirm the issue. In my system, memory keeps increasing until the crash... Looking like this is jax/jaxlib issue. I'll isolate and report to them.

[fehiepsi]

@ziatdinovmax Thanks for the report! Now you can upgrade jaxlib to fix the issue.