[mccall_q]: JAX conversion by bishmaybarik · Pull Request #624 · QuantEcon/lecture-python.myst

bishmaybarik · 2025-09-25T03:03:15Z

This PR is created to convert the mccall_q.md lecture codes into JAX equivalents.

The JAX equivalent of all codes should be written
Fixing the Headings and Sub-Headings to satisfy Quantecon style guidelines
Check for grammatical errors and also see if things are easy to follow through
Adjust the figure legends, as their current position above the curves makes the graphs difficult to interpret.
Rectify the codes with @HumphreyYang 's suggestions
Comment out the codes neatly

github-actions · 2025-09-25T05:54:49Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_q

github-actions · 2025-09-26T04:59:02Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_q

earlier the legend overlapped on the curves, making it difficult to interpret the graph. this version fixes it

github-actions · 2025-09-26T07:32:51Z

📖 Netlify Preview Ready!

Preview URL: https://pr-624--sunny-cactus-210e3e.netlify.app (211832b)

📚 Changed Lecture Pages: mccall_q

github-actions · 2025-09-26T07:48:16Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_q

github-actions · 2025-09-26T07:58:31Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_q

Copilot

Pull Request Overview

This PR converts the McCall job search model lecture from NumPy/Numba implementation to JAX, updating the code to use JAX's functional programming paradigm and JIT compilation. The conversion includes modernizing the code structure, fixing grammatical errors, and improving figure presentation.

Key changes include:

Complete conversion from NumPy/Numba to JAX with functional programming approach
Replacement of class-based structure with NamedTuple and standalone functions
Implementation of JAX's immutable array operations and control flow primitives

Copilot · 2025-09-29T12:51:31Z

lectures/mccall_q.md

+        qtable, s, accept_count, t, key = state
+        # for first interaction, just continue since error is large
+        # for subsequent interactions, compute actual error
+        error = jnp.where(t==0, δ + 1, jnp.max(jnp.abs(qtable - state[0])))


The error calculation references state[0] but state is a tuple where the first element is qtable. This creates a circular reference where the error is computed as the difference between qtable and itself, which will always be zero after the first iteration. This should reference a previous qtable value stored separately.

Copilot · 2025-09-29T12:51:31Z

lectures/mccall_q.md

-
    for n in range(N):
