Diffusions¶
-
sgmcmcjax.diffusions.badodab(dt, a: float = 0.01) → Tuple[Callable, Tuple[Callable, Callable], Callable][source]¶ Splitting scheme for the 3-equation Langevin diffusion. See https://arxiv.org/abs/1505.06889 This is a more stable discretisation than SG-NHT
- Parameters
dt (float) – step size
a (float, optional) – initial value of alpha. Defaults to 0.01.
- Returns
An (init_fun, update_fun, get_params) triple.
- Return type
Tuple[Callable, Tuple[Callable, Callable], Callable]
-
sgmcmcjax.diffusions.baoab(dt, gamma: float, tau: float = 1.0) → Tuple[Callable, Tuple[Callable, Callable], Callable][source]¶ BAOAB splitting scheme for the underdampled Langevin diffusion. https://aip.scitation.org/doi/abs/10.1063/1.4802990
- Parameters
dt (float) – step size
gamma (float) – friction coefficient
tau (float, optional) – temperature. Defaults to 1.
- Returns
An (init_fun, (update1, update2), get_params) triple.
- Return type
Tuple[Callable, Tuple[Callable, Callable], Callable]
-
sgmcmcjax.diffusions.psgld(dt, alpha: float = 0.99, eps: float = 1e-05) → Tuple[Callable, Callable, Callable][source]¶ Preconditioned SGLD diffusion See algorithm 1 in paper: https://arxiv.org/pdf/1512.07666.pdf
- Parameters
dt ([type]) – step size
alpha (float, optional) – decay weights for gradients. Defaults to 0.99.
eps ([type], optional) – controls extreme in curvature. Defaults to 1e-5.
- Returns
An (init_fun, update_fun, get_params) triple.
- Return type
Tuple[Callable, Callable, Callable]
-
sgmcmcjax.diffusions.sghmc(dt, alpha: float = 0.01, beta: float = 0) → Tuple[Callable, Callable, Callable, Callable][source]¶ diffusion for stochastic gradient HMC. See paper: https://arxiv.org/abs/1402.4102. Uses the parametrisation in section G (appendix)
- Parameters
dt (float) – step size
alpha (float, optional) – friction coefficient. Defaults to 0.01.
beta (float, optional) – estimation of the stochastic gradient noise. Defaults to 0.
- Returns
An (init_fun, update_fun, get_params, resample_momentum) triple.
- Return type
Tuple[Callable, Callable, Callable, Callable]
-
sgmcmcjax.diffusions.sgld(dt) → Tuple[Callable, Callable, Callable][source]¶ SGLD diffusion https://www.ics.uci.edu/~welling/publications/papers/stoclangevin_v6.pdf
This is an Euler-Maruyam solver for an overdamped Langevin diffusion
- Parameters
dt (float) – step size
- Returns
An (init_fun, update_fun, get_params) triple.
- Return type
Tuple[Callable, Callable, Callable]
-
sgmcmcjax.diffusions.sgldAdam(dt, beta1: float = 0.9, beta2: float = 0.999, eps: float = 1e-08) → Tuple[Callable, Callable, Callable][source]¶ ‘ADAM’-like SGMCMC diffusion. See appendix in paper: https://arxiv.org/abs/2105.13059v1
- Parameters
dt (float) – step size
beta1 (float, optional) – weights for the first moment of the gradients. Defaults to 0.9.
beta2 (float, optional) – weights for the second moment of the gradients. Defaults to 0.999.
eps (float, optional) – small value to avoid instabilities. Defaults to 1e-8.
- Returns
An (init_fun, update_fun, get_params) triple.
- Return type
Tuple[Callable, Callable, Callable]
-
sgmcmcjax.diffusions.sgnht(dt, a: float = 0.01) → Tuple[Callable, Callable, Callable][source]¶ Euler solver for the SG-NHT diffusion See algorithm 2 in http://people.ee.duke.edu/~lcarin/sgnht-4.pdf
- Parameters
dt (float) – step size
a (float, optional) – diffusion factor. Defaults to 0.01.
- Returns
An (init_fun, update_fun, get_params) triple.
- Return type
Tuple[Callable, Callable, Callable]
-
sgmcmcjax.diffusions.welling_teh_schedule(a: float, b: float, gamma: float = 0.55) → Callable[source]¶ Polynomial schedule from https://www.ics.uci.edu/~welling/publications/papers/stoclangevin_v6.pdf
-
sgmcmcjax.diffusion_util.diffusion_factory(is_palindrome: Optional[bool] = False, is_sghmc: Optional[bool] = False) → Callable[source]¶ Diffusion factory Returns a decorator to make a diffusion factory work on PyTrees
- Parameters
is_palindrome (Optional[bool], optional) – whether or not the diffusion is a palidrome method (ie: splitting scheme) Defaults to False.
is_sghmc (Optional[bool], optional) – whether or not the sampler will include momentum resampling. Defaults to False.
- Returns
decorator to apply to a diffusion factory
- Return type
Callable