Fix: Remove @jax.jit decorator from excess demand function to resolve shape errors #770
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
- Upgrade anaconda to 2025.12 - Upgrade quantecon-book-theme to 0.15.1 - Add jupyter-book version constraint (<2.0) - Configure Dependabot for conda ecosystem with jupyter-book restriction - Temporarily disable build cache for full execution check
The @jax.jit decorator on e() causes shape mismatch errors when the function is called with different dimensional inputs (2D vs 3D) throughout the notebook. Removing the decorator allows JAX to handle the function more flexibly. The newton() solver function still uses @jax.jit internally, so performance is maintained.
Contributor
Author
|
Closing as looked more closely at the issue and incorporated a fix in #768 |
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.
Issue
The
e()excess demand function had a@jax.jitdecorator that caused shape mismatch errors when called with different dimensional inputs throughout the notebook:This resulted in:
TypeError: dot_general requires contracting dimensions to have the same shape, got (3,) and (3000,)Solution
Removed the
@jax.jitdecorator from thee()function definition at line 556 in newton_method.md.The
newton()solver function still uses@jax.jitinternally for its update functionq(), so performance is maintained while allowing JAX to handle varying input dimensions flexibly.Testing
This fix allows the notebook to execute successfully with the upgraded anaconda environment (2025.12).
Related