Convert wealth dynamics lecture to JAX - main conversion complete

Co-authored-by: mmcky <[email protected]>

Author: Copilot