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

Build a Custom ReLU Activation Function

yesterday easy (<10 mins) | autograd jax activation-function custom-function

Implement the Rectified Linear Unit (ReLU) activation function in JAX. Then, use `jax.grad` to find its derivative. The ReLU function is defined as: $ReLU(x) = max(0, x)$ **Verification:** - For...

Vectorized Operations with vmap

yesterday easy (<10 mins) | jax vmap vectorization batching

You have a function that processes a single data point. Your goal is to use `jax.vmap` to apply this function to a whole batch of data without writing an explicit loop. For example, consider a...

Debugging JAX code with `jax.debug.print`

yesterday easy (<10 mins) | debugging jax jit

In JAX, standard Python `print` statements don't always work as expected within `jit`-compiled functions because they execute at trace time. [11] The solution is to use `jax.debug.print`. [11, 23]...