Oh I see.
I don’t think there’s an issue with JAX, it’s just that we don’t support version 0.5.1 yet. We do have plans for adding support for newer versions of JAX though. This will take a few months so in the meantime I’d recommend using a lower version of JAX. For the demo and code you shared v0.4.33 should work. Let me know if this works for you Johanna!