[mccall_fitted_vfi] JAX conversion

[longye-tian]

Hi @mmcky @HumphreyYang
This PR updates the Job Search III (Fitted Value Function Iteration) lecture to use JAX instead of NumPy/Numba, following modern QuantEcon standards and fixing style guide violations.

Reference to Issue #502

Changes

Style Fixes (per QuantEcon style guide)

• Fixed section heading capitalization (Title Rule):
  • ## The Algorithm → ## The algorithm
  • ### Value Function Iteration → ### Value function iteration
  • ### Fitted Value Function Iteration → ### Fitted value function iteration

JAX Conversion

• Replaced NumPy/Numba imports with JAX equivalents
• Converted @jitclass to NamedTuple with factory function pattern
• Updated random number generation to use JAX's PRNG keys
• Replaced imperative loops with jax.lax.while_loop
• Implemented functional programming patterns in exercises using jax.vmap
• Updated all array operations from NumPy to JAX

Text Consistency

• Updated documentation references from numpy.interp to jnp.interp
• Fixed code references: np.inf → jnp.inf, np.linspace → jnp.linspace

[github-actions]

🚀 Deployed on https://68a7f735988a4b6753cfd684--nostalgic-wright-5fa355.netlify.app

[HumphreyYang]

Many thanks @longye-tian, it looks really nice!

I just pushed some minor fix of typos.

Please see some very minor suggestions below (most are just cosmetic)!

[Copilot]

Pull Request Overview

This PR modernizes the Job Search III (Fitted Value Function Iteration) lecture by converting from NumPy/Numba to JAX, implementing functional programming patterns and fixing style guide violations. The conversion improves compatibility with modern QuantEcon standards while maintaining the lecture's educational content.

Key changes:

• Converts NumPy/Numba code to JAX with functional programming patterns
• Fixes section heading capitalization per QuantEcon style guide
• Updates random number generation to use JAX's PRNG system

[mmcky]

Thanks @longye-tian for an excellent PR, thanks @HumphreyYang for comments. I just had a few questions.

1. the figure in exercise 2 is slightly different. This might be OK -- but just thought I would check if this is due to a random seed, or if it has any implications for the exercise (LHS = this PR, RHS = Live site)

2. The timing is slower on JAX version (GitHub actions build - gpu) than the numba version. (LHS = this PR, RHS = Live site) (Status: Local testing shows jax should run approximately 2 x faster than numba version, so this timing is an issue with GA environment). #567 (comment)

• review collab comments to either agree or resolve.

[mmcky]

@HumphreyYang I see 95f5ef6 the result is 10s so 2s faster than gpu and is the same as the numba implementation.

[HumphreyYang]

@HumphreyYang I see 95f5ef6 the result is 10s so 2s faster than gpu and is the same as the numba implementation.

Hi @mmcky,

This is an interesting behaviour. I am curious could it be because of the jb build issue.

I will run some more benchmark.

[mmcky]

@HumphreyYang I see 95f5ef6 the result is 10s so 2s faster than gpu and is the same as the numba implementation.

Hi @mmcky,

This is an interesting behaviour. I am curious could it be because of the jb build issue.

I will run some more benchmark.

So I ran the main branch and this branch locally

state	time
main (numba)	6s
job_separation (jax)	3s

so it does appear to be a GitHub Actions thing.
Surprising as we are running this on an g4dn.2xlarge which has 8 vCPUs with 64Gb RAM.

[HumphreyYang]

Hi @mmcky,

Please find the gist result from my end:

https://gist.github.com/HumphreyYang/a488df0e8ef01b8d91af6e6309d5e45f

It should run around 2x faster than the numba version in both CPU and GPU.

There are while loops in the code so I expect GPU wouldn't help much.

[HumphreyYang]

@HumphreyYang I see 95f5ef6 the result is 10s so 2s faster than gpu and is the same as the numba implementation.

Hi @mmcky,
This is an interesting behaviour. I am curious could it be because of the jb build issue.
I will run some more benchmark.

So I ran the main branch and this branch locally

state time
main (numba) 6s
job_separation (jax) 3s
so it does appear to be a GitHub Actions thing. Surprising as we are running this on an g4dn.2xlarge which has 8 vCPUs with 64Gb RAM.

Great! Our number matches! It looks like 8vCPUs are not helping : )

[mmcky]

@longye-tian I have updated the timing comment above (with strikethrough) based on our investigations and opened #568

[HumphreyYang]

1. the figure in exercise 2 is slightly different. This might be OK -- but just thought I would check if this is due to a random seed, or if it has any implications for the exercise (LHS = this PR, RHS = Live site)

Thanks @mmcky for the detailed review:

JAX and NumPy use different pseudorandom number generators (PRNGs), so even with the same seed, they produce different random numbers. Below is an illustration:
https://gist.github.com/HumphreyYang/f8aa7b322c48723bac568290de8a08fc

[mmcky]

1. the figure in exercise 2 is slightly different. This might be OK -- but just thought I would check if this is due to a random seed, or if it has any implications for the exercise (LHS = this PR, RHS = Live site)

Thanks @mmcky for the detailed review:

JAX and NumPy use different pseudorandom number generators (PRNGs), so even with the same seed, they produce different random numbers. Below is an illustration: https://gist.github.com/HumphreyYang/f8aa7b322c48723bac568290de8a08fc

Thanks @HumphreyYang so once we change to jax that graph will become static again (with the new shape)? Also just to confirm, the new shape doesn't change the interpretation in any way re: context of the solution?

[HumphreyYang]

1. the figure in exercise 2 is slightly different. This might be OK -- but just thought I would check if this is due to a random seed, or if it has any implications for the exercise (LHS = this PR, RHS = Live site)

Thanks @mmcky for the detailed review:
JAX and NumPy use different pseudorandom number generators (PRNGs), so even with the same seed, they produce different random numbers. Below is an illustration: https://gist.github.com/HumphreyYang/f8aa7b322c48723bac568290de8a08fc

Thanks @HumphreyYang so once we change to jax that graph will become static again (with the new shape)? Also just to confirm, the new shape doesn't change the interpretation in any way re: context of the solution?

Hi @mmcky, yes, the new PRNGs should output the same solution everytime and the interpretation is not changed. Everything is in the same magnitude as the old code so the interpretations are the same.

[mmcky]

thanks @longye-tian for an excellent PR.

thanks @HumphreyYang for your comments and reviews.

@jstac this will serve as another reference in the numba -> jax conversion work.

[mmcky]

@jstac I tagged your review for this one so you have 👀 on what the team has put together.

I'll merge this on Monday.