jax_sgmc.scheduler.random_thinning

jax_sgmc.scheduler.random_thinning(step_size_schedule, burn_in_schedule, selections, key=None)

Random thinning weighted by the step size.

Randomly select samples not subject to burn in. The probability of selection is proportional to the step size to deal with the issue of the decaying step size. This only works for static step size and burn in schedules.

Parameters:

• step_size_schedule (specific_scheduler) – Static step size schedule

• burn_in_schedule (specific_scheduler) – Static burn in schedule

• selections (int) – Number of selected samples

• key (Optional[Array]) – PRNGKey for drawing selections

Return type:

specific_scheduler

Returns:

Returns a scheduler marking the accepted samples.