Adding a discussion of JAX

[jstac]

Hi @mmcky , I've ported this over from the JAX GPU lecture series and tweaked it to change the focus to using JAX on CPUs.

Could you please read it carefully and make any necessary modifications -- or delegate this task? It's probably necessary to look at the text relating to timing of code execution and check that it still makes sense.

[mmcky]

thanks @jstac

I will need to update the repo infrastructure to enable this series to access hardware with a gpu.

Tasks

• @mmcky debug missing software and target jax running on cpu
• @mmcky read the lecture carefully in the context of this lecture series.

[mmcky]

@jstac just read the cpu target.

I will just have to update the software and debug. Thanks.

[jstac]

Thanks Matt. As you say, this is just for CPU.

JAX is already faster on the CPU so we can unify many of our lectures around this platform.

[github-actions]

🚀 Deployed on https://688c10959c083e5ca9af0a1c--epic-agnesi-957267.netlify.app

[mmcky]

@jstac this is now running well.

I have just completed a review. Nice, quick introduction to JAX. The difference in generating random numbers is pretty interesting and a minor inconvenience when using jax (having to use keys) but would assist with reproducibility with the lectures :-)

There was only one issue re: timing. There are a couple of instances where you write about removing compile time and the total run time is higher (due to higher sys time). The actual user time is lower but this is confusing for new programmers.

and

Perhaps we should remove the second part and just focus on overall execution time?

[jstac]

Thanks @mmcky , nice catch. Please go ahead with your suggestion.