[ge_arrow] Update to JAX and compare runtime

[xuanguang-li]

Updated the ge_arrow.md to JAX and complemented the styling consistent with the operation manual.

Key changes:

• Rewrite the RecurCompetitive class as a NamedTuple.
• Complete all computations inside the compute_rc_model function. Inside this function, arguments of sub-functions can be written in the same way as the definitions in the theory part.
• Partially jitted the main computation function, and used jax.lax.fori_loop to conduct loops.
• Fixed some typos and styling.

Update: Runtime Comparison Between JAX (GPU), JAX (CPU), and NumPy

Methodology: nearly the same as in #654

• The JAX version uses the code in this PR, while the NumPy version uses the code in main.
• The runtime for JAX (GPU) is measured using Google Colab T4 GPU runtime.
• Runtime is collected using qe.timeit over 1,000 iterations.
• Each iteration consists of solving Example 3 (solving the wealth distribution for 100 kinds of transition matrices).

Results:

• Average runtime: JAX (CPU) > NumPy > JAX (GPU).
• More details are shown in the attached plots.

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-717--sunny-cactus-210e3e.netlify.app (b58cafe)

📚 Changed Lecture Pages: aiyagari, cake_eating_numerical, career, coleman_policy_iter, egm_policy_iter, ge_arrow, harrison_kreps, ifp, ifp_advanced, jv, lake_model, lqcontrol, mccall_correlated, mccall_fitted_vfi, mccall_model, mccall_model_with_separation, mccall_q, odu, optgrowth_fast, two_auctions, wald_friedman_2

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-717--sunny-cactus-210e3e.netlify.app (bf55c41)

📚 Changed Lecture Pages: aiyagari, cake_eating_numerical, career, coleman_policy_iter, egm_policy_iter, ge_arrow, harrison_kreps, ifp, ifp_advanced, jv, lake_model, lqcontrol, mccall_correlated, mccall_fitted_vfi, mccall_model, mccall_model_with_separation, mccall_q, odu, optgrowth_fast, two_auctions, wald_friedman_2

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-717--sunny-cactus-210e3e.netlify.app (253dea6)

📚 Changed Lecture Pages: aiyagari, cake_eating_numerical, career, coleman_policy_iter, egm_policy_iter, ge_arrow, harrison_kreps, ifp, ifp_advanced, jv, lake_model, lqcontrol, mccall_correlated, mccall_fitted_vfi, mccall_model, mccall_model_with_separation, mccall_q, odu, optgrowth_fast, two_auctions, wald_friedman_2

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-717--sunny-cactus-210e3e.netlify.app (be07276)

📚 Changed Lecture Pages: aiyagari, cake_eating_numerical, career, coleman_policy_iter, egm_policy_iter, ge_arrow, harrison_kreps, ifp, ifp_advanced, jv, lake_model, lqcontrol, mccall_correlated, mccall_fitted_vfi, mccall_model, mccall_model_with_separation, mccall_q, odu, optgrowth_fast, two_auctions, wald_friedman_2

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-717--sunny-cactus-210e3e.netlify.app

Commit: 12bace3

📚 Changed Lectures

• aiyagari
• cake_eating_numerical
• career
• coleman_policy_iter
• egm_policy_iter
• ge_arrow
• harrison_kreps
• ifp
• ifp_advanced
• jv
• lake_model
• lqcontrol
• mccall_correlated
• mccall_fitted_vfi
• mccall_model
• mccall_model_with_separation
• mccall_q
• odu
• optgrowth_fast
• two_auctions
• wald_friedman_2

---

Build Info

• Workflow: Build Project [using jupyter-book]