ModuleNotFoundError: No module named 'blackjax.optimizers'

[xaviergonzalez]

Bug Description

I am trying to use blackjax in colab. Yesterday, it was working fine. Now, when I try to import blackjax, I get the following error: ModuleNotFoundError: No module named 'blackjax.optimizers'

Steps/Code to Reproduce

! pip install blackjax
import blackjax

Expected Results

No error messages, and blackjax loads.

Actual Results

ModuleNotFoundError: No module named 'blackjax.optimizers'

Versions

Python 3.7.13 (default, Apr 24 2022, 01:04:09)
[GCC 7.5.0]
Jax 0.3.8
Jaxlib 0.3.7

[karm-patel]

Same error!

Versions:
Python 3.7.13
blackjax 0.8.1
jax 0.3.13
jaxlib 0.3.10

[rlouf]

Can you try with blackjax 0.8.0 ? @junpenglao

[xaviergonzalez]

Can you try with blackjax 0.8.0 ? @junpenglao

I'm getting the same error with the following code

! pip install blackjax==0.8.0
import blackjax

[rlouf]

This should be solved in 0.8.2 !

[murphyk]

It does work if you first install jaxopt, like this:

But if you install blackjax before jaxopt it will fail

[rlouf]

This is really strange, I'll try to reproduce it locally. I'll also include jax and jaxopt in the dependencies in the next release.