add optax optimizer, and make examples directory

[murphyk]

• I generalized build_adam_optimizer to make build_optax_optimizer that takes in any optax optimizer.
• I created examples subdirectory
• I added examples/gauss_sample_sgld.py which is basically the same as your 'quick start' demo (sample from D-dim. Gaussian)
• I added examples/gauss_optimize.py which is similar to gauss_sample_sgld but uses the adam optimizer
• I was not able to make gauss_optimize.py test my new optax_optimizer because I get the attempted relative import with no known parent package error :(

[murphyk]

Related code: https://github.com/probml/pyprobml/blob/master/scripts/sgmcmc_utils.py https://github.com/probml/pyprobml/blob/master/scripts/sgmcmc_nuts_demo.py

[jeremiecoullon]

Thanks! I fixed the relative import error and also added those 2 extra scripts in the example directory