os_egm_jax#

  • Series: lecture-python.myst

  • File: lectures/os_egm_jax.md

  • Audit date: 2026-08-26

  • Corpus snapshot: e25fdf2345

  • Categories audited: writing, math, code, figures, links, admonitions (JAX out of scope)

  • Overall score: 8.6 / 10

  • Priority: NONE

Score breakdown#

Category

Score

One-line note

Writing

8/10

qe-writing-003 ×3; qe-writing-007 ×2.

Math

10/10

no mechanical violations detected.

Code

6.5/10

qe-code-002 ×7; qe-code-001 ×4.

JAX

out of scope

JAX rules target lecture-jax.

Figures

7/10

qe-fig-006 ×2; qe-fig-005 ×2; qe-fig-008 ×1.

References

N/A

no citations in this lecture.

Links

10/10

no mechanical violations detected.

Admonitions

10/10

no mechanical violations detected.

Issues#

Critical#

None found.

High severity#

  • [qe-code-002] — Use Unicode symbols for Greek letters in code. Count: 7. Lines: 158, 160, 375, 376, 377, 378. Example: spelled-out mu.

Medium severity#

  • [qe-code-001] (reviewer) — Follow PEP8 unless closer to mathematical notation. Count: 4. Lines: 151, 286, 209, 390. Example: s names two unrelated quantities inside one nine-line function. 151 unpacks β, μ, s, s_grid, shocks, α = model, where s is the shock scale parameter (declared as such at 97 and 106 and used at 121), and 155 then defines def compute_c(s) whose s is a savings value - so the inner parameter shadows the outer binding, f(s, α) at 157 means savings, and the outer s is never used at all. The sibling lecture avoided this exactly: {doc}`os_egm` calls the shock scale ν (182, 196) and keeps s for savings throughout (122, 245, 253), and the field right beside the collision here is still s_grid, “exogenous savings grid” (98). The JAX rewrite renamed ν to s and produced the shadow. Second, u_crra (286-287) is dead code that would return NaN at the first value the exercise uses: \((c^{1-\gamma} - 1)/(1-\gamma)\) is \(0/0\) at \(\gamma = 1\) and 366 sets γ_values = [1.0, 1.05, 1.1, 1.2]. The exercise survives only because the solver calls u_prime_crra and u_prime_inv_crra, both of which are well defined at \(\gamma = 1\) (\(c^{-1}\) and \(x^{-1}\)) - but 281-283 presents all three functions as needed, so a reader who evaluates the first one at the advertised value gets NaN with no warning. Third, both solvers unpack four values from jax.lax.while_loop and return two, discarding i and error (209, 358), and neither carries the two safeguards the NumPy sibling has - {doc}`os_egm`:306-307 prints the error each iteration under verbose and 309-310 warns when max_iter is hit - so a JAX solve that exhausts its 1000-iteration budget returns silently and looks converged. Fourth, v_star (72-80) is defined and never called, as in the sibling; and 390 and 393 disagree about line width inside one loop, linewidth=2 on the \(\gamma = 1\) branch and nothing on the others, which is the drafted qe-fig-008 finding at 393.

  • [qe-fig-005] — Descriptive figure names for cross-referencing. Count: 2. Lines: 224, 384. Example: code-cell figure without mystnb figure metadata.

  • [qe-fig-006] — Lowercase axis labels. Count: 2. Lines: 395, 396. Example: axis label State x.

  • [qe-writing-003] (reviewer) — Maintain logical flow. Count: 3. Lines: 244, 409, 89. Example: the lecture exists to be faster and never measures against anything. 43 promises “improved performance”, 244 asserts “The JAX implementation is very fast thanks to JIT compilation and vectorization”, 252-257 attributes that speed to four named causes, and the single timing at 246-249 produces a number with nothing to compare it to: the NumPy timing lives in a different file ({doc}`os_egm`:346-348) and even uses a different precision setting, qe.Timer() there against qe.Timer(precision=8) here, so a reader cannot line the two up without running both notebooks. Second, 409-411 compares policies[1.0] with policies[γ] element by element - a comparison at matched savings grid points, since all four solves share model_crra.s_grid - immediately after 402-404 has explained that the endogenous \(x\) grids differ across \(\gamma\). The comparison is well defined and the conclusion at 414 is correct, but nothing says the difference is taken at common \(s\) rather than common \(x\), which is precisely the reading the preceding two sentences prime. Third, 89-91 states a real design rule and the exercise then breaks it without comment: the Model “stores only the data” because “Utility and production functions will be defined globally to work with JAX’s JIT compiler”, and yet \(\gamma\) - a preference parameter - is threaded through as an explicit argument to K_crra and solve_model_crra (303, 337) rather than added to the model, leaving two conventions and no stated reason to prefer either.

  • [qe-writing-007] (reviewer) — Use visual elements to enhance understanding. Count: 2. Lines: 252, 384. Example: the four bullet points at 252-257 - JIT compilation of the whole solver, vectorisation via vmap, jax.lax.while_loop in place of a Python loop, and “efficient JAX array operations throughout” - attribute the lecture’s entire payoff to four causes and quantify none of them. Four bars, or four rows of a table, timing the pure-Python version, vmap alone, vmap plus jit, and the full solver, would make this the most useful figure in the pair; the first and last of those four already exist in {doc}`os_egm`:346-348 and in this lecture at 246-249. Second, the exercise’s figure (384-399) plots four policies on four different \(x\) ranges, an artefact the prose then has to apologise for at 402-404 (“the plots for \(\gamma > 1\) do not cover the entire x-axis range shown”). Plotting consumption against the shared exogenous \(s\) grid, or adding that as a second panel, would remove the artefact and would show the endogenous-grid dependence that 404 describes - which is the one genuinely EGM-specific point in the exercise.

