All Katas - ML Katas

Custom Gradient with `jax.custom_vjp`

yesterday hard (<30 mins) | autograd jax custom-gradient vjp

Implement a function with a custom gradient using `jax.custom_vjp`. This is useful for numerical stability or for defining gradients for non-differentiable operations. A good example is a function...

Implement a Convolutional Layer

yesterday hard (<30 mins) | cnn jax convolution deep-learning

Implement a 2D convolutional layer from scratch in JAX. This will involve using `jax.lax.conv_general_dilated`. You will need to manage the kernel initialization and the forward pass logic....

Data Parallelism with `pmap`

yesterday hard (<30 mins) | jax pmap parallelism distributed-training

Parallelize a training step across multiple devices (e.g., multiple CPU cores if you don't have GPUs/TPUs) using `jax.pmap`. This is a fundamental technique for scaling up training. **Task:** Take...

Combine `vmap` and `pmap`

yesterday hard (<30 mins) | jax vmap pmap parallelism

For more complex parallelism patterns, you can combine `vmap` and `pmap`. For instance, you can use `pmap` for data parallelism across devices, and `vmap` for model ensembling on each device....