metric_gaussian_kl.py: Simplify MPI sample prop

Use the same logic for returning the samples from the various MPI tasks
as for working with them in each task.
9 jobs for KL_sample_generator in 17 minutes and 4 seconds