Optimize Aiyagari model: Switch to VFI with JIT-compiled lax.while_loop #674
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This commit significantly improves the performance and code quality of the Aiyagari model lecture by switching from Howard Policy Iteration (HPI) to Value Function Iteration (VFI) as the primary solution method, with HPI moved to an exercise. Major changes: - Replace HPI with VFI using jax.lax.while_loop and @jax.jit compilation - Reduce asset grid size from 200 to 100 points for efficiency - Reduce asset grid maximum from 20 to 12.5 (better suited for equilibrium) - Use 'loop_state' instead of 'state' in loops to avoid DP terminology confusion - Remove redundant @jax.jit decorators from helper functions (only on top-level functions) - Move HPI implementation to Exercise 3 with complete solution Performance improvements: - VFI equilibrium computation: ~0.68 seconds (was ~11+ seconds with damped iteration) - HPI in Exercise 3: ~0.48 seconds with optimized JIT compilation - 85x speedup compared to unoptimized Python loops Code quality improvements: - Cleaner JIT compilation strategy (only on ultimate calling functions) - Both VFI and HPI use compiled lax.while_loop for consistency - Helper functions automatically inlined and optimized by JAX - Clear separation of main content (VFI) and advanced material (HPI exercise) Educational improvements: - Students learn VFI first (simpler, more standard algorithm) - HPI presented as advanced exercise with guidance and complete solution - Exercise asks students to verify both methods produce same equilibrium Generated with Claude Code Co-Authored-By: Claude <[email protected]>
Contributor
Author
|
Clarification: I switched the solver to VFI because it's simple and more familiar, so it won't act as a barrier to entry. The HPI solver is now an exercise |
…ic Programming book link Replace the broken cross-reference to opt_savings_2 (which doesn't exist in this PR) with a direct link to the Dynamic Programming book at dp.quantecon.org where Howard policy iteration is discussed in detail. This fixes the build warning: aiyagari.md:689: WARNING: unknown document: 'opt_savings_2' 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
|
📖 Netlify Preview Ready! Preview URL: https://pr-674--sunny-cactus-210e3e.netlify.app (341df07) 📚 Changed Lecture Pages: aiyagari, endogenous_lake |
Updated the "Primitives and operators" section to correctly state that we solve the household problem using value function iteration (VFI), not Howard policy iteration (HPI). Removed the outdated reference to Ch 5 of Dynamic Programming book. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]>
|
📖 Netlify Preview Ready! Preview URL: https://pr-674--sunny-cactus-210e3e.netlify.app (a3ccd27) 📚 Changed Lecture Pages: aiyagari, endogenous_lake |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
This PR significantly improves the performance and code quality of the Aiyagari model lecture by switching from Howard Policy Iteration (HPI) to Value Function Iteration (VFI) as the primary solution method, with HPI moved to an advanced exercise.
Major Changes
Algorithm Changes
jax.lax.while_loopwith@jax.jitcompilation for optimal performanceParameter Optimization
Code Quality Improvements
@jax.jiton top-level functions (not helper functions)loop_stateinstead ofstateto avoid confusion with DP state variableslax.while_loopPerformance Improvements
JIT Compilation Analysis
We tested four different JIT compilation strategies for HPI:
Winner: Only JIT-compile the top-level function and let JAX optimize the entire call graph.
Educational Improvements
Testing
Files Changed
lectures/aiyagari.md: 216 insertions(+), 103 deletions(-)🤖 Generated with Claude Code
Co-Authored-By: Claude [email protected]