Problems setting the environment about s-fsvi HOT 3 CLOSED

Thanks for catching this. I added chex==0.0.7 to the environment.yml file.

from s-fsvi.

Hi,

I don't have a solution for you, but the error looks like it's likely a JAX issue. To check whether this is a JAX installation issue, try running

import jax
from jax.lib import xla_bridge
print(f"Device: {xla_bridge.get_backend().platform}")

seed=0
main_rng = jaxrandom.PRNGKey(seed)

If this runs without errors and shows that you're running JAX on GPU, your JAX setup is likely working correctly. If you're getting an error, I would recommend checking if you're using the jax, jaxlib, optax, and chex versions specified in the environment.yml file.

from s-fsvi.

Thanks for your quick answer! I managed to solve the problem. It wasn't a jax problem, but when I checked the other libraries I noticed that chex is not listed in the environment.yml, this leads to install the wrong version of chex.
I tried with chex==0.0.7, which is the version I found in the colab_requirements.txt, and this fixes the problem.

from s-fsvi.