Multi-Path Pathfinder #783
Draft
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 update adds the multi pathfinder capabilities, parallel Pathfinder runs, LBFGS optimisation using Optax, and importance sampling. - Multi-Pathfinder: Supports parallel and vectorized sampling strategies, improving scalability. - LBFGS Optimizer: Refactored using Optax for better numerical stability and performance, with enhanced inverse Hessian estimation. - Importance Sampling: Added support for various methods, improving sampling accuracy. - Testing: Expanded tests for both Pathfinder and LBFGS to ensure correctness and stability.
This was referenced Mar 8, 2025
junpenglao
reviewed
Apr 7, 2025
Comment on lines
+487
to
+496
| def psislw_wrapper(logiw_array): | ||
| def psislw(logiw_array): | ||
| result_logiw, result_k = az.psislw(np.array(logiw_array)) | ||
| return np.array(result_logiw), np.array(result_k) | ||
|
|
||
| return jax.pure_callback( | ||
| psislw, | ||
| (jnp.zeros_like(logiw_array), jnp.zeros((), dtype=jnp.float64)), | ||
| logiw_array, | ||
| ) |
There was a problem hiding this comment.
We should add a pure JAX implementation of psislw (in a seperate PR), as pure_callback will be quite inefficient.
There was a problem hiding this comment.
We can. I have found a JAX implementation here https://gist.github.com/adamhaber/0556671340e0daa9e2c6e3fd535cd992 by @adamhaber as a good starting point
|
Hi @aphc14, could you isolate the change replacing jaxopt with optax in a seperate PR? |
|
Should this PR include the optax code changes while referencing the separate optax PR? Or should this PR revert the changes from optax to jaxopt? |
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.
Implement Multi-Pathfinder and Enhance Pathfinder and LBFGS Optimizers
High-Level Description
Multi-Pathfinder: Introduces the
multi_pathfinderfunction, enabling parallel and vectorized sampling strategies. Initial tests indicate thatmulti_pathfinderis functioning as expected. Feel free to test it out and provide feedback.LBFGS Optimizer: The optimizer has been refactored to use Optax.
single pathfinder fix: Fixed inaccurate calculations of S Z matrices, phi, log densities.
alpha_recover: Decoupling of alpha recover out of LBFGS optimisation.
Importance Sampling: Added support for various importance sampling methods.
Testing: Expanded and updated tests for both Pathfinder and LBFGS.
Current Status
This is a draft PR that requires some tidying up. There are existing linting errors from mypy that need to be addressed. Your feedback and testing are welcome to help refine these changes.
Checklist
maincommit.pre-committo check for any issues.resolves #763, #213, #461, #749, #704
related #465, #387