Maximum Likelihood Estimation in Jax

[djlacombe]

Hello!

I wanted to make the community aware of a Jupyter notebook I uploaded to GitHub that illustrates maximum likelihood estimation of the normal linear model in Jax.

The GitHub repo is available here: Maximum Likelihood in Jax

I just started using Jax, so there may be some errors or more efficient ways of coding all of this but I thought this might be interesting for those that use Jax.

Any comments or corrections would be appreciated.

Thanks.

[jakevdp]

Very cool, thanks for sharing! Related to this, one thing you may wish to check out is the recently released jaxopt package, which some members of the JAX team have been helping with. It has some interesting enhancements over the standard scipy-inspired minimization API.

[djlacombe]

Thank you in return! Again, I'm a novice Jaxer, but I'm very interested in these issues. I'd also like the notebook to be widely shared and perhaps incorporated into the docs as an example.