ML Katas

by: newest upvotes saves

Vectorizing with `vmap`

this year easy (<30 mins) | vmap batching

Another of JAX's powerful features is its ability to automatically vectorize functions with `vmap`. For example, `vmap` can be used to automatically batch a function that was written for a single...

1