ML Katas

In this exercise, you will implement a full training step for the regression problem that you have been working on.

1. Instantiate your model parameters, W and b, and your data x and y_true. Remember to use a PRNG key.
2. Compute the gradients of your loss function with respect to your model parameters.
3. Update your parameters using gradient descent. That is, $W_{new}=W_{old}-\alpha {\nabla }_{W}L$ and $b_{new}=b_{old}-\alpha {\nabla }_{b}L$ for some learning rate $\alpha$. You can set $\alpha =0.01$.
4. Wrap this all in a function that you jit compile. How much of a speedup do you see?