Expose model logp and gradlogp
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 条评论