Low severity#

  • [qe-fig-008] — Use lw=2 for line charts. Count: 1. Lines: 393. Example: plot() without lw=.

Strengths#

  • The lecture does not restate the model, and says so precisely once: 59 (“For details on the savings problem and the endogenous grid method (EGM), please see {doc}`os_egm`”) and 63-67 list the three specialisations being inherited, so the file is exactly as long as the JAX content requires and no shorter.

  • The one structural change from the NumPy version is explained rather than just made: 89-91 says the Model now stores only data and that the utility and production functions are global “to work with JAX’s JIT compiler” - the reader learns why the sibling’s Callable fields disappeared instead of noticing it later.

  • The JIT timing is done correctly, which is easy to get wrong: solve_model_time_iter is @jax.jit (183) and is called once at 219 before the timed call at 248, so compilation happens outside the qe.Timer() block, and jax.block_until_ready(c) at 249 forces the asynchronous dispatch to complete before the timer stops - without it the measurement would be of dispatch rather than of work.

  • The vmap structure mirrors the algebra: compute_c (155-160) is written for a single savings point exactly as {eq}`egm_getc` of {doc}`os_egm` is written for a single \(s_i\), with the same comment naming the integral it approximates, and 163-164 then vectorises it in two lines - so the reader sees the scalar operator and its vectorisation separately.

  • The exercise is a genuine numerical experiment rather than a re-run: it asks for CRRA policies at \(\gamma\) approaching 1 from above and for the demonstration that they converge to the log case, gives a hint explaining why values close to 1 are needed (274, so the endogenous grids have comparable coverage), and then checks the convergence numerically at 408-411 rather than resting on the figure.

  • The exercise’s solution anticipates the reader’s objection to its own figure: 402-404 explains that the \(\gamma > 1\) curves stop short because the endogenous grid \(x = s + \sigma(s)\) moves with the policy - which is the most EGM-specific observation anywhere in the pair.

  • Timing uses qe.Timer(precision=8) (247) rather than a %%time magic, and the precision is raised deliberately because the quantity being measured is small - which is what qe-code-004 asks for, done with attention.