Make fit_pathfinder more similar to fit_laplace and pm.sample #447

velochy · 2025-04-01T19:00:35Z

Add initvals to parameters
Add observations and constants to return value (idata)
Add both fit_laplace and fit_pathfinder to docs and make fit link to both
Clean up init.py include structure

This arose out of https://discourse.pymc.io/t/pymc-extras-fit-compatibility-with-pm-sample/16763/5

zaxtax · 2025-04-08T09:13:10Z

This looks alright to me. Thoughts @ricardoV94 ?

fonnesbeck · 2025-04-08T11:08:00Z

CC @aphc14

ricardoV94 · 2025-04-09T10:45:18Z

pymc_extras/__init__.py

-from pymc_extras.inference.find_map import find_MAP
-from pymc_extras.inference.fit import fit
-from pymc_extras.inference.laplace import fit_laplace
+from pymc_extras.inference import *


Can we use explicit imports?

Yes, of course. I'll fix that when I have a few minutes

ricardoV94 · 2025-04-09T10:46:13Z

pymc_extras/inference/pathfinder/pathfinder.py

        The inference data containing the results of the Pathfinder algorithm.

    References
    ----------
    Zhang, L., Carpenter, B., Gelman, A., & Vehtari, A. (2022). Pathfinder: Parallel quasi-Newton variational inference. Journal of Machine Learning Research, 23(306), 1-49.
    """

+    if initvals is not None:


Shouldn't mutate the model, make a copy perhaps if there's no better way to just forward the initvals

I would normally agree. However, I tried it, but model.copy() does not produce a working model sometimes - most notably when any transformations are used.

Should I use some other copy function?

@velochy you want the clone_model function https://github.com/pymc-devs/pymc/blob/2842401f95de74ab37b7750cff455af28cddaffa/pymc/model/fgraph.py#L374

Thank you for that @zaxtax
Updated the code. Also added a test to check if initvals are used correctly.

velochy · 2025-04-09T12:27:44Z

@ricardoV94 it currently also fails a test on including blackjax, which is currently included in pathfinder.py
How would you like to handle it? Do I bring the include lower down so it is only included when used, or do I add it to the requirements of the package?

aphc14 · 2025-04-09T14:39:37Z

@ricardoV94 it currently also fails a test on including blackjax, which is currently included in pathfinder.py How would you like to handle it? Do I bring the include lower down so it is only included when used, or do I add it to the requirements of the package?

Yup, that is a good idea. closes #443

aphc14 · 2025-04-09T15:04:07Z

pymc_extras/inference/pathfinder/pathfinder.py

        The inference data containing the results of the Pathfinder algorithm.

    References
    ----------
    Zhang, L., Carpenter, B., Gelman, A., & Vehtari, A. (2022). Pathfinder: Parallel quasi-Newton variational inference. Journal of Machine Learning Research, 23(306), 1-49.
    """

+    if initvals is not None:
+        for rv_name, ivals in initvals.items():
+            model.set_initval(model.named_vars[rv_name], ivals)


Might need to ensure that ivals is a support point for the RV. For example, x ~ Uniform(-1, 1) would have nan initial values with model.set_initval(model.named_vars["x"], 2)

While in the ideal world, I would agree, in practice
a) It is very nontrivial to do as I understand, as the limits are not specified anywhere where they are easy to take
b) pm.sample does no such checks, and the goal of this PR is to be compatible with that

Yup, seems fair enough. Thanks for this submission @velochy

zaxtax · 2025-04-09T15:48:49Z

Maybe we make a requirements-test.txt file like arviz has? https://github.com/arviz-devs/arviz/blob/main/requirements-test.txt

…

On Wed, 9 Apr 2025, 14:28 Margus Niitsoo, ***@***.***> wrote: @ricardoV94 <https://github.com/ricardoV94> it currently also fails a test on including blackjax, which is currently included in pathfinder.py How would you like to handle it? Do I bring the include lower down so it is only included when used, or do I add it to the requirements of the package? — Reply to this email directly, view it on GitHub <#447 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUM66WC34I7BYIQZWT32YUG5LAVCNFSM6AAAAAB2HVSQBWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDOOBZGUZTOMRUG4> . You are receiving this because you commented.Message ID: ***@***.***> *velochy* left a comment (pymc-devs/pymc-extras#447) <#447 (comment)> @ricardoV94 <https://github.com/ricardoV94> it currently also fails a test on including blackjax, which is currently included in pathfinder.py How would you like to handle it? Do I bring the include lower down so it is only included when used, or do I add it to the requirements of the package? — Reply to this email directly, view it on GitHub <#447 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUM66WC34I7BYIQZWT32YUG5LAVCNFSM6AAAAAB2HVSQBWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDOOBZGUZTOMRUG4> . You are receiving this because you commented.Message ID: ***@***.***>

