[wealth_dynamics] Update lecture to use JAX instead of numba

[mmcky]

Please update lecture/wealth_dynamics lecture and convert the code to make use of the JAX python package. This includes looking at code that has high complexity or long run times, and areas where numba is being used to accelerate code.

Our style guide has a number of guidelines we should follow when making these updates

1. JAX code conversions
2. JAX admonitions and sequence generation

Please only use JAX where it would be beneficial to do, solving a problem well suited to JAX.