ML Katas

by: newest upvotes saves

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

1