When running under MPI we serialise only the sampler_state on the first rank, and ignore the one on the other ranks.
This is wrong.
When you load it, you get the same seed and chains on every rank (which is wrong!)

I want to get netket working with proper checkpointing, so getting serialisation to work correctly is necessary.

This seems to be what needs to be done to support proper full serialisation of variational states and the stream of samples they generate.

To avoid raising errors, if you attempt to load a saved state there are three cases:

• If you have same number of ranks and of chains, then just load it and all is good
• If you have different number of ranks , but same number of total chains, then just load the chains and use some random seed (we cannot restore perfectly, but at least the chains are already thermalised)
• If you have different number of chains, then simply load the weights and throw a warning.

I'd like some opinions..