[markov_asset] Update to JAX#654

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

update_markov_assetQuantEcon/lecture-python.myst:update_markov_assetCopy head branch name to clipboard