Refactor: Add JAX Hessian computation and correct Fitter.solve #2
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 builds upon the JAX integration by adding JAX-based Hessian computation for both MFAModel and InstMFAModel. It also includes the correction to Fitter.solve() from the previous iteration. Key changes: - MFAModel and InstMFAModel now compute the Hessian using `jax.hessian` on their respective core JAX objective functions when the JAX pathway is active. These are JIT-compiled. - The JAX-computed Hessian is provided to SciPy optimizers if the solver is not 'slsqp' or 'ralg'. - Corrected Fitter.solve() in fit.py to properly pass JAX usage flags. Existing functionality: - JAX-based objective and gradient for MFAModel (steady-state). - Structural JAX support for InstMFAModel (instationary), with the core instationary JAX calculation (`core_calculate_inst_mdvs_jax`) still being a non-functional placeholder. - Data preparation in Calculator and Fitter/InstFitter for JAX pathways. Testing: - Test script execution remains blocked by environment issues. Full validation of JAX steady-state path (including Hessian) and instationary path structure could not be completed. Further work: - Implement `core_calculate_inst_mdvs_jax` for instationary models. - Thoroughly test and benchmark all JAX pathways.
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.
This commit builds upon the JAX integration by adding JAX-based Hessian computation for both MFAModel and InstMFAModel. It also includes the correction to Fitter.solve() from the previous iteration.
Key changes:
jax.hessianon their respective core JAX objective functions when the JAX pathway is active. These are JIT-compiled.Existing functionality:
core_calculate_inst_mdvs_jax) still being a non-functional placeholder.Testing:
Further work:
core_calculate_inst_mdvs_jaxfor instationary models.