[mccall_model] JAX conversion of the Job Search Model I

[longye-tian]

I've updated the Job Search Model I, align with the change we've had before on Job Search Model II and III.

Overview

Converted lectures/mccall_model.md from NumPy/Numba to JAX implementation, following the patterns established in mccall_model_with_separation.md and mccall_fitted_vfi.md.

Key Changes

1. Library Updates

• Replaced: numpy → jax.numpy (jnp)
• Replaced: numba.jit → jax.jit
• Removed: numba.experimental.jitclass
• Added: jax, jax.numpy, typing.NamedTuple

2. Data Structure Changes

• Before: Used Numba's @jitclass for the McCall model class
• After: Used NamedTuple for immutable model structure, compatible with JAX's functional paradigm

3. Array Operations

• Converted: All NumPy arrays to JAX arrays (jnp)
• Updated: Array mutations to use JAX's functional update syntax (.at[].set())

4. Loop Transformations

• Before: Standard Python while loops with Numba JIT
• After: jax.lax.while_loop with condition and body functions

5. Random Number Generation

• Before: NumPy's global random state (np.random.seed())
• After: JAX's explicit PRNG key system (jax.random.PRNGKey)

6. Vectorization

• Before: Manual loops over parameter grids
• After: jax.vmap for automatic vectorization

[HumphreyYang]

Many thanks @longye-tian!

I thought @shlff self-assigned this lecture in the issue here: QuantEcon/meta#173.

Just CCing @shlff here so we are not working on the same lecture separately : )

[longye-tian]

Many thanks @longye-tian!

I thought @shlff self-assigned this lecture in the issue here: QuantEcon/meta#173.

Just CCing @shlff here so we are not working on the same lecture separately : )

Thank you for your notice @HumphreyYang .

So Sorry @shlff . I forgot to check this list...

[github-actions]

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: back_prop, bayes_nonconj, divergence_measures, imp_sample, likelihood_ratio_process, likelihood_ratio_process_2, likelihood_var, mccall_model, mle, newton_method, qr_decomp

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-619--sunny-cactus-210e3e.netlify.app (6f9619b)

📚 Changed Lecture Pages: mccall_model

[HumphreyYang]

Many thanks @longye-tian! Nice changes!!

Just some minor suggestions in the review.

Three minor general suggestions:

1. There are some comments like this in the code:

    # == Now compute the reservation wage == #

It would be nice to convert them to more standard comments like

# Now compute the reservation wage

2. We used jax.random many times. It might be a good idea to give it an alias in the import:

import jax.random as jr

and then use

jr.method_name()

3. There are a few more lines that are longer than the box in the preview. It would be nice if you could follow the pattern in the review comment and cut them to fewer than 80 characters so they fit into the code cell in the preview nicely!

[longye-tian]

Hi @HumphreyYang ,

Thank you for your detailed review. I've updated the lecture accordingly.

Best,
Longye

[github-actions]

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_model

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-619--sunny-cactus-210e3e.netlify.app (94d2339)

📚 Changed Lecture Pages: mccall_model

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-619--sunny-cactus-210e3e.netlify.app (01c5e80)

📚 Changed Lecture Pages: mccall_model

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-619--sunny-cactus-210e3e.netlify.app (3711f5b)

📚 Changed Lecture Pages: mccall_model

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-619--sunny-cactus-210e3e.netlify.app (1a748dc)

📚 Changed Lecture Pages: mccall_model

[github-actions]

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_model

[github-actions]

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_model

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-619--sunny-cactus-210e3e.netlify.app (75dc290)

📚 Changed Lecture Pages: mccall_model