0

I have a problem that I do not know how to tackle it. I have a nb on Colab that was working perfectly on CPU Here is the working nb 19th Jan 23

Then, I was looking to get a new nb working in the same condition (CPU)

  1. Dealing with JAX 0.3.25 version I have pip install
!pip install jax==0.3.25
!pip install jaxlib==0.3.25

then also pip install my repo at it was Jan 203, using the git commit "tag"

!pip install -q git+https://<git-repo>@a3ecb7f

Now proceeding the same as the Jan nb, crashes as you can see HERE nb 22 March 23

My repo is a fork of an other repo done Nov. 2022.

Is there a reason why the March nb may crash due to some hidden changes? I'm not so confortable with git so may I was doing something wrong on Calob.

Help is warmly welcome.

Jean-Eric
  • 372
  • 2
  • 14
  • The second colab link (showing the error) is not public. Speculating: a number of things have changed since january, including updated JAX versions, updated Python version in the Colab runtime, etc. – jakevdp Mar 22 '23 at 19:44

1 Answers1

0

I'm not sure exactly what the bug is, but I happen to know that in January the version of JAX on Colab was 0.3.25. I found that installing this version makes your notebook run without error:

!pip install jax==0.3.25 jaxlib==0.3.25

It's hard to say exactly what JAX change may have led to the error you're seeing. JAX's CHANGELOG summarizes the changes in recent versions.

jakevdp
  • 77,104
  • 11
  • 125
  • 160
  • Hi @jakevdp Do you manage to run my nb w/o error! After your edit, I just do exactly as you that is to say make the pip install on a single line, while I was doing it in two separate lines. Does it make a difference ??? – Jean-Eric Mar 23 '23 at 09:31
  • After my edit, I just do exactly as you and it works with 0.3.25. I have investigate all the JAX releases and this is the 0.4.6 that fails ( I have opened the nb privilage) – Jean-Eric Mar 23 '23 at 10:18
  • I have not manage to get an install of the 0.4.5 to see if it is ok or not – Jean-Eric Mar 23 '23 at 10:21