[kesten_processes] Kesten Processes lecture update

[kp992]

This PR:

• Updates the lecture with latest style guide
• Replaces numba code with JAX
• minor typo fixes

[github-actions]

🚀 Deployed on https://68bb3aa8cd486a0a579d9ae4--nostalgic-wright-5fa355.netlify.app

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-591--sunny-cactus-210e3e.netlify.app (5b0ce92)

📚 Changed Lecture Pages: kesten_processes

[HumphreyYang]

Many thanks @kp992, it looks really nice!

Just some minor comments!

[HumphreyYang]

I think one trick to avoid many Nones is by multiplying list and then unwrap:

    vectorized_single_draw = vmap(
        generate_single_draw,
        in_axes=(0, *[None]*9),
    )

[mmcky]

Why are there so many None needed here? Seems like a crazy function interface choice (coming from someone who has never used vmap before.

[HumphreyYang]

I think they want to take into account mapping across all possible dimensions and output of all possible shapes so they ended up having this massive choice set.

[kp992]

Okay, I will change this a bit to keep less None

[mmcky]

@kp992 @HumphreyYang I haven't come across this error before

Unable to initialize backend 'cuda': INVALID_ARGUMENT: Unexpected option name passed to PJRT_Client_Create: use_tfrt_gpu_client (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend

I don't see this happening in other PRs -- but looks like an install issue?

It is coming from

@jit
def generate_draws(
    key=random.PRNGKey(123),
    μ_a=-0.5,
    σ_a=0.1,
    μ_b=0.0,
    σ_b=0.5,
    μ_e=0.0,
    σ_e=0.5,
    s_bar=1.0,
    T=500,
    M=1_000_000,
    s_init=1.0,
):
    """
    JAX-jit version of the generate_draws function.
    Returns:
        Array of M draws
    """
    # Create M different random keys for parallel execution
    keys = random.split(key, M)

    draws = vectorized_single_draw(
        keys, μ_a, σ_a, μ_b, σ_b, μ_e, σ_e, s_bar, T, s_init
    )

    return draws

[HumphreyYang]

Hi @mmcky,

The code runs well on my GPU, I think this might be related to the environment. I will retrigger the action and see if it helps.

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-591--sunny-cactus-210e3e.netlify.app (6e44c3c)

📚 Changed Lecture Pages: kesten_processes

[HumphreyYang]

Hi @mmcky,

This is fixed now. It's a bit tricky to find but it is caused by a minor typo :tags: [no-execute] -> :tags: [skip-execution].

[kp992]

The code runs well on my GPU, I think this might be related to the environment.

Thanks for fixing this. Even for me it worked fine. Probably we can work on making debug messages easier.

[HumphreyYang]

The code runs well on my GPU, I think this might be related to the environment.

Thanks for fixing this. Even for me it worked fine. Probably we can work on making debug messages easier.

Many thanks @kp992, I was expecting bug like this to be raised by jupyter but not jax. It's interesting it is raised as a device interface issue : )

[mmcky]

The code runs well on my GPU, I think this might be related to the environment.

Thanks for fixing this. Even for me it worked fine. Probably we can work on making debug messages easier.

Many thanks @kp992, I was expecting bug like this to be raised by jupyter but not jax. It's interesting it is raised as a device interface issue : )

The no-execute mean that that cell was still executed -- as it is an invalid tag. Nice work @HumphreyYang

[mmcky]

thanks @kp992 you read my mind! Was just about to make the edit :)

Will merge this soon. Thanks for your PR.

[mmcky]

@kp992 once this finishes building -- can you please do a final output check of the preview link and give me your thumbs up.

[github-actions]

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: kesten_processes

[kp992]

Looks good to me @mmcky

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-591--sunny-cactus-210e3e.netlify.app (9bead7e)

📚 Changed Lecture Pages: kesten_processes

[jstac]

@mmcky I'm not seeing any JAX on the 'Changed lecture pages' link above.

[mmcky]

@mmcky I'm not seeing any JAX on the 'Changed lecture pages' link above.

@jstac the jax code in this lecture is contained in the exercise solutions - but do let me know if you don't see it on your end.

[mmcky]

We have imports in the solution

We typically have a rule that all imports are at the top of the lecture, however I am in two minds re: solutions. Half of me thinks keeping imports allowed in solutions can be useful when introducing a new package that hasn't been used in the lecture itself. Any thoughts on this style-guide item @jstac

[jstac]

Ah, I see. Thanks @mmcky

We should shift this lecture to JAX throughout.

All simulation of paths in JAX should use the generate_paths functionality from the lake model PR #589

[HumphreyYang]

Many thanks @kp992, @jstac, and @mmcky!

I will have a go at converting them into the new template!

[mmcky]

thanks @HumphreyYang I have started on the documentation in the manual using the previous PR.

I will use this one as a final cross check and then do a final edit of https://github.com/QuantEcon/QuantEcon.manual/pull/73/files.

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-591--sunny-cactus-210e3e.netlify.app (7d00832)

📚 Changed Lecture Pages: kesten_processes

[mmcky]

thanks @HumphreyYang for these changes.

@jstac -- I am in favour of consistency, but this change seems to be adding quite a few lines of code (making it more complex) -- so I just wanted to run this template by you to make sure we are heading down the right path here.

[jstac]

@mmcky This is, in a sense, the cost of using JAX. JAX wasn't built for iteration and has no immutable arrays.

If we're going to keep using the lax.scan logic, then, given it's slightly unusual syntax, I think it's better to pull that logic out into a specialized function and explain it.

I recommend not hiding this code, since it's essential to the lecture.

I'm very open to other ideas and approaches.

The other option is refrain from JAX conversions for lectures that iterate a low dimensional object forward many times. Admittedly, numpy and numba are better for this. I'm open to this discussion too.

[mmcky]

Thanks @jstac.

If we're going to keep using the lax.scan logic, then, given it's slightly unusual syntax, I think it's better to pull that logic out into a specialized function and explain it.

Completely agree.

I recommend not hiding this code, since it's essential to the lecture.

👍

I'm very open to other ideas and approaches.

I guess I would like to have tabbed code blocks available - giving the choice (But that is more to maintain, so committing is probably optimal)

[mmcky]

The other option is refrain from JAX conversions for lectures that iterate a low dimensional object forward many times. Admittedly, numpy and numba are better for this. I'm open to this discussion too.

Are you thinking of something like if lecture > 10s runtime then use jax otherwise use numba kind of rule to figure out when to apply it?

I will mull over this tonight.

[mmcky]

@HumphreyYang you've done a great job with this conversion. It is going to be a really useful pattern to finalise the docs.

[github-actions]

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: kesten_processes

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

Why previews are down (repo-wide findings)

1. This branch is stale — 178 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 ~178 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-591--sunny-cactus-210e3e.netlify.app/kesten_processes.html

---

This PR touches: kesten_processes.md. Last CI build: success@2025-09-29. Branch: 178 commits behind main as of 2026-07-08.