Jax error

[TheOneTrueGuy]

I get this on several different branches such as camenduru as well as the original
AttributeError: module 'jax.random' has no attribute 'KeyArray'

have tried various updates, upgrades and using an earlier version via:
pip install "jax[cuda12_local]==0.4.23" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
none of which has worked. This used to work great but now the stable-diffusion 1.5 model is lost and I will have to find a new one to arrange the download from. Will try SDXL too but it looks like it is super touchy.

Thanks to any and all for any kind of help. Good luck in your own efforts too.