ML Katas

by: newest upvotes saves

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

1