No module named 'jax.linear_util' - Githubissues

keyonvafa / tbip

Text-Based Ideal Points

MIT License

43 stars 15 forks source link

No module named 'jax.linear_util' #17

Open paride92 opened 3 months ago

paride92 commented 3 months ago

Dear all,

I am trying to run your script in google colab https://colab.research.google.com/github/pyro-ppl/numpyro/blob/64fcc1ca64dae20b2f5624fb563fc18ab927b9ae/notebooks/source/tbip.ipynb

However, when I reach the "Perform Inference" section, I receive this error: ModuleNotFoundError: No module named 'jax.linear_util'

Do you have any suggestions on how to fix the issue? I tried to change the version of numpyro and jax, but I could not solve the issue.

Best regards