Skip to content

Conversation

@jstac
Copy link
Contributor

@jstac jstac commented Dec 2, 2025

Summary

This PR fixes several bugs in the ifp_advanced lecture that were preventing the code from running:

  • Fixed typo: jnp.block_until_ready(a_start)a_star.block_until_ready() (incorrect variable name)
  • Added missing from functools import partial import needed for JAX JIT static arguments
  • Fixed compute_asset_stationary JIT compilation by adding @partial(jax.jit, static_argnums=(3, 4, 5)) to handle dynamic shape parameters
  • Removed invalid verbose parameter from solve_model calls
  • Refactored solve_model to use jax.lax.while_loop for better performance and proper JIT compilation
  • Cleaned up code structure with improved function signatures and inline comments

Testing

Successfully ran the lecture through jupytext conversion and Python execution:

  • All code cells execute without errors
  • Gini coefficients computed correctly
  • Model solving and simulation complete in ~15 seconds

🤖 Generated with Claude Code

- Fix typo: jnp.block_until_ready(a_start) → a_star.block_until_ready()
- Add missing functools.partial import for JAX JIT static arguments
- Fix compute_asset_stationary JIT compilation with static_argnums
- Remove invalid verbose parameter from solve_model calls
- Refactor solve_model to use jax.lax.while_loop for better performance
- Clean up code structure and improve function signatures

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
@github-actions
Copy link

github-actions bot commented Dec 2, 2025

📖 Netlify Preview Ready!

Preview URL: https://pr-757--sunny-cactus-210e3e.netlify.app (8a1b015)

📚 Changed Lecture Pages: ifp_advanced

@jstac
Copy link
Contributor Author

jstac commented Dec 2, 2025

The lecture now runs successfully from start to finish. Key fixes:

  1. Variable typo: Fixed a_starta_star in the block_until_ready() call
  2. JAX JIT compilation: Added proper static argument handling for compute_asset_stationary to support dynamic num_households, T, and seed parameters
  3. Function signature cleanup: Refactored solve_model to use jax.lax.while_loop instead of Python loops for better JIT performance
  4. Import fix: Added functools.partial import for the @partial decorator

All code executes cleanly with expected output. The model solves and simulates successfully, producing Gini coefficients for different parameter values.

@jstac jstac merged commit 60235eb into main Dec 2, 2025
1 check passed
@jstac jstac deleted the fixadv branch December 2, 2025 11:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants