ML Katas

Your task is to implement a simple linear regression model from scratch using JAX. You'll need to define the model, a loss function (like Mean Squared Error), and a gradient descent update rule. This exercise will help you get comfortable with JAX's functional style, its automatic differentiation capabilities with jax.grad, and the use of jax.jit for performance optimization.

Model: $y=wx+b$

Loss (MSE): $L(w,b)=\frac{1}{N}{\sum }_{i=1}^N(y_i-(w{x}_i+b){)}^2$

Verification: - Your loss should decrease over training iterations. - The final learned parameters w and b should be close to the true values used to generate the data.