Fix MultiProcessGpuTest test
Fix MultiProcessGpuTest test

Since MultiProcessGpuTest was using 'shell=True', only the first element
of the args was executed (i.e. python). Therefore the spawn processes
never executed jax code.
Fix the test and make sure 'jax.distributed' initialize by checking
jax.device_count().

google/jax,google/jax,google/jax,google/jax