[mccall_model] JAX conversion of the Job Search Model I #619

longye-tian · 2025-09-17T01:45:59Z

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 · 2025-09-17T02:29:38Z

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 · 2025-09-17T02:47:25Z

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 · 2025-09-17T03:16:18Z

📖 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 · 2025-09-17T03:39:45Z

📖 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:

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

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()

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!

lectures/mccall_model.md

Co-authored-by: Humphrey Yang <39026988+HumphreyYang@users.noreply.github.com>

longye-tian · 2025-09-18T02:12:57Z

Hi @HumphreyYang ,

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

Best,
Longye

github-actions · 2025-09-18T02:21:58Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_model

github-actions · 2025-09-18T02:22:10Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_model

github-actions · 2025-09-18T02:22:35Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_model

github-actions · 2025-09-18T02:22:38Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_model

github-actions · 2025-09-18T02:29:55Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_model

github-actions · 2025-09-18T02:32:29Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_model

github-actions · 2025-09-18T02:34:01Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_model

github-actions · 2025-09-18T02:38:17Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_model

lectures/mccall_model.md

github-actions · 2025-09-21T02:58:46Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_model

longye-tian · 2025-09-21T21:07:30Z

Hi @mmcky ,

I've updated the lecture and revised according to very good suggestions from @HumphreyYang and @shlff .

What do you think?

Best,
Longye

jstac · 2025-10-31T20:26:14Z

@mmcky FYI, I'm getting a 404 on the netlify link above.

I'll probably push some changes now, which might generate a fresh link that works.

@longye-tian @HumphreyYang Actually I disagree with import jax.random as jr, although I appreciate the suggestion. The only reason is that it's nonstandard, so requires a search if someone is trying to work out what it means.

I'll revert the change now and review the lecture.

github-actions · 2025-10-31T21:04:28Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_model

mmcky · 2025-10-31T23:19:23Z

@mmcky FYI, I'm getting a 404 on the netlify link above.

thanks for reporting @jstac. I can't seem to replicate this on my end. The recent ones appear to work for me. Would you mind pasting the link that is 404'ing for you?

HumphreyYang · 2025-10-31T23:28:38Z

@longye-tian @HumphreyYang Actually I disagree with import jax.random as jr, although I appreciate the suggestion. The only reason is that it's nonstandard, so requires a search if someone is trying to work out what it means.

Many thanks @jstac, I think there is an entry in the style guide that suggests us to use import jax.random as jr so I proposed the change in this PR:
https://manual.quantecon.org/styleguide/jax-conversion.html#random-number-generation

I have opened a PR on this in the style sheet:
https://github.com/QuantEcon/QuantEcon.manual/pull/76

jstac · 2025-10-31T23:46:55Z

thanks for reporting @jstac. I can't seem to replicate this on my end. The recent ones appear to work for me. Would you mind pasting the link that is 404'ing for you?

All links working again after my last push. Do they go stale after a while?

Many thanks @jstac, I think there is an entry in the style guide that suggests us to use import jax.random as jr so I proposed the change in this PR:
https://manual.quantecon.org/styleguide/jax-conversion.html#random-number-generation

Ah, thanks for letting me know @HumphreyYang

I have opened a PR on this in the style sheet:
https://github.com/QuantEcon/QuantEcon.manual/pull/76

Nice work.

jstac · 2025-11-01T00:29:49Z

Please hold off on any edits to this PR.

Update mccall_model.md

e860f3d

longye-tian added jax-conversion lecture labels Sep 17, 2025

Merge branch 'main' into job_1

6f9619b

longye-tian marked this pull request as ready for review September 17, 2025 03:49

longye-tian requested review from HumphreyYang, mmcky and shlff September 17, 2025 03:49

HumphreyYang requested changes Sep 18, 2025

View reviewed changes

longye-tian and others added 7 commits September 18, 2025 12:06

Update lectures/mccall_model.md

a1c0a5e

Co-authored-by: Humphrey Yang <39026988+HumphreyYang@users.noreply.github.com>

Update lectures/mccall_model.md

f66be67

Co-authored-by: Humphrey Yang <39026988+HumphreyYang@users.noreply.github.com>

Update lectures/mccall_model.md

94d2339

Co-authored-by: Humphrey Yang <39026988+HumphreyYang@users.noreply.github.com>

Update lectures/mccall_model.md

01c5e80

Co-authored-by: Humphrey Yang <39026988+HumphreyYang@users.noreply.github.com>

Update lectures/mccall_model.md

3711f5b

Co-authored-by: Humphrey Yang <39026988+HumphreyYang@users.noreply.github.com>

Update lectures/mccall_model.md

db5a05f

Co-authored-by: Humphrey Yang <39026988+HumphreyYang@users.noreply.github.com>

Update mccall_model.md

75dc290

Update mccall_model.md

1a748dc

shlff reviewed Sep 19, 2025

View reviewed changes

lectures/mccall_model.md Show resolved Hide resolved

Update mccall_model.md

d9cae86

misc

e75d48a

jstac added the do-not-merge label Nov 1, 2025

jstac mentioned this pull request Nov 1, 2025

[mccall_model] Update McCall model lecture --- take 2! #663

Merged

Uh oh!

[mccall_model] JAX conversion of the Job Search Model I #619

Are you sure you want to change the base?

[mccall_model] JAX conversion of the Job Search Model I #619

Uh oh!

Conversation

longye-tian commented Sep 17, 2025

Overview

Key Changes

1. Library Updates

2. Data Structure Changes

3. Array Operations

4. Loop Transformations

5. Random Number Generation

6. Vectorization

Uh oh!

HumphreyYang commented Sep 17, 2025

Uh oh!

longye-tian commented Sep 17, 2025

Uh oh!

github-actions bot commented Sep 17, 2025

Uh oh!

github-actions bot commented Sep 17, 2025

Uh oh!

HumphreyYang left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

longye-tian commented Sep 18, 2025

Uh oh!

github-actions bot commented Sep 18, 2025

Uh oh!

github-actions bot commented Sep 18, 2025

Uh oh!

github-actions bot commented Sep 18, 2025

Uh oh!

github-actions bot commented Sep 18, 2025

Uh oh!

github-actions bot commented Sep 18, 2025

Uh oh!

github-actions bot commented Sep 18, 2025

Uh oh!

github-actions bot commented Sep 18, 2025

Uh oh!

github-actions bot commented Sep 18, 2025

Uh oh!

Uh oh!

github-actions bot commented Sep 21, 2025

Uh oh!

longye-tian commented Sep 21, 2025

Uh oh!

jstac commented Oct 31, 2025

Uh oh!

github-actions bot commented Oct 31, 2025

Uh oh!

mmcky commented Oct 31, 2025

Uh oh!

HumphreyYang commented Oct 31, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jstac commented Oct 31, 2025

Uh oh!

jstac commented Nov 1, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

HumphreyYang commented Oct 31, 2025 •

edited

Loading