Improve and simplify JAX lectures#532

jstac merged 1 commit intomainQuantEcon/lecture-python-programming:mainfrom

jxmodsQuantEcon/lecture-python-programming:jxmodsCopy head branch name to clipboard