[inventory_dynamics] Update code to JAX and latest style guide by kp992 · Pull Request #623 · QuantEcon/lecture-python.myst

kp992 · 2025-09-23T05:26:41Z

This PR:

Removes numba and numpy support.
Updates all the code to use JAX.
fixes minor typos and title errors.

github-actions · 2025-09-23T18:00:55Z

📖 Netlify Preview Ready!

Preview URL: https://pr-623--sunny-cactus-210e3e.netlify.app (2b7c4fe)

📚 Changed Lecture Pages: inventory_dynamics

Copilot

Pull Request Overview

This PR updates the inventory dynamics lecture to use JAX instead of numba and numpy, modernizing the codebase and improving performance. The changes align with the latest style guide and fix minor formatting issues.

Key changes:

Replaces numba-based Firm class with JAX-based functions using NamedTuple
Converts simulation functions to use JAX's functional programming paradigm with jax.lax.scan and jax.vmap
Updates all numpy operations to JAX numpy equivalents

github-actions · 2025-09-30T10:21:42Z

📖 Netlify Preview Ready!

Preview URL: https://pr-623--sunny-cactus-210e3e.netlify.app (54d1cac)

📚 Changed Lecture Pages: inventory_dynamics

HumphreyYang

Many thanks @kp992! These are really nice! Just two very minor observations.

I think we should remove the grid in the figures to maintain consistency in the style:

for ax in axes:
    ax.grid(alpha=0.4)

HumphreyYang · 2025-10-21T05:22:39Z

+    # Generate independent random keys for each firm
+    firm_keys = random.split(key, (num_firms, sim_length))
+    # Run simulation for all firms
+    restock_indicators = vectorized_simulate(firm_keys)
+    # Compute frequency (fraction of firms that restocked > 1 times)
+    frequency = jnp.mean(restock_indicators)
+    return frequency


I think it might look better if we add some spaces in between:

Suggested change

# Generate independent random keys for each firm

firm_keys = random.split(key, (num_firms, sim_length))

# Run simulation for all firms

restock_indicators = vectorized_simulate(firm_keys)

# Compute frequency (fraction of firms that restocked > 1 times)

frequency = jnp.mean(restock_indicators)

return frequency

# Generate independent random keys for each firm

firm_keys = random.split(key, (num_firms, sim_length))

# Run simulation for all firms

restock_indicators = vectorized_simulate(firm_keys)

# Compute frequency (fraction of firms that restocked > 1 times)

frequency = jnp.mean(restock_indicators)

return frequency

jstac · 2026-07-08T06:03:31Z

🤖 Status note for a future session — from a maintainer investigation on 2026-07-08 into why open-PR previews 404. Context only, not instructions.

Netlify preview: https://pr-623--sunny-cactus-210e3e.netlify.app/ currently returns 404.

Why previews are down (repo-wide findings)

1. This branch is stale — 174 commits behind main. A preview build compiles the whole site from this branch. This branch's lectures/house_auction.md still has unpinned !pip install prettytable, which now breaks on a wcwidth incompatibility. main fixed this on 2026-06-28 by pinning prettytable<3.18 (#939). This alone fails any rebuild of this branch until it's updated to main.

2. The arviz failure was a red herring — do NOT pin arviz or rewrite plotting. A 2026-07-07 rebuild also failed in ar1_bayes/ar1_turningpts with an arviz_plots figsize ValueError. That was a transient bug in an intermediate arviz-plots 1.x release, already fixed in arviz 1.2.0. Verified locally on a clean latest-stack venv: the real az.plot_trace(trace) cell (pymc + numpyro InferenceData) runs green. The lectures use only 1.x-compatible arviz APIs (plot_trace, summary, from_numpyro, compare).

Recommended first step for this PR

