[kalman] Update to JAX

[xuanguang-li]

This pull request makes updates to the kalman.md notebook, modernizing the codebase by switching from NumPy to JAX.

Key changes include:

• Replaced all numpy (np) operations with jax.numpy (jnp) and updated imports accordingly.
• Refactored the bivariate Gaussian plotting logic: replaced the custom bivariate_normal and gen_gaussian_plot_vals functions with a vectorized, JAX-based implementation using jax.vmap.
• Standardized the use of random seeds in simulations for reproducible results.
• Adding mystnb metadata blocks to each figure-generating code cell.
• Fixed minor typos and improved the clarity.

[github-actions]

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: back_prop, bayes_nonconj, kalman, lqcontrol, mle, qr_decomp

[github-actions]

📖 Netlify Preview Ready!

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

📚 Changed Lecture Pages: back_prop, bayes_nonconj, kalman, lqcontrol, mle, qr_decomp

[mmcky]

Thanks @xuanguang-li -- @HumphreyYang this is now ready for review.

[github-actions]

📖 Netlify Preview Ready!

Preview URL: https://pr-611--sunny-cactus-210e3e.netlify.app (533e30f)

📚 Changed Lecture Pages: back_prop, bayes_nonconj, kalman, lqcontrol, mle, qr_decomp

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

Why previews are down (repo-wide findings)

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

---

This PR touches: kalman.md. Last CI build: success@2025-09-18. Branch: 193 commits behind main as of 2026-07-08.

[xuanguang-li]

@copilot resolve the merge conflicts in this pull request

[Copilot]

@copilot resolve the merge conflicts in this pull request

Resolved by merging origin/main into this branch and fixing the lectures/kalman.md conflicts. Commit: eeb3b74.

[jstac]

Thanks very much for this @xuanguang-li, and apologies that it has sat open for a while. Careful work went into it, and I want to explain properly why we've decided not to merge it — the reasoning turned out to be more general than this one lecture, and it has been genuinely useful to us.

Why we're keeping kalman.md in NumPy

We do have a general policy of moving towards JAX for high-performance code, and for lectures that lead towards it. But this lecture isn't in that category, and working through the diff made clear that JAX here costs more than it buys.

There's no bottleneck to attack. The computations are a 100×100 plotting grid, a T=600 loop whose cost is dominated by scipy.integrate.quad, and a T=50 horse race. None of these is a constraint on build time or on the reader. JAX's advantages — XLA fusion, GPU parallelism, vmap — need something substantial to work on.

The numerics don't actually move to JAX. The arrays become jnp, but the work is still done by scipy.linalg.inv, scipy.stats.norm.pdf, quad, and QuantEcon.py's Kalman and LinearStateSpace — all NumPy-based. So each call converts device arrays back to host arrays and back again. That's a real cost with no offsetting benefit, and it means the lecture takes on JAX's constraints without becoming a JAX lecture. Things like jnp.sqrt(v) where v is a Python float, or jnp.zeros(T) passed to fill_between, are symptoms of this.

Readability goes backwards in one place in particular. multivariate_normal(μ.ravel(), C).pdf(pos) is one line and is transparently the mathematics. The replacement gen_gaussian_plot_vals is ~30 lines with an explicit inv, det, and a vmap. For a plotting grid this size, that's a clear loss on the criterion we care most about — code that is close to the math. Plot-generating code is somewhere we should be relaxed about performance.

One pattern is worth flagging for future work, because it's an easy trap and appears in the exercise solutions:

z = jnp.empty(T)
for t in range(T):
    z = z.at[t].set(...)

Outside jit there is no compiler to elide the copy, so each .at[].set() copies the whole array and this becomes O(T²) where NumPy is O(T). When a loop body has to stay in Python — because it calls SciPy or steps a stateful object like Kalman — plain np.empty plus z[t] = ... is both faster and clearer.

Two smaller things. The GPU admonition implies hardware this lecture doesn't benefit from, which may mislead readers. And Z_F.at[jnp.where(Z_F < 1e-10)].set(0.1), added to make the contours render, sets near-zero densities up to 0.1 — that's mathematically wrong, and it only became necessary because of the plotting rewrite.

Separately, and independently of the JAX question: the branch appears to be based on an earlier version of kalman.md, so it reverts a number of unrelated prose and notation improvements — the X_t / Y_t random-variable notation, the Bishop2006 proof note, and several expanded explanations. Worth knowing for any future branch off this file.

What came out of it

This PR was valuable well beyond itself. It prompted us to consolidate our thinking about when a lecture should use JAX, which we had never written down clearly. Our guidance was scattered across two manual pages that repeated each other, and neither answered the question this PR raises.

That's now fixed: QuantEcon/QuantEcon.manual#113 merges them into a single JAX page with a checklist for when to convert, an idiomatic style guide, and a Numba migration section. It also corrects a few things we had been recommending that we shouldn't have been, and it records the .at[].set() pattern above along with several others. It should make the next conversion much easier to scope and to review.

Please do take a look and comment there if you disagree with any of it, or if the checklist would have steered you differently — that feedback would be genuinely useful before it lands.

We'll close this one on that basis, but I hope you'll pick up another conversion where JAX will really pay off. Value function iteration, large simulations, and anything wanting autodiff are all good candidates, and your work here shows you'd do them well. Thanks again.

[jstac]

Many thanks @xuanguang-li . I appreciate all your hard work.

The above comment was written by Claude under my guidance. I'm sorry it took so long to get to this PR and that I've decided to close it in the end. I think it's been valuable for all of us.

The key takeaway is that we lose on readability without gaining performance, so we'll keep this in NumPy.

CC @mmcky @HumphreyYang @kp992 @longye-tian

[xuanguang-li]

Thanks @jstac! I agree with your point and Claude's review.

One thing that surprised me is the behavior of the .at[].set() function. I tested it separately, and it is indeed incredibly slow without jit. It seems that there is a fixed dispatch overhead every time .at[].set() is called, which dominates the runtime when the array size is relatively small (say, $T < 10^6$). Although repeatedly copying the array inside the loop results in an $O(T^2)$ algorithm, I don't think this is the main bottleneck. When $T$ is small, the array copies themselves are very fast.

I've attached the benchmark results as well, and I think this is a valuable lesson from working on this PR.

T	JAX (µs/op)	NumPy (µs/op)	Implied GB/s
1,000	313.41	0.051	0.0
10,000	275.91	0.051	0.1
100,000	281.74	0.050	1.4
1,000,000	1147.85	0.050	3.5
10,000,000	4149.13	0.050	9.6

[jstac]

Thanks @xuanguang-li . Nice analysis. Yes, it's best avoided, except maybe in a jitted function, since it creates an entire new array just to change one element! This is now noted in the manual.