Challenge: Optax

Let's test our understanding of Optax with a coding challenge.

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.


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

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

Get hands-on with 1000+ tech skills courses.