Need a value that reflects every element processed by a mapped axis? jax.lax.psum performs an all-reduce sum across that axis. Give vmap a name such as "batch", then use the same name inside psum. section 6.2.5
Normalize values across a batch
arr = jnp.array(range(50)) norm = jax.vmap( lambda x: x / jax.lax.psum(x, axis_name="batch"), axis_name="batch" ) result = norm(arr) Each mapped function receives one element, while psum supplies the batch-wide total. Dividing by it produces values whose sum is 1.0. section 6.2.5
The same pattern scales beyond one process. With pmap, each device first computes a local sum with jnp.sum; psum then exchanges those sums across hosts and returns the global sum to every worker. section 7.4
Listing 6.17 Using collective ops and axis_name
arr = jnp.array(range(50)) #1
arr
>>> Array([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
>>> 11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
>>> 21, 22, 23, 24, 25, 26, 27, 28, 29, 30,
>>> 31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
>>> 41, 42, 43, 44, 45, 46, 47, 48, 49],
>>> dtype=int32)
norm = jax.vmap(
lambda x: x/jax.lax.psum(x, axis_name='batch'), #2
axis_name='batch') #3
norm(arr) #4
>>> Array([0. , 0.00081633, 0.00163265, 0.00244898, 0.00326531,
>>> 0.00408163, 0.00489796, 0.00571429, 0.00653061, 0.00734694,
>>> 0.00816326, 0.00897959, 0.00979592, 0.01061224, 0.01142857,
>>> 0.0122449 , 0.01306122, 0.01387755, 0.01469388, 0.0155102 ,
>>> 0.01632653, 0.01714286, 0.01795918, 0.01877551, 0.01959184,
>>> 0.02040816, 0.02122449, 0.02204082, 0.02285714, 0.02367347,
>>> 0.0244898 , 0.02530612, 0.02612245, 0.02693878, 0.0277551 ,
>>> 0.02857143, 0.02938776, 0.03020408, 0.03102041, 0.03183673,
>>> 0.03265306, 0.03346939, 0.03428571, 0.03510204, 0.03591837,
>>> 0.03673469, 0.03755102, 0.03836735, 0.03918367, 0.04 ],
>>> dtype=float32)
jnp.sum(norm(arr)) #5
>>> Array(1., dtype=float32)
#1 Generates an array for illustration
#2 Uses psum() op with axis_name=’batch’
#3 Uses axis_name='batch' in vmap()
#4 Applies our normalization function
#5 Checks normalized values
Runnable listing showing axis_name, psum inside vmap, and the normalized output. section 6.2.5