velochy · 2025-04-09T16:00:06Z

Yup, that is a good idea. closes #443

Considering you have a standing PR on this, does it not make more sense to just have your quickfix merged separately?
I don't think it conflicts with anything I've done here

zaxtax · 2025-04-09T16:01:38Z

Definitely!

…

On Wed, 9 Apr 2025, 18:00 Margus Niitsoo, ***@***.***> wrote: Yup, that is a good idea. closes #443 <#443> Considering you have a standing PR on this, does it not make more sense to just have your quickfix merged separately? I don't think it conflicts with anything I've done here — Reply to this email directly, view it on GitHub <#447 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUIEXAGQCBW56OPD6R32YU7ZZAVCNFSM6AAAAAB2HVSQBWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDOOJQGIYDINJWHE> . You are receiving this because you commented.Message ID: ***@***.***> *velochy* left a comment (pymc-devs/pymc-extras#447) <#447 (comment)> Yup, that is a good idea. closes #443 <#443> Considering you have a standing PR on this, does it not make more sense to just have your quickfix merged separately? I don't think it conflicts with anything I've done here — Reply to this email directly, view it on GitHub <#447 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUIEXAGQCBW56OPD6R32YU7ZZAVCNFSM6AAAAAB2HVSQBWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDOOJQGIYDINJWHE> . You are receiving this because you commented.Message ID: ***@***.***>

aphc14 · 2025-04-10T14:49:05Z

The latest test failure is due to pathfinder using pt.batched_dot operations. PR #453 should solve this.

velochy · 2025-04-14T15:04:00Z

Is there anything I should still do?
Tests are failing for things unrelated to this PR to the best of my understanding, and I have addressed all of the comments

zaxtax · 2025-04-14T15:19:39Z

Try to rebase your PR. If it still fails for unrelated reasons, that's not on you to resolve

…

On Mon, 14 Apr 2025, 17:04 Margus Niitsoo, ***@***.***> wrote: Is there anything I should still do? Tests are failing for things unrelated to this PR to the best of my understanding, and I have addressed all of the comments — Reply to this email directly, view it on GitHub <#447 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUNBG4YRRB3WV4M534T2ZPE7LAVCNFSM6AAAAAB2HVSQBWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDQMBSGAZDCNBVGI> . You are receiving this because you commented.Message ID: ***@***.***> *velochy* left a comment (pymc-devs/pymc-extras#447) <#447 (comment)> Is there anything I should still do? Tests are failing for things unrelated to this PR to the best of my understanding, and I have addressed all of the comments — Reply to this email directly, view it on GitHub <#447 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUNBG4YRRB3WV4M534T2ZPE7LAVCNFSM6AAAAAB2HVSQBWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDQMBSGAZDCNBVGI> . You are receiving this because you commented.Message ID: ***@***.***>

zaxtax · 2025-04-14T18:51:34Z

pymc_extras/inference/pathfinder/pathfinder.py

-        "identity" : Applies log importance weights directly without resampling.
-        None : No importance sampling weights. Returns raw samples of size (num_paths, num_draws_per_path, N) where N is number of model parameters. Other methods return samples of size (num_draws, N).
+
+        - **"psis"** : Pareto Smoothed Importance Sampling (default). Usually most stable.


Can you remove these ** characters? They aren't consistent the style used elsewhere.

zaxtax · 2025-04-14T18:55:36Z

Can you also move the import of jax into the relevant functions in the pathfinder file?

zaxtax · 2025-04-15T19:06:51Z

I'm not sure why the build source distribution test is still unhappy. Can you investigate why importing jax still happens?

…ue for pathfinder and cleaned relevant docs a bit

velochy · 2025-04-15T20:17:53Z

I'm not sure why the build source distribution test is still unhappy. Can you investigate why importing jax still happens?

I think it was because of the import from pymc.sampling.jax. Moved that down as well. Hopefully tests will pass now

velochy · 2025-04-16T12:55:05Z

What is still needed to get it merged? Review from @jessegrabowski ?

zaxtax · 2025-04-16T14:18:14Z

Just giving others one last chance to chime in before I merge it 😀

…

