Use jax.local_device_count() instead of jax.device_count() to get the number of accelerators attached to a host.
Use jax.local_device_count() instead of jax.device_count() to get the number of accelerators attached to a host.

PiperOrigin-RevId: 407302483
Change-Id: I556eab295b6dc9eb56f6ce1700326ba072bba92f

deepmind/ferminet