Repository navigation
MLX MCLMC sampler - #743
jessegrabowski wants to merge 59 commits into
Conversation
eb91d4c to
9311d33
Compare
| _log = logging.getLogger(__name__) | ||
|
|
||
| # The dynamics, the integrator schedules, and the adaptation are ported from the blackjax | ||
| # isokinetic sampler, so its source is the reference for anything ambiguous here. |
Codecov Report❌ Patch coverage is ❌ Your patch status has failed because the patch coverage (12.31%) is below the target coverage (50.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## main #743 +/- ##
===========================================
+ Coverage 51.60% 85.92% +34.32%
===========================================
Files 73 113 +40
Lines 8003 11063 +3060
===========================================
+ Hits 4130 9506 +5376
+ Misses 3873 1557 -2316
🚀 New features to boost your workflow:
|
|
CC @reubenharry in case you'd like to take a look |
517a8a7 to
539b4d0
Compare
|
Very exciting! A quick comment on this code comment:
This is misleading. The MH step is not what makes standard HMC decorrelate - rather, the resampling of momentum that happens after the MH step is. In fact, HMC with MH can be extended to have Langevin noise. Let me know if there are any specific questions that I can help with! |
Deliberately not registered with fit() or exported from pymc_extras.inference, because importing it requires mlx, which installs only on Apple Silicon.
The previous estimate omitted the -1/(n-1) correction and the trailing rho_hat_even term, which damped the seed-to-seed variation in the adapted L.
Blackjax bumps num_steps2 only for the reported tuning-step total; the adaptation itself runs the unbumped count in both the scan and the readjust.
Blackjax applies handle_nans in the kernel, so it protects sampling as well as adaptation; sample() previously let a nan log-density propagate for the rest of the run.
MLX has no reverse rule for Inverse or Cholesky, so differentiating the funcified logp with mx.value_and_grad failed on any model using them; PyTensor already knows those gradients symbolically.
A finite log-density with a non-finite gradient means the start sits on a boundary or removable singularity; left unchecked it turned the whole run into nan draws with a collapsed step size.
Folding a nan gradient into the Adam moments poisoned every later step, so the ascent silently stopped moving; this holds the whole optimizer state instead, as optax's apply_if_finite does.
The textbook form loses up to 13% in float32 at small delta by subtracting two near-equal log 2s, and the cancellation-free rewrite leans on expm1, which MLX computes less accurately above roughly 0.15.
Hand-written special functions such as gammaln expand the graph enough to exceed the limit, so a model can hit it without doing anything unusual.
…ction Ports nuts-rs's windowed sqrt(Var[x]/Var[grad]) diagonal and its diagonal-plus-low-rank fit; only the low-rank form can precondition a posterior whose ridges are not axis-aligned, and it is rejected in favour of the diagonal whenever the window does not support it.
Freezing dim lengths and data is the caller's choice, not the sampler's: it rules out pm.set_data, and the frozen copy was what everything downstream reported results against.
Read the Docs builds on Linux, where mlx does not install, and autosummary has to import the function to render its page.
…does The trainer reseeded a compiled function by overwriting its RNG input storage with jax_typify hardcoded, which only worked on JAX. pymc's rule replaces it: a linker that copies RNGs at compile time gets a fresh compile per seed, and every other linker reseeds the cached function's shared variables through the graph.
…ery compiled step
|
The finite branch can multiply the cap instead of passing it through unchanged: step_size_max = mx.where(is_finite, state.step_size_max * 1.02, state.step_size * 0.8)Once the cap can relax, the Var[E] controller sets the step size on its own merits again. |
|
The existing guards key on non-finite steps: Carrying the best point through the loop: best_logdensity, best_position = -mx.inf, position
...
improved = usable & (logdensity > best_logdensity)
best_logdensity = mx.where(improved, logdensity, best_logdensity)
best_position = mx.where(improved, position, best_position)
...
final_logdensity = logp_and_grad(position)[0]
position = mx.where(
mx.isfinite(final_logdensity) & (final_logdensity >= best_logdensity), position, best_position
)Where the ascent behaves the last point is the best point and this changes nothing. On the model above the controller starts at -11,213 instead. An update-norm threshold does not substitute for it. Adam normalizes the magnitude out of the update, so over a 200-step ascent a run that collapses and one that does not have p99 update norms of 0.1937 and 0.1890 against |
This branch has not been deployed
4 participants
It's based on the blackjax implementation by @junpenglao . Highly experimental, but promising results out of the gate. I'll update with some results on the pymc-models-catalog soon.