rlouf / mcx

Express & compile probabilistic programs for performant inference on CPU & GPU. Powered by JAX.
https://rlouf.github.io/mcx
Apache License 2.0
325 stars 17 forks source link

Add Adam and RMSProp optimizers #15

Closed rlouf closed 4 years ago

rlouf commented 4 years ago

Optimizers may be needed during warmup, but it is also fun to show they can be written as kernels. This PR adds Adam and RMSProp.

rlouf commented 4 years ago

Good job me :+1: merging.