ML Katas

Now that you have all of the basic building blocks, it's time to build a simple multi-layer perceptron (MLP). In this exercise, you are to build a 2-layer MLP for a regression problem.

1. Implement the MLP. It should take as input a vector x and a list of weight matrices and bias vectors. The first layer should be an affine transformation followed by a ReLU activation and the second layer should be another affine transformation.
2. Implement a loss function that computes the mean-squared error.
3. Implement a jit'ed training step, as in the previous exercise.
4. Implement a training loop that runs for 1000 steps. Plot the loss as a function of the training step. It should be decreasing.