How to load params in single machine.

[matrix97317]

Hi，everyone:

I have saved multi-process parameters in a distributed environment. such as:

Now, I hope to load these parameters in a single machine (non-distributed environment). How to do this?

[matrix97317]

@cpgaffney1 @selamw1

[selamw1]

Hi @matrix97317

To load multi-process parameters onto a single machine, you can use Orbax's PyTreeCheckpointer. The key is to provide an "abstract target" during restoration. This template, created with jax.eval_shape, instructs Orbax to reconstruct the distributed array chunks into complete, single-device arrays.
See this gist as a reference example demonstrating this.