AttributeError: module 'jax._src.random' has no attribute 'KeyArray'

[dominikstrb]

After pip-installing dynamax into a fresh virtualenv with Python 3.12.7, I get an AttributeError: module 'jax._src.random' has no attribute 'KeyArray' when importing LinearGaussianSSM from dynamax.linear_gaussian_ssm.

The reason seems to be that dynamax.types defines PRNGKey = prng.KeyArray and jax.random has removed KeyArray in version 0.4.24.

I have not been able to find a constellation of jax versions that works because version 0.4.23 does not seem to be compatible with scipy>=1.6.0 because 'scipy.linalg' has no attribute 'tril'.

Proposed solution:
Either the requirements for dynamax need to define concrete versions of jax, scipy (and possibly other dependencies) that work with the current codebase, or the type PRNGKey needs to be redefined in dynamax.types (although I cannot judge to what extent what part of dynamax relies on that type.

[dominikstrb]

I was able to get it working by installing the following versions:

• scipy <=1.12.0,>= 1.6.0
• jax==0.4.23
• jaxlib==0.4.23
• optax==0.2.2
• chex==0.1.86

[gileshd]

Hi @dominikstrb!

Apologies for this problem. There is actually an open PR, #379, which updates the PRNGKey code to work with the JAX updates - we are hoping to get this merged soon.

In the meantime installing directly from that branch should do the trick and avoid having to pin jax==0.4.23.