Thanks @HumphreyYang and thanks @shizejin for reporting!

(@HumphreyYang always puts in great PRs, so the whole issue is clear at a glance. This is super valuable to me.)

I was trying to avoid timeit because people who don't understand ipython get the impression that the code is slow.

But it seems this is causing problems.

@HumphreyYang , let's do as you recommend. Would you mind to put the fix in?

Many thanks @jstac for the suggestion! I will submit a FIX for this.

I was trying to avoid timeit because people who don't understand ipython get the impression that the code is slow.

I think this is a very valid concern as well. The reader sees that the number is small but notices that the code takes longer to run.

I will add a brief note below the first use of %timeit, explaining that it runs the line multiple times to provide a robust comparison.

Given this concern, should we change %timeit globally or only in the introductory lecture?

I think the issue is less noticeable when the runtime differences are large.

I agree @HumphreyYang . Let's change it only in the introductory lecture and with the note that you mention. Thanks.

Fixed by using JAX==0.6.0