Closed colbybanbury closed 10 months ago
I did some basic tests with the v0.1.4 optax release and confirmed it works, but someone with more vision on the whole repo would know the preferred fix.
Very late answer, but we'll push some larger updates to the codebase within the next 1-2 months, that will also fix this.
Should be fixed at head now.
A recent change in Optax is breaking due to jax.Array Migration
Optax commit
Error:
I believe the fix would either be to pin optax to a release. e.g.
git+https://github.com/deepmind/optax.git@v0.1.4
hereor increase the minimum Jax version e.g.
pip install "jax[tpu]>=0.4.1"
here