PRNKey error from dynamax

[rklees]

I installed sts-jax using the latest development branch. When running the sts_electric_demo notebook, I get an error message

ImportError: cannot import name 'PRNGKey' from 'dynamax.types' (/Users/rolandklees/miniforge3/envs/stsjax/lib/python3.11/site-packages/dynamax/types.py)

I also suffer from the same error message when running my own code. Is there anyone who knows how to fix this?

[PaoloRanzi81]

Hi, I have stumbled upon the same issue. I think is the newest version of JAX is causing it (in particular the Random Number Generator). After installing everything in a sort of "recovery mode" like here, I did some trial-and-error.

I did manually change 3 scripts (learning.py, sts_model.py and sts_ssm.py) in "sts-jax" using more or less the following pattern:

from dynamax.types import Scalar
# from dynamax.types import PRNGKey, Scalar # old version (not working anymore)

key = jr.key(0),
# key: PRNGKey = jr.PRNGKey(0), # old version (not working anymore)

I hope it helps!