jax.process_index() returns the integer index of the current JAX process. In a multi-host program, the same Python program runs independently on every host, so the index identifies which process is executing the code. section 7.4
Initialize the distributed cluster before performing JAX computations, then use the process index to specialize work or designate one process for coordination. A single-process TPU setup reports index 0. section 7.4 section 3.2.5
Identify the worker and vary its random seed
jax.distributed.initialize()\nprint("Running worker:", jax.process_index())\nrng_key = random.PRNGKey(42 + 10 * jax.process_index())\n\nif jax.process_index() == 0:\n print("Global device count:", jax.device_count()) section 7.4
Listing 7.30 Program for TPU Pod slice
import jax
import jax.numpy as jnp
from jax import random
jax.distributed.initialize() #1
print('== Running worker: ', jax.process_index()) #2
def dot(v1, v2):
return jnp.vdot(v1, v2)
rng_key = random.PRNGKey(42 + 10*jax.process_index()) #3
vs = random.normal(rng_key, shape=(2_000_000,3)) #3
v1s = vs[:1_000_000,:] #3
v2s = vs[1_000_000:,:] #3
# The total number of TPU cores in the Pod
device_count = jax.device_count() #4
# The number of TPU cores attached to this host
local_device_count = jax.local_device_count() #5
if jax.process_index() == 0: #6
print('-- global device count:', jax.device_count())
#print('global devices:', jax.devices())
print('-- local device count:', jax.local_device_count())
#print('local devices:', jax.local_devices())
print('-- JAX version:', jax.__version__)
v1sp = v1s.reshape(
(local_device_count,
v1s.shape[0]//local_device_count,
v1s.shape[1]))
v2sp = v2s.reshape(
(local_device_count,
v2s.shape[0]//local_device_count,
v2s.shape[1]))
if jax.process_index() == 0:
print('-- v1sp shape: ', v1sp.shape) # (8, 125000, 3)
dots = jax.pmap(jax.vmap(dot))(v1sp,v2sp) #7
if jax.process_index() == 0:
print('-- dots shape: ', dots.shape) # (8, 125000)
global_sum = jax.pmap( #8
lambda x: jax.lax.psum(jnp.sum(x), axis_name='p'), #9
axis_name='p'
)(dots)
if jax.process_index() == 0:
print('-- global_sum shape: ', global_sum.shape) # (8,)
print(f'== Worker {jax.process_index()} global sum: {global_sum}') #10
dots = dots.reshape((dots.shape[0]*dots.shape[1]))
if jax.process_index() == 0:
print('-- result shape: ', dots.shape) # (1000000,)
local_sum = jnp.sum(dots) #11
print(f'== Worker {jax.process_index()} local sum: {local_sum}')
print(f'== Worker {jax.process_index()} done')
Shows process_index() in a complete TPU Pod program, alongside global and local device counts, per-process data, and collective computation. section 7.4
