Skip to content

Fix jitter_rvs handling in find_MAP - #791

Closed
raashish1601 wants to merge 1 commit into
pymc-devs:mainfrom
raashish1601:fix/687-find-map-jitter
Closed

raashish1601 wants to merge 1 commit into
pymc-devs:mainfrom
raashish1601:fix/687-find-map-jitter

Conversation

@raashish1601

Copy link
Copy Markdown

Fixes #687

_make_initial_point had two problems with jitter_rvs:

  • None was turned into an empty list, so nothing was jittered, while the find_MAP and fit_laplace docstrings say None jitters all variables.
  • With freeze_model=True (the default), find_MAP works on a frozen copy of the model, but the variables in jitter_rvs come from the original model. None of them matched, so an explicit list was ignored too.

Now None jitters all free RVs, and an explicit list is matched against the model's free RVs by name. jitter_rvs=[] still means no jitter.

With the example from the issue, all three calls now end at w * z ≈ 0.99 instead of staying at the w = z = 0 stationary point.

Tests: added test_find_MAP_jitter_rvs (None / all / one variable, frozen and not frozen) and test_find_MAP_no_jitter. Four of the new cases fail on main. The rest of tests/inference/laplace_approx passes locally (apart from two JAX-backend cases, since jax isn't installed here), and ruff is clean.

@ricardoV94

Copy link
Copy Markdown
Member

Work on find_map should move to pymc after pymc-devs/pymc#8429, the one here will just redirect/be deprecated.

@raashish1601

Copy link
Copy Markdown
Author

Thanks, makes sense. Should I close this one (and #792, which also touches find_MAP) now, or leave them until pymc-devs/pymc#8429 is merged?

@raashish1601

Copy link
Copy Markdown
Author

pymc-devs/pymc#8429 is merged now. Its find_MAP takes jitter=True/False for all optimized variables, so the jitter_rvs problem from #687 doesn't carry over there. I'll leave this and #792 for you to merge or close, depending on when pymc-extras switches to the pymc version.

@ricardoV94 ricardoV94 closed this Oct 8, 2026
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.

find_MAP: jitter has multiple issues

2 participants