Improve integration with external samplers#2203
Conversation
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
juanitorduz
changed the title
|
Thank you for the feedback @fehiepsi . Let me know how this iteration looks. I am happy to keep iterating until we get the best level of abstraction :) |
|
Would people be happy merging this? Certainly it looks good from my end |
|
I would 😄! |
|
Let's wait a bit to rethink whether the design choices are reasonable. Looking like we are changing other internal codes due to the change. It indicates that we will likely break current user code. Let's see if we can come up with simpler solutions. |
|
Great! Always open for feedback 🤗 |
External samplers v2
|
@fehiepsi you were right, there was no need to add such a large change to the internal API 🙈 . In a463116, I simplified the approach and minimized the changes. Now this PR is way less invasive 🤗 fyi @reubenharry |
|
Yeah, this seems pleasantly concise |
|
@fehiepsi, what do you think about the current status? Any additional feedback or concerns? :) |
Adopts master's ty-only type-checking toolchain (mypy removed, ty.toml deleted, ty config consolidated in pyproject.toml). Resolved the lone pyproject.toml conflict by keeping master's [tool.ty.*] config and adding numpyro/infer/util.py to [tool.ty.src].include. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Per fehiepsi's review on PR pyro-ppl#2203, drop the get_log_density_fn helper, its LogDensityInfo return type, and the PositionDict alias; build the log density inline by negating model_info.potential_fn (initialize_model already returns a single-position potential_fn/postprocess_fn). - numpyro/infer/util.py: remove get_log_density_fn + LogDensityInfo; narrow w.message to Warning in find_valid_initial_params so util.py type-checks under ty (newly added to [tool.ty.src]). - numpyro/_typing.py, numpyro/infer/__init__.py, docs/source/utilities.rst: drop the removed symbols. - notebooks/source/other_samplers.ipynb: inline the pattern and re-execute. - test/infer/test_external_helpers.py: inline the pattern; keep constrain_fn and end-to-end MCMCKernel coverage. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Reword the markdown lead-in to the Pathfinder log-density cell to say we build the log-density function by negating the potential energy, avoiding the awkward doubled 'negated' phrasing. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Address fehiepsi review: remove get_log_density_fn, inline via initialize_model
|
Thanks @fehiepsi I simplified it as suggested and resolved the merge conflicts (sorry for the commit mess, but we will squash it anyway :) ) |
|
Looks great, thanks! |
3 participants
Motivated by #2124 and specifically by #2124 (comment), we suggest a cleaner API to integrate with other samplers. This includes MCLMC from Blackjax. This is illustrated in the notebook.
(most from the code changes come from the notebook)