All Katas - ML Katas

Implement a Simple Linear Regression in JAX

yesterday easy (<30 mins) | autograd optimization jax jit linear-regression

Your task is to implement a simple linear regression model from scratch using JAX. You'll need to define the model, a loss function (like Mean Squared Error), and a gradient descent update rule....

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...

Manual Gradient Descent Step

this month medium (<30 mins) | autograd optimization gradient descent

Simulate one step of gradient descent for a simple quadratic loss. ### Problem Given a scalar parameter $w$ initialized at 5.0, minimize the loss $L(w) = (w - 3)^2$ using PyTorch. - **Input:**...