Update to 0.1.0. This update includes - a planned deprecation of meth:utils.update_step - jit multiple training steps together - full support for pmap across devices (GPU/TPU) of the same model - refactoring to make library simpler, flatter - bug fixes and minor improvements