jax

Fix test_gpu_distributed_initialize() to get a more proper GPUs number

#503

Merged