Intro notebook fails on colab
Closed this issue · 5 comments
I am running docs/notebooks/jax-cosmo-intro on colab with a CPU-only runtime and the following line:
F = - hessian_loglik(params)
gives a mysterious error:
TypeError: Integers cannot be raised to negative powers, got integer_pow(ShapedArray(int32[513,1]), -2)
Any idea why this is being traced with an integer array here?
I just checked that I get the same error with a GPU runtime.
Yes, thank you! I spent some time to track down the issue, apparently it was something broken in jax, checkout this issue:
jax-ml/jax#3861
It should work again with the latest version of JAX.... this happens sometime... but the JAX devs are super responsive and usually very quickly find fixes
Nice sleuthing (and great to see the jax devs so responsive) - thanks!
I've just pushed updated the colab notebook, by adding a jax update at the beginning and printing out versions of both jax and jax-cosmo. This should help next time around ^^'
@all-contributors please add @dkirkby for bug
I've put up a pull request to add @dkirkby! 🎉