Update this branch to main (merge or rebase — pulls in #939 plus ~174 other commits), then let CI rebuild. On today's latest libraries the site builds clean, so the preview should return. house_auction is the known blocker; updating also picks up other since-merged fixes — rebuild and address any remaining per-lecture failures. Verify with:

curl -sI https://pr-623--sunny-cactus-210e3e.netlify.app/inventory_dynamics.html

This PR touches: inventory_dynamics.md. Last CI build: success@2025-09-30. Branch: 174 commits behind main as of 2026-07-08.

github-actions · 2026-07-22T04:44:53Z

📖 Netlify Preview Ready!

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

Commit: 8139bb3

📚 Changed Lectures

inventory_dynamics

Build Info

Workflow: Build Project [using jupyter-book]

kp992 · 2026-07-22T05:35:12Z

Thanks @jstac for updating this.

@jit

Convert the lecture from Numba (jitclass / @jit(parallel=True)) to JAX: - Replace the `Firm` jitclass with a `NamedTuple` and pure update functions. - Drive the time loop with `jax.lax.fori_loop`, which reads like the Python for-loops it replaces, and generate a fresh per-period shock with `random.fold_in(key, t)`. Each path therefore takes a single key, so the calling code (and every `jax.vmap`) needs only a one-dimensional array of keys instead of a 2-D key grid. - Use the modern typed-key API (`random.key`) rather than `random.PRNGKey`. - Vectorize the Monte Carlo draw loops with `jax.vmap` (including the 50,000-draw marginal-distribution cell, which as a Python loop over a jitted call ran ~28s on GPU; vmapped it is ~0.3s). - Use `sample.min()/.max()` in `plot_kde` instead of Python `min()/max()`, which iterate a JAX array element-by-element (~4s per call). - Add a note that `X.at[t].set(...)` is compiled to an in-place update under `jax.jit`, so the functional style costs no per-period allocation. Fix a correctness bug in the restock-frequency exercise: the previous `simulate_firm_restocks` captured the global `x_init` (= 50 by that point in the lecture) rather than taking it as an argument, so `compute_freq`'s `x_init=70` was silently ignored. The exercise asks for firms starting at X_0 = 70; the code answered for X_0 = 50, reporting ~0.96 instead of the correct ~0.45. Passing `x_init` explicitly through the vmapped call fixes it. Also: give the two exercise-solution helpers distinct names so they no longer shadow each other, and minor typo, terminology, and sentence-case title fixes. Whole-lecture execution drops from ~57s to ~5s on an RTX 4080. Co-authored-by: Humphrey Yang <39026988+HumphreyYang@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

jstac · 2026-07-22T07:08:40Z

Many thanks @kp992 , merging.

mmcky · 2026-07-23T03:47:41Z

\translate-resync

Triggering the first sync now that the zh-cn workflow is live (#979, merged today). This PR merged before the workflow existed, so its change to inventory_dynamics.md never reached QuantEcon/lecture-python.zh-cn — translate status currently reads it as OUTDATED. Verifies the end-to-end path for QuantEcon/project-translation#14 step 4.

mmcky · 2026-07-23T03:48:46Z

✅ Translation sync completed (zh-cn)

Target repo: QuantEcon/lecture-python.zh-cn
Translation PR: QuantEcon/lecture-python.zh-cn#198
Files synced (1):

lectures/inventory_dynamics.md

kp992 mentioned this pull request Sep 23, 2025

Lecture Improvement Project QuantEcon/meta#173

Open

kp992 added jax-conversion improve labels Sep 23, 2025

kp992 added the ready label Sep 24, 2025

kp992 requested review from HumphreyYang and mmcky September 24, 2025 00:38

mmcky added review and removed ready labels Sep 24, 2025

mmcky requested a review from Copilot September 24, 2025 01:04

Copilot AI reviewed Sep 24, 2025

View reviewed changes

HumphreyYang closed this Sep 30, 2025

HumphreyYang reopened this Sep 30, 2025

HumphreyYang reviewed Oct 21, 2025

View reviewed changes

Chihiro2000GitHub mentioned this pull request Jul 21, 2026

Update NumPy random seed usage → recommended Generator approach across lecture repos QuantEcon/meta#299

Open

56 tasks

jstac force-pushed the inventory_dyns branch from 54d1cac to 27876ed Compare July 22, 2026 04:30

jstac force-pushed the inventory_dyns branch from 27876ed to 8139bb3 Compare July 22, 2026 06:03

jstac merged commit c6d14eb into main Jul 22, 2026
1 check passed

jstac deleted the inventory_dyns branch July 22, 2026 07:08

mmcky mentioned this pull request Jul 23, 2026

🌐 [translation-sync] [inventory_dynamics] Update code to JAX and latest style guide QuantEcon/lecture-python.zh-cn#198

Merged

mmcky mentioned this pull request Jul 23, 2026

[perm_income.md] Update np.random → Generator API #975

Merged

Uh oh!

Uh oh!

Conversation

kp992 commented Sep 23, 2025

Uh oh!

github-actions Bot commented Sep 23, 2025

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Sep 30, 2025

Uh oh!

HumphreyYang left a comment

Choose a reason for hiding this comment

Uh oh!

HumphreyYang Oct 21, 2025

Choose a reason for hiding this comment

Uh oh!

jstac commented Jul 8, 2026

Why previews are down (repo-wide findings)

Recommended first step for this PR

Uh oh!

github-actions Bot commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📖 Netlify Preview Ready!

📚 Changed Lectures

Uh oh!

kp992 commented Jul 22, 2026

Uh oh!

jstac commented Jul 22, 2026

Uh oh!

Uh oh!

mmcky commented Jul 23, 2026

Uh oh!

mmcky commented Jul 23, 2026

✅ Translation sync completed (zh-cn)

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

github-actions Bot commented Jul 22, 2026 •

edited

Loading