ML Katas

by: newest upvotes saves

Implement a Basic Optimizer with Optax

yesterday medium (<30 mins) | optimization jax optax sgd

Use Optax, JAX's optimization library, to create a simple Stochastic Gradient Descent (SGD) optimizer. You'll need to define a model, a loss function, and then use `optax.sgd` to update the...

1