Issues (or typos?) when running JAX code with multiple GPUs about ferminet HOT 6 CLOSED

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.

from ferminet.

That does make sense - might be just a small change from the internal version that we never tested. We'd be happy to look at a PR.

from ferminet.

Sorry, these crept in when removing a small internal dependency and as a typo respectively. Happy to take a look at a PR, equally happy to release an updated version early next week with fixes.

from ferminet.

Got it. Will work on a fix PR then

from ferminet.

By the way, for case 2 (the logging one)，I got confused by the error message. It should be the case that constants.pmap always returns an array (of the same values in this case) whose size is equal to num_devices. It's just that when the array is of size 1, then logging can handle it automatically while if the size the greater than 1, it would complain. So I simply take the first element in my PR.

from ferminet.

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

from ferminet.