Skip to content

Address fehiepsi review: remove get_log_density_fn, inline via initialize_model#2

Merged
juanitorduz merged 4 commits into
external-samplers-master-mergefrom
external-samplers-feedback
Jun 30, 2026
Merged

Address fehiepsi review: remove get_log_density_fn, inline via initialize_model#2
juanitorduz merged 4 commits into
external-samplers-master-mergefrom
external-samplers-feedback

Conversation

@juanitorduz

@juanitorduz juanitorduz commented Jun 30, 2026

Copy link
Copy Markdown
Owner

Addresses fehiepsi's two latest comments on pyro-ppl#2203. Scoped so the diff shows only the review-feedback changes (6 files).

Base branch

Base is `external-samplers-master-merge` = `external-samplers` + a routine merge of `origin/master` (which moved the repo to the ty-only toolchain: mypy removed, `ty.toml` deleted, ty config in `pyproject.toml`; the lone `pyproject.toml` conflict was resolved by keeping master's `[tool.ty.*]` and adding `numpyro/infer/util.py` to `[tool.ty.src].include`). Keeping that merge out of this diff is intentional, so reviewers see just the feedback changes here.

Addressing the review

Both comments asked to drop a thin abstraction; done via full removal + inline:

  • Removed `get_log_density_fn`, its `LogDensityInfo` return type, and the `PositionDict` alias.
  • The log density is now built inline by negating `model_info.potential_fn` (`initialize_model` already returns a single-position `potential_fn`/`postprocess_fn` with `dynamic_args=False`), per "we can add 2 lines of code in the notebook instead."
  • Updated `numpyro/infer/init.py`, `docs/source/utilities.rst`, the notebook prose, and the tests. Kept the `constrain_fn` (`batch_ndims`) work and the end-to-end `MCMCKernel` coverage.
  • `numpyro/infer/util.py` is now type-checked by ty; this surfaced 3 pre-existing diagnostics in `find_valid_initial_params` (`w.message` typed `Warning | str`), fixed with a one-line `assert isinstance(w.message, Warning)` narrowing.

Verification

  • `make lint` green (ruff check, ruff format, header check, `ty check`).
  • `prek run --all-files` all hooks pass.
  • `pytest test/infer/test_external_helpers.py test/infer/test_infer_util.py` → 61 passed.
  • `notebooks/source/other_samplers.ipynb` fully re-executed (Pathfinder + MCLMC + flowMC), 0 errors, outputs refreshed.

🤖 Generated with Claude Code

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>
@juanitorduz
juanitorduz changed the base branch from external-samplers to external-samplers-master-merge June 30, 2026 21:43
@juanitorduz juanitorduz changed the title Address fehiepsi review + merge master (ty-only toolchain) Address fehiepsi review: remove get_log_density_fn, inline via initialize_model Jun 30, 2026
juanitorduz and others added 3 commits June 30, 2026 23:46
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>
@juanitorduz
juanitorduz merged commit 21ebd13 into external-samplers-master-merge Jun 30, 2026
@juanitorduz
juanitorduz deleted the external-samplers-feedback branch June 30, 2026 21:54
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.

1 participant