On Wed, 16 Apr 2025, 14:55 Margus Niitsoo, ***@***.***> wrote: What is still needed to get it merged? Review from @jessegrabowski <https://github.com/jessegrabowski> ? — Reply to this email directly, view it on GitHub <#447 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUMGP6MZSWI6CUJPSID2ZZHMBAVCNFSM6AAAAAB2HVSQBWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDQMBZGUYDINBVGA> . You are receiving this because you were mentioned.Message ID: ***@***.***> *velochy* left a comment (pymc-devs/pymc-extras#447) <#447 (comment)> What is still needed to get it merged? Review from @jessegrabowski <https://github.com/jessegrabowski> ? — Reply to this email directly, view it on GitHub <#447 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUMGP6MZSWI6CUJPSID2ZZHMBAVCNFSM6AAAAAB2HVSQBWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDQMBZGUYDINBVGA> . You are receiving this because you were mentioned.Message ID: ***@***.***>

zaxtax · 2025-04-16T21:34:05Z

@velochy now let's get #444 merged in too 😄

zaxtax requested a review from jessegrabowski April 8, 2025 09:15

ricardoV94 reviewed Apr 9, 2025

View reviewed changes

aphc14 reviewed Apr 9, 2025

View reviewed changes

velochy force-pushed the FitCompatibility branch from 1ee9bcf to 370ffe0 Compare April 9, 2025 16:13

velochy force-pushed the FitCompatibility branch 3 times, most recently from f53445a to b81684d Compare April 14, 2025 16:51

zaxtax reviewed Apr 14, 2025

View reviewed changes

velochy force-pushed the FitCompatibility branch from b81684d to 3416258 Compare April 14, 2025 19:41

Added initvals to parameters, constants and observations to returnval…

ef44c82

…ue for pathfinder and cleaned relevant docs a bit

velochy force-pushed the FitCompatibility branch from 3416258 to ef44c82 Compare April 15, 2025 20:16

zaxtax self-requested a review April 16, 2025 03:43

zaxtax approved these changes Apr 16, 2025

View reviewed changes

zaxtax merged commit 53df690 into pymc-devs:main Apr 16, 2025
5 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Make fit_pathfinder more similar to fit_laplace and pm.sample #447

Make fit_pathfinder more similar to fit_laplace and pm.sample #447

velochy commented Apr 1, 2025 •

edited

Loading

zaxtax commented Apr 8, 2025

fonnesbeck commented Apr 8, 2025

ricardoV94 Apr 9, 2025

velochy Apr 9, 2025

velochy Apr 9, 2025

ricardoV94 Apr 9, 2025

velochy Apr 9, 2025

zaxtax Apr 14, 2025

velochy Apr 14, 2025

velochy commented Apr 9, 2025

aphc14 commented Apr 9, 2025

aphc14 Apr 9, 2025

velochy Apr 9, 2025

aphc14 Apr 10, 2025

zaxtax commented Apr 9, 2025 via email

velochy commented Apr 9, 2025

zaxtax commented Apr 9, 2025 via email

aphc14 commented Apr 10, 2025

velochy commented Apr 14, 2025

zaxtax commented Apr 14, 2025 via email •

edited

Loading

zaxtax Apr 14, 2025

velochy Apr 14, 2025

zaxtax commented Apr 14, 2025

zaxtax commented Apr 15, 2025

velochy commented Apr 15, 2025

velochy commented Apr 16, 2025

zaxtax commented Apr 16, 2025 via email

zaxtax commented Apr 16, 2025 •

edited

Loading

Make fit_pathfinder more similar to fit_laplace and pm.sample #447

Make fit_pathfinder more similar to fit_laplace and pm.sample #447

Conversation

velochy commented Apr 1, 2025 • edited Loading

zaxtax commented Apr 8, 2025

fonnesbeck commented Apr 8, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

velochy commented Apr 9, 2025

aphc14 commented Apr 9, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zaxtax commented Apr 9, 2025 via email

velochy commented Apr 9, 2025

zaxtax commented Apr 9, 2025 via email

aphc14 commented Apr 10, 2025

velochy commented Apr 14, 2025

zaxtax commented Apr 14, 2025 via email • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zaxtax commented Apr 14, 2025

zaxtax commented Apr 15, 2025

velochy commented Apr 15, 2025

velochy commented Apr 16, 2025

zaxtax commented Apr 16, 2025 via email

zaxtax commented Apr 16, 2025 • edited Loading

velochy commented Apr 1, 2025 •

edited

Loading

zaxtax commented Apr 14, 2025 via email •

edited

Loading

zaxtax commented Apr 16, 2025 •

edited

Loading