Jax SGMC

JaxSGMC brings Stochastic Gradient Markov chain Monte Carlo (SGMCMC) samplers to JAX. Inspired by optax, JaxSGMC is built on a modular concept to increase reusability and accelerate research of new SGMCMC solvers. Additionally, JaxSGMC aims to promote probabilistic machine learning by removing obstacles in switching from stochastic optimizers to SGMCMC samplers.

To get started quickly using SGMCMC samplers, JaxSGMC provides some popular pre-built samplers in jax_sgmc.alias:

• SGLD (rms-prop)

• SGHMC

• reSGLD

• SGGMC

• AMAGOLD

• OBABO

Getting Started

• Installation
  • Basic Setup
  • Additional Packages
  • Installation from Source
• Quickstart
  • Data Generation
  • Data Loader
  • (Log-)Likelihood and (Log-)Prior
  • Solvers in alias.py
  • Comparison with NumPyro

Reference Documentation

• Data Loading
  • Steps for Setting up a Data Chain
  • Important Notes
  • Numpy Data Loader
  • Tensorflow Data Loader
• Compute Potential from Likelihood
  • Setup DataLoaders
  • Unbatched Likelihood
  • Batched Likelihood
  • Likelihoods with States
• Saving of Samples
  • Seting up Saving
  • Extending Saveable PyTree Types
• Setup Schedulers
  • Specific Schedulers
  • Basic Scheduler
• Setup Custom Solver
  • Overview
  • Load Reference Data
  • Transform Log-likelihood to Potential
  • RMSprop Adaption
  • Integrator and Solver
  • Scheduler
  • Save samples in numpy Arrays
  • Run Solver
  • Plot Results
  • Large Models: Save Data to HDF5

Advanced Topics

• Extend Adapted Quantities
  • Extension of Adaption Strategies
  • Extension of Quantities
• Extending Schedulers
  • Integration with Base Scheduler

Examples

• Image Classification on CIFAR-10

API Documentation

• jax_sgmc
  • jax_sgmc.adaption
    • Adaption Strategies
    • Developer Information
  • jax_sgmc.alias
    • Solvers
  • jax_sgmc.data
    • jax_sgmc.data.core
    • jax_sgmc.data.numpy_loader
    • jax_sgmc.data.tensorflow_loader
    • jax_sgmc.data.hdf5_loader
  • jax_sgmc.integrator
    • Overview
    • Integrators
    • Integrator States
    • Utility
  • jax_sgmc.io
    • Data Collectors
    • Saving Strategies
    • Pytree to Dict Transformation
  • jax_sgmc.potential
    • Stochastic Potential
    • Full Potential
  • jax_sgmc.scheduler
    • Base Scheduler
    • Specific Schedulers
  • jax_sgmc.solver
    • MCMC
    • Solvers
    • Solver States
  • jax_sgmc.util
    • jax_sgmc.util.Array
    • jax_sgmc.util.tree_multiply
    • jax_sgmc.util.tree_add
    • jax_sgmc.util.tree_scale
    • jax_sgmc.util.list_vmap
    • jax_sgmc.util.list_pmap
    • jax_sgmc.util.pytree_list_to_leaves
    • jax_sgmc.util.pytree_leaves_to_list

Indices and tables

• Index

• Module Index

• Search Page