-        if n%(N/10)==0:
+        if n % (N // 10) == 0:


Integer division N // 10 can result in zero when N < 10, causing a division by zero error in the modulo operation. This should use max(1, N // 10) to ensure the divisor is never zero.

Suggested change

if n % (N // 10) == 0:

if n % max(1, N // 10) == 0:

Copilot · 2025-09-29T12:51:31Z

lectures/mccall_q.md

+      max_epochs = int(jnp.max(epochs_to_plot))  # Convert to Python int
+      # iterate on epoch numbers
+      for n in range(max_epochs + 1):
+          if n%(max_epochs/10)==0:


Division by zero error can occur when max_epochs is 0. The condition n%(max_epochs/10)==0 will fail. This should use max_epochs > 0 and n%(max_epochs/10)==0 or n%(max(1, max_epochs/10))==0.

Suggested change

if n%(max_epochs/10)==0:

if n % max(1, max_epochs // 10) == 0:

HumphreyYang

Many thanks @bishmaybarik! Nice changes -- they are really good improvements. Here are some minor comments.

lectures/mccall_q.md

bishmaybarik · 2025-09-30T10:13:42Z

hi @HumphreyYang , these are extremely useful suggestions -- thanks a lot! I'll make all the necessary changes and push the updated version once ready.

vectorized computation instead of repeatedly resetting arrays with for loops

github-actions · 2025-10-01T06:48:20Z

📖 Netlify Preview Ready!

Preview URL: https://pr-624--sunny-cactus-210e3e.netlify.app (9143f6b)

📚 Changed Lecture Pages: mccall_q

github-actions · 2025-10-01T07:07:21Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_q

github-actions · 2025-10-01T07:15:27Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_q

github-actions · 2025-10-01T07:17:23Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_q

github-actions · 2025-10-06T05:49:25Z

📖 Netlify Preview Ready!

Preview URL: https://pr-624--sunny-cactus-210e3e.netlify.app (06e07c0)

📚 Changed Lecture Pages: mccall_q

github-actions · 2025-10-06T05:58:49Z

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: mccall_q

Copilot

Pull Request Overview

Copilot reviewed 1 out of 1 changed files in this pull request and generated 1 comment.

Copilot · 2025-10-06T08:37:02Z

lectures/mccall_q.md

+        qtable, s, accept_count, t, key = state
+        # for first interaction, just continue since error is large
+        # for subsequent interactions, compute actual error
+        error = jnp.where(t==0, δ + 1, jnp.max(jnp.abs(qtable - state[0])))


The error calculation references state[0] which should be the previous qtable. This creates a circular reference since state is the current state tuple. Consider storing the previous qtable separately or using a different approach to calculate the error.

bishmaybarik · 2025-10-13T02:27:11Z

hi @HumphreyYang and @mmcky , may I know if you have any thoughts on these updates? It would be great to hear from you!

update import block

edd7608

bishmaybarik self-assigned this Sep 25, 2025

bishmaybarik added 2 commits September 25, 2025 09:07

convert mccall model from numba to jax

61004bc

complete jax conversion of q-learning implementation

c8814f3

bishmaybarik added in-work jax-conversion labels Sep 25, 2025

fix headings and subheadings to match style guide

c575e39

bishmaybarik added 3 commits September 26, 2025 12:31

rectify grammatical errors in the lecture

211832b

fix legend and place it below graph

d7b49a6

earlier the legend overlapped on the curves, making it difficult to interpret the graph. this version fixes it

update jupytext version

ef3e791

bishmaybarik marked this pull request as ready for review September 26, 2025 07:30

bishmaybarik requested review from HumphreyYang and mmcky September 26, 2025 07:31

bishmaybarik added review and removed in-work labels Sep 26, 2025

HumphreyYang requested a review from Copilot September 29, 2025 12:50

Copilot AI reviewed Sep 29, 2025

View reviewed changes

HumphreyYang reviewed Sep 30, 2025

View reviewed changes

bishmaybarik added 4 commits October 1, 2025 11:48

use seed equals 0 as default

9143f6b

use vectorized computation

a72d9b3

vectorized computation instead of repeatedly resetting arrays with for loops

replace eps (epsilon) with the unicode character

f1a7dca

remove unnecessary comment

f53c4ef

bishmaybarik added 10 commits October 1, 2025 12:58

put default values of variables into the class

89d1b09

shorten code length to have each line under 80 characters

088f91f

follow copilot's suggestion to prevent division by 0

75db66d

shorten lengthy notes

6613e32

replace loop with jax.vmap for vectorized value update

d6034fe

make PRNG key explicit in plot_epochs calls for consistent results

a65f940

prevent division by zero as per copilot

659554f

replace valfunc_VFI with valfunc_vfi throughout the lecture

dbcfd78

fix a typo in the lecture

06e07c0

add minor comments into the code

d5d2565

bishmaybarik requested a review from HumphreyYang October 6, 2025 05:42

bishmaybarik requested a review from Copilot October 6, 2025 08:36

Copilot AI reviewed Oct 6, 2025

View reviewed changes

	if n%(max_epochs/10)==0:
	if n % max(1, max_epochs // 10) == 0:

Uh oh!

Conversation

bishmaybarik commented Sep 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Sep 25, 2025

Uh oh!

github-actions bot commented Sep 26, 2025

Uh oh!

github-actions bot commented Sep 26, 2025

Uh oh!

github-actions bot commented Sep 26, 2025

Uh oh!

github-actions bot commented Sep 26, 2025

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Uh oh!

Copilot AI Sep 29, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Sep 29, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Sep 29, 2025

Choose a reason for hiding this comment

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!

Uh oh!

bishmaybarik commented Sep 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Oct 1, 2025

Uh oh!

github-actions bot commented Oct 1, 2025

Uh oh!

github-actions bot commented Oct 1, 2025

Uh oh!

github-actions bot commented Oct 1, 2025

Uh oh!

github-actions bot commented Oct 6, 2025

Uh oh!

github-actions bot commented Oct 6, 2025

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Uh oh!

Copilot AI Oct 6, 2025

Choose a reason for hiding this comment

Uh oh!

bishmaybarik commented Oct 13, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

bishmaybarik commented Sep 25, 2025 •

edited

Loading

bishmaybarik commented Sep 30, 2025 •

edited

Loading