Closed adam-hartshorne closed 2 months ago
This is unexpected (the tests pass with jax-version 0.4.28). Can you post a pip freeze
?
Sorry, my bad, I had forgotten I had changed Pycharm to use a virtual environment to an earlier version of JAX to test something.
I can confirm it works fine on v0.4.28
Thanks for checking and clarifying! :)
Using JAX / JaxLib v0.4.28 and upgrading to v0.3.4 results in the following error (which doesn't exist when using v0.3.3