[lake_model] Update lecture to JAX and check compliance to style sheet (#589)

* updates to JAX

* Add back quantecon install

* updates

* updates

* updates

* updates

* update jax install

* update static_argnames

* minor update

* update according to feedback

* updates according to feedback

* misc

* updates according to feedback

---------

Co-authored-by: Matt McKay <[email protected]>
Co-authored-by: John Stachurski <[email protected]>

Author: 3 people