ITADN

Expose model logp and gradlogp

#2920Openmarcobonici 创建于 2026-05-24
M
marcobonicicommented
Reactant/ProbProg currently makes it possible to define a probabilistic model and run HMC/NUTS through `ProbProg.mcmc`. I tried and it seems to work fine (and performance is amazing once you but there does not appear to be a first-class public API for evaluating the model log density and its gradient at an arbitrary parameter vector. It would be useful to expose an API like: ```julia logp, gradlogp = ProbProg.logp_and_grad(model, θ, data...; constraints/observations...) ``` or a compiled callable: ```julia f = ProbProg.compile_logp_and_grad(model, observations, args...) f(θ) # -> (logp, gradlogp) ``` This would make ProbProg-defined models usable with external tools such as `Optim.jl`, `Optimization.jl`, and `Pathfinder.jl`, without requiring users to duplicate the model as a separate hand-written log-density function. ## Current situation The current public API provides nearby pieces, but not the desired abstraction: - `ProbProg.generate` / `generate_` can evaluate the model weight for constrained choices. This weight is effectively the model-defined log density. - `ProbProg.mcmc` / `mcmc_logpdf` internally compute gradients for HMC/NUTS and return the final gradient in `MCMCState.gradient` and `MCMCState.potential_energy`. - `ProbProg.mcmc_logpdf` supports custom log-density functions, but this requires writing an external logp separate from the ProbProg model. The missing piece is a direct, reusable evaluator for a ProbProg model: ```julia θ -> logp θ -> (logp, gradlogp) ``` where `logp` is derived from the model's own `ProbProg.sample(...; logpdf=...)` sites. External inference and optimization tools commonly require a callable log density and gradient. Examples include: ```julia Optim.optimize(f, g!, θ0, Optim.LBFGS()) ``` and Pathfinder-style workflows where one wants to initialize MCMC chains from an approximate posterior: ```julia Pathfinder.multipathfinder(optfun, ndraws; init=...) ``` The former is something we use to compute MLE/MAP (we used that with a student of mine [here](https://arxiv.org/abs/2508.11811), we are planning to do that again), the latter is something I use to get good initial guesses for my chains. At the moment, users who want this workflow appear to need to either: 1. write a separate hand-coded log posterior, duplicating the ProbProg model, or 2. manually build a wrapper around `ProbProg.generate` and differentiate it with Enzyme. Option 1 is error-prone and defeats the purpose of having one model definition. Option 2 works experimentally (I made that work, but it seems fragile, at least in my implementation!), but it is not a documented or ergonomic public API. **Caveat** : this might be also a poor understanding of the codebase on my side. cc @wsmoses
0 条评论