Issues (or typos?) when running JAX code with multiple GPUs

Hi there. Got two issues when running the JAX code with multiple GPUs:

1. https://github.com/deepmind/ferminet/blob/b46077f1a4687b3ac03b583d9726f59ab4d914d9/ferminet/train.py#L293-L297 It would hit too many values to unpack error when num_devices is greater than 1. My understanding is that we should do

key, *subkeys = jax.random.split(key, num_devices+1)

instead (note the extra asterisk), in which case the following explicit broadcast is not necessary any more for single GPU case.

1. https://github.com/deepmind/ferminet/blob/b46077f1a4687b3ac03b583d9726f59ab4d914d9/ferminet/train.py#L372-L373 constants.pmap gives a tuple of an array instead of just an array in this case when num_devices is greater than 1 (not sure why, probably just JAX’s API). This would cause logging to complain. It’s easy to fix though.

Let me know if it makes sense. Also if you like, I can submit a tiny PR to fix them

[jsspencer]

That’s the correct thing to do. pmean is a collective reduce (ie like MPI_AllReduce instead of MPI_Reduce, if you’re familiar with MPI) and the result is sharded all devices. The logging call requires the data on the host, which transfers the array back, resulting in an array of length num_devices, with elements, as you say, indentical due to the pmean. The same thing is done in several places during the main training loop.

[jsspencer]

Fixed in #17. Thanks for spotting this and sending the patch!