jax_sgmc.alias

Popular solvers ready to use.

While JaxSGMC has been designed to be flexible, starting with full flexibility can be complicated. Therefore, this file contains some popular solvers with preset properties, which can be applied directly to the problem or used as a guide to set up a custom solver.

Solvers

sgld(potential_fn, data_loader[, ...])	Stochastic Gradient Langevin Dynamics.
re_sgld(potential_fn, data_loader[, ...])	Replica Exchange Stochastic Gradient Langevin Diffusion.
amagold(stochastic_potential_fn, ...[, ...])	Amortized Metropolis Adjustment for Efficient Stochastic Gradient MCMC.
sggmc(stochastic_potential_fn, ...[, ...])	Stochastic gradient guided monte carlo.
sghmc(potential_fn, data_loader[, ...])	Stochastic Gradient Hamiltonian Monte Carlo.
obabo(potential_fn, data_loader[, ...])	Langevin Monte Carlo with partial momentum refreshment.