Andrew Elkadi

The version of chex being pointed to is deprecated. Running `!pip install --upgrade jaxlib` `!pip install --upgrade chex` should fix these issues. Hope this helps!

Hello, Running the following in a cell should fix this issue: `!pip install --upgrade jaxlib` `!pip install --upgrade chex` It appears the version of chex being pointed to is deprecated......