[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 <mmcky@users.noreply.github.com>
Co-authored-by: John Stachurski <john.stachurski@gmail.com>

Author: 3 people