Linear Regression is one of the most fundamental algorithms in machine learning. In this post, we’ll implement it from scratch using JAX — Google’s high-performance numerical computing library that provides automatic differentiation and GPU/TPU acceleration.