Implementation of SAM as a jax/optax GradientTransformation, with additional adaptive and periodic extensions. This codebase does not presently implement the layer-wise extensions specified in the latter report for large-batch training. See demo.py for a worked example of how to use the interface.