Skip to content

Improve serialization of variational states

Vicentini Filippo requested to merge pv/serialize_vstate into master

Created by: PhilipVinc

This PR addressees two issues in the serialization of variational states:

  • The chunk_size was not serialised
  • When serialising a state that was sampled, the 'current' sampler_state was saved. This lead to a situation where it Wass impossible to regenerate the old samples, because the previous PRNGKey was lost. This PR now keeps around the old sampler state so that it can be serialised.

Both, but in particular the latter, are necessary to actually implement checkpointing.

Merge request reports