[kalman_2] Update to JAX#644

xuanguang-li wants to merge 7 commits intomainQuantEcon/lecture-python.myst:mainfrom

update_kalman2QuantEcon/lecture-python.myst:update_kalman2Copy head branch name to clipboard