Further reading
- Linear models in machine learning: https://blog.dataiku.com/top-machine-learning-algorithms-how-they-work-in-plain-english-1
- The JAX ecosystem: https://moocaholic.medium.com/jax-a13e83f49897
- A brief tutorial that further explores JAX: https://colinraffel.com/blog/you-don-t-know-jax.html