Challenge: Optax

Time to switch back to the implementation side of things. As we mentioned earlier, implementing an optimizer in Optax is pretty straightforward. Let’s elaborate on those points in this guided challenge.

Prerequisites

We’ll use the linear regression example from the earlier challenge, but we’ll make it simpler by reducing it to a linear function ($b$ will be $0$).

Note: Usually, we implement neural networks by treating bias as a part of the weight vector, $w$ itself. (It’s prepended as $w_0$)