[newton_method] Update lecture code and fix typos (#564)

* update code

* fix typos

* update lecture to use JAX

* apply suggestions

* update code line length

* minor updates

---------

Co-authored-by: Matt McKay <mmcky@users.noreply.github.com>
Co-authored-by: Humphrey Yang <u6474961@anu.edu.au>

Author: 3 people