[markov_asset] Update to JAX and runtime comparison

[xuanguang-li]

This pull request modernizes the markov_asset.md notebook by refactoring the codebase to use JAX and JAX NumPy.

• Refactored all numerical computations and linear algebra from numpy and numpy.linalg to jax.numpy and jax.numpy.linalg.
• Replaced the mutable AssetPriceModel class with an immutable, type-safe NamedTuple version and introduced a create_ap_model factory function for consistent construction.
• Updated function signatures and internal logic to use the new model structure.
• Replaced manual Python loops with JAX functional control flow constructs such as jax.lax.while_loop and jax.lax.fori_loop.
• Fixed minor bugs in parameter initialization, state handling, and tolerance defaults.

Update: Runtime Comparison Between JAX (GPU), JAX (CPU), and NumPy

Methodology:

• The JAX version uses the code in this PR, while the NumPy version uses the code in main.
• The runtime for JAX (GPU) is measured using Google Colab.
• Runtime is collected using qe.timeit over 1,000 iterations.
• Each iteration consists of one function that solves the consol price and one that solves the call option price.

Results:

• Average runtime: JAX (CPU) > NumPy > JAX (GPU).
• More details are shown in the attached plots.

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-654--sunny-cactus-210e3e.netlify.app (0317e4a)

📚 Changed Lecture Pages: harrison_kreps, markov_asset, two_auctions

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-654--sunny-cactus-210e3e.netlify.app (76c64b2)

📚 Changed Lecture Pages: harrison_kreps, markov_asset, two_auctions

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-654--sunny-cactus-210e3e.netlify.app (02dfe6e)

📚 Changed Lecture Pages: harrison_kreps, markov_asset, two_auctions

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-654--sunny-cactus-210e3e.netlify.app (211ec24)

📚 Changed Lecture Pages: harrison_kreps, markov_asset, two_auctions

[github-actions]

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: markov_asset

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-654--sunny-cactus-210e3e.netlify.app (4adc2f5)

📚 Changed Lecture Pages: markov_asset

[github-actions]

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: markov_asset, markov_asset_numpy

[jstac]

🤖 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-654--sunny-cactus-210e3e.netlify.app/ currently returns 404.

Why previews are down (repo-wide findings)

1. This branch is stale — 164 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).

Note on recent timeline activity

This PR was close/reopened on 2026-07-07 by a maintainer session purely to trigger a rebuild test — not a content change. That rebuild failed on the stale-branch issue above. Apologies for the notification churn.

Recommended first step for this PR

Update this branch to main (merge or rebase — pulls in #939 plus ~164 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-654--sunny-cactus-210e3e.netlify.app/markov_asset.html

---

This PR touches: markov_asset.md, markov_asset_numpy.md. Last CI build: failure@2026-07-07. Branch: 164 commits behind main as of 2026-07-08.

[github-actions]

📖 Netlify Preview Ready!

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

Commit: 533f572

📚 Changed Lectures

• markov_asset
• markov_asset_numpy

---

Build Info

• Workflow: Build Project [using jupyter-book]

[jstac]

This seems to be added by mistake.

[jstac]

Many thanks @xuanguang-li , much appreciated.

By the way, you added a file to this PR that I assume is a mistake: https://github.com/QuantEcon/lecture-python.myst/pull/654/changes#diff-200962c576f4fdd31f65eea1717388dafb176a3bb5ec8630dbe0840a61f49f57

You can leave this PR with me. I'll add any necessary fixes and then merge it.

[jstac]

Hi @xuanguang-li . (CC @mmcky @HumphreyYang )

I've now reviewed this carefully, including testing with Claude on my CPU and GPU.

I always ask Claude to read our current style guide before assessing. Here it recommends keeping the lecture in NumPy and I agree. JAX is slower on the GPU, and only somewhat faster if we pin it to using the CPU.

This comes at some readability cost, with the pin to the CPU, lax.while loop, etc.

For that reason, I've decided to close the PR without merging. This lecture can stay in NumPy.

FYI, here is Claude's report:

The GPU is ~5–6× slower than NumPy for this lecture, plus a ~0.3 s one-time compile hit. Root cause: call_option's jax.lax.while_loop
runs 182 iterations to converge, and each iteration is a tiny 25×25 matmul + a scalar convergence check. That's a latency-bound,
sequential workload — 182 back-to-back GPU kernel launches with device round-trips per iteration. GPUs are throughput devices; this is
close to their worst case. NumPy runs all 182 iterations in-process with no launch overhead.

This resolves the apparent contradiction with the PR description (which reported GPU as fastest): that benchmark was almost certainly
measuring JAX's async dispatch, not completed compute — a classic JAX benchmarking trap where you time how fast work is queued rather
than finished. Once you sync the device, the GPU advantage disappears and reverses.

Bottom line: this empirically confirms style-guide point #4 — markov_asset is not computationally intensive (tiny 25×25 / 5×5
matrices), so JAX/GPU delivers no benefit and actively hurts (compile latency + slower steady-state). The author's "very long runtime
on GPU" is real, expected, and inherent to the problem size, not a bug to fix.

My recommendation to you

I'd lean toward not converting this lecture to JAX at all — keep it NumPy per the style guide's explicit carve-out — rather than
merging a slower, less-readable (checkify-heavy) version. If the migration tracker (meta#299) mandates conversion regardless, then at
minimum items 1–3 from the saved response still need fixing first.