quantax.utils.to_replicated_numpy

quantax.utils.to_replicated_numpy(array: Array) → ndarray

Gather a (possibly distributed) array into a contiguous numpy array holding the full data, identical on every process.

In multi-host jobs the array is first replicated and then brought to the host with jax.experimental.multihost_utils.global_array_to_host_local_array().