flax

2e1232e1 - fix temperature_sampler + add test