Port Numba-accelerated code over to Google JAX

[jstac]

This is a meta issue for conversion of code to JAX. The target is lectures that are currently accelerated by Numba and could alternatively be accelerated by JAX.

One example is https://python.quantecon.org/wealth_dynamics.html

This could be shifted to JAX (as a model for this, see the solution to the last exercise in https://python-programming.quantecon.org/jax_intro.html)

@thomassargent30 suggests that, at least for now, we keep the Numba code and add the JAX code as a separate implementation, so that we can compare the speed. I suggest that we put the JAX code first in each case.

Also, for each lecture, we should only add the JAX code if it's actually faster.

As a workflow, I suggest that we

1. create a separate issue for each lecture to be converted
2. close that issue via PR
3. close this meta issue

@Smit-create and @HumphreyYang , I would like you to lead this project, if you are willing. I will supervise. We can discuss this via zoom before you get started, or during, if it's helpful. Just request by email and we'll set it up.

[Smit-create]

Thanks for opening this issue. Moving to JAX is a great option. We can meet and discuss the probable lectures that can be optimized using JAX and start porting them one by one.

[jstac]

The solution to the main exercise in https://python.quantecon.org/kesten_processes.html can also be ported to JAX in a similar way.

This notebook might be a useful guide for the simulation method: https://github.com/QuantEcon/cbc_workshops/blob/main/day_4/european_option_jax.ipynb

[jstac]

As mentioned by @HumphreyYang , we can check the backend in each lecture with

print(f"jax backend: {jax.devices()[0].platform}")

Are we happy with this convention?

A second point is that I think we need to force the use of 64 bit floats, since some of the calculations need a lot of precision to converge. Therefore I suggest we add

# Use 64 bit floats with JAX in order to match NumPy/Numba code
jax.config.update("jax_enable_x64", True)

in each lecture after we import JAX. CC @Smit-create and @mmcky regarding these conventions.

[HumphreyYang]

A second point is that I think we need to force the use of 64 bit floats, since some of the calculations need a lot of precision to converge. Therefore I suggest we add

# Use 64 bit floats with JAX in order to match NumPy/Numba code
jax.config.update("jax_enable_x64", True)

in each lecture after we import JAX. CC @Smit-create and @mmcky regarding these conventions.

Hi @jstac, coercing 64-bit computation is a good idea but in my case, it results in an exponential increase in memory use (around 40 GB) resulting in a RESOURCE_EXHAUSTED error and a very significant runtime increase even when memory is uncapped. I think other large simulations will have the same issue, and downscaling these examples may not be the optimal path. Thus, I think we should probably only update this config when it is necessary (i.e. in Newton's lecture where convergence is dependent on the precision).

[jstac]

Roger that @HumphreyYang , makes sense.

[jstac]

@mmcky , when you have time please note the above CC. I guess we'll do this on a case-by-case basis.