Add up one vector per device across 8 devices and give every device the total, building the all-reduce yourself from neighbour-to-neighbour sends with jax.lax.ppermute.
Row d of x is a vector of length C that belongs to device d. Build a one-axis mesh named ring over all 8 devices and use jax.shard_map with x split by rows, so each device starts with its own (1, C) row. The goal is the result that jax.lax.psum would give,
resultd=e=0∑7xefor every device d,where xe is row e of x. Cut every row into 8 chunks of C/8 consecutive values and connect the devices in a ring in which every device d sends to device (d + 1) mod 8, or one in which every device d sends to device (d - 1) mod 8. The ring algorithm runs two phases of 7 steps each. In the reduce-scatter phase, every step passes one chunk to the next device, which adds its own copy of that chunk before passing the running sum on, so after 7 steps each device owns one chunk that holds the sum over all 8 devices. In the all-gather phase, the finished chunks travel around the ring for 7 more steps until every device holds all 8 of them.
Return a float32 JAX array of shape (8, C), split by rows so that device d holds row d, in which every row equals the sum of the 8 rows of x.
Input: x = [[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0],[1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0],[2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0],[3.0,3.0,3.0,3.0,3.0,3.0,3.0,3.0],[4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0],[5.0,5.0,5.0,5.0,5.0,5.0,5.0,5.0],[6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0],[7.0,7.0,7.0,7.0,7.0,7.0,7.0,7.0]]
Output: [[28.0,28.0,28.0,28.0,28.0,28.0,28.0,28.0],[28.0,28.0,28.0,28.0,28.0,28.0,28.0,28.0],[28.0,28.0,28.0,28.0,28.0,28.0,28.0,28.0],[28.0,28.0,28.0,28.0,28.0,28.0,28.0,28.0],[28.0,28.0,28.0,28.0,28.0,28.0,28.0,28.0],[28.0,28.0,28.0,28.0,28.0,28.0,28.0,28.0],[28.0,28.0,28.0,28.0,28.0,28.0,28.0,28.0],[28.0,28.0,28.0,28.0,28.0,28.0,28.0,28.0]]
Explanation: Device d holds a row of eight copies of d, so each of its 8 chunks is the single value d. After reduce-scatter each device owns one chunk equal to 0 + 1 + 2 + 3 + 4 + 5 + 6 + 7 = 28, and after all-gather every device holds all 8 finished chunks, so every row is eight copies of 28.
Input: x = [[-2.0,-1.0,0.0,1.0,2.0,-2.0,-1.0,0.0],[1.0,2.0,-2.0,-1.0,0.0,1.0,2.0,-2.0],[-1.0,0.0,1.0,2.0,-2.0,-1.0,0.0,1.0],[2.0,-2.0,-1.0,0.0,1.0,2.0,-2.0,-1.0],[0.0,1.0,2.0,-2.0,-1.0,0.0,1.0,2.0],[-2.0,-1.0,0.0,1.0,2.0,-2.0,-1.0,0.0],[1.0,2.0,-2.0,-1.0,0.0,1.0,2.0,-2.0],[-1.0,0.0,1.0,2.0,-2.0,-1.0,0.0,1.0]]
Output: [[-2.0,1.0,-1.0,2.0,0.0,-2.0,1.0,-1.0],[-2.0,1.0,-1.0,2.0,0.0,-2.0,1.0,-1.0],[-2.0,1.0,-1.0,2.0,0.0,-2.0,1.0,-1.0],[-2.0,1.0,-1.0,2.0,0.0,-2.0,1.0,-1.0],[-2.0,1.0,-1.0,2.0,0.0,-2.0,1.0,-1.0],[-2.0,1.0,-1.0,2.0,0.0,-2.0,1.0,-1.0],[-2.0,1.0,-1.0,2.0,0.0,-2.0,1.0,-1.0],[-2.0,1.0,-1.0,2.0,0.0,-2.0,1.0,-1.0]]
Input: x = [[7.5,-3.5,-3.0,2.0,6.5,0.0,-0.5,-4.5,-3.0,-1.0,-5.0,-8.0,8.0,-4.0,-5.0,-6.5],[-2.5,-3.5,-5.0,-7.5,0.5,1.0,-6.0,4.0,6.0,8.0,5.5,-0.5,3.0,5.0,-1.0,5.5],[8.0,-7.0,4.5,6.5,2.5,-4.0,-5.0,7.0,8.0,0.5,-2.0,7.0,-2.5,-1.5,-3.0,-4.5],[-0.5,-4.0,-3.5,-7.5,-2.0,2.0,5.0,4.0,1.5,3.0,-5.5,-3.0,-5.5,2.0,-3.0,1.0],[-2.5,2.0,8.0,5.5,-4.0,0.0,0.5,7.5,1.5,7.5,-7.0,6.5,-7.5,-2.5,-1.0,-3.0],[-7.5,-1.5,4.5,-3.0,1.5,-0.5,-3.5,1.0,6.5,5.5,3.5,-4.0,0.0,-2.5,2.0,-2.0],[-4.0,4.0,-2.0,-4.0,-1.0,-1.5,-8.0,7.5,2.5,-2.5,3.5,6.5,3.5,-4.5,3.0,-1.5],[5.5,-5.5,3.5,1.5,4.5,-5.5,8.0,2.0,3.0,8.0,3.0,-6.0,-5.0,4.0,5.0,-1.0]]
Output: [[4.0,-19.0,7.0,-6.5,8.5,-8.5,-9.5,28.5,26.0,29.0,-4.0,-1.5,-6.0,-4.0,-3.0,-12.0],[4.0,-19.0,7.0,-6.5,8.5,-8.5,-9.5,28.5,26.0,29.0,-4.0,-1.5,-6.0,-4.0,-3.0,-12.0],[4.0,-19.0,7.0,-6.5,8.5,-8.5,-9.5,28.5,26.0,29.0,-4.0,-1.5,-6.0,-4.0,-3.0,-12.0],[4.0,-19.0,7.0,-6.5,8.5,-8.5,-9.5,28.5,26.0,29.0,-4.0,-1.5,-6.0,-4.0,-3.0,-12.0],[4.0,-19.0,7.0,-6.5,8.5,-8.5,-9.5,28.5,26.0,29.0,-4.0,-1.5,-6.0,-4.0,-3.0,-12.0],[4.0,-19.0,7.0,-6.5,8.5,-8.5,-9.5,28.5,26.0,29.0,-4.0,-1.5,-6.0,-4.0,-3.0,-12.0],[4.0,-19.0,7.0,-6.5,8.5,-8.5,-9.5,28.5,26.0,29.0,-4.0,-1.5,-6.0,-4.0,-3.0,-12.0],[4.0,-19.0,7.0,-6.5,8.5,-8.5,-9.5,28.5,26.0,29.0,-4.0,-1.5,-6.0,-4.0,-3.0,-12.0]]
me = jax.lax.axis_index("ring") is each device's position as a traced int.
jax.lax.ppermute(x, "ring", [(i, (i + 1) % 8) for i in range(8)]) sends to the next device.
chunks[i] and chunks.at[i].set(v) both accept a traced i.
Accepts: array
Add up one vector per device across 8 devices and give every device the total, building the all-reduce yourself from neighbour-to-neighbour sends with jax.lax.ppermute.
Row d of x is a vector of length C that belongs to device d. Build a one-axis mesh named ring over all 8 devices and use jax.shard_map with x split by rows, so each device starts with its own (1, C) row. The goal is the result that jax.lax.psum would give,
resultd=e=0∑7xefor every device d,where xe is row e of x. Cut every row into 8 chunks of C/8 consecutive values and connect the devices in a ring in which every device d sends to device (d + 1) mod 8, or one in which every device d sends to device (d - 1) mod 8. The ring algorithm runs two phases of 7 steps each. In the reduce-scatter phase, every step passes one chunk to the next device, which adds its own copy of that chunk before passing the running sum on, so after 7 steps each device owns one chunk that holds the sum over all 8 devices. In the all-gather phase, the finished chunks travel around the ring for 7 more steps until every device holds all 8 of them.
Return a float32 JAX array of shape (8, C), split by rows so that device d holds row d, in which every row equals the sum of the 8 rows of x.
Input: x = [[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0],[1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0],[2.0,2.0,2.0,2.0,2.0,2.0,2.0,2.0],[3.0,3.0,3.0,3.0,3.0,3.0,3.0,3.0],[4.0,4.0,4.0,4.0,4.0,4.0,4.0,4.0],[5.0,5.0,5.0,5.0,5.0,5.0,5.0,5.0],[6.0,6.0,6.0,6.0,6.0,6.0,6.0,6.0],[7.0,7.0,7.0,7.0,7.0,7.0,7.0,7.0]]
Output: [[28.0,28.0,28.0,28.0,28.0,28.0,28.0,28.0],[28.0,28.0,28.0,28.0,28.0,28.0,28.0,28.0],[28.0,28.0,28.0,28.0,28.0,28.0,28.0,28.0],[28.0,28.0,28.0,28.0,28.0,28.0,28.0,28.0],[28.0,28.0,28.0,28.0,28.0,28.0,28.0,28.0],[28.0,28.0,28.0,28.0,28.0,28.0,28.0,28.0],[28.0,28.0,28.0,28.0,28.0,28.0,28.0,28.0],[28.0,28.0,28.0,28.0,28.0,28.0,28.0,28.0]]
Explanation: Device d holds a row of eight copies of d, so each of its 8 chunks is the single value d. After reduce-scatter each device owns one chunk equal to 0 + 1 + 2 + 3 + 4 + 5 + 6 + 7 = 28, and after all-gather every device holds all 8 finished chunks, so every row is eight copies of 28.
Input: x = [[-2.0,-1.0,0.0,1.0,2.0,-2.0,-1.0,0.0],[1.0,2.0,-2.0,-1.0,0.0,1.0,2.0,-2.0],[-1.0,0.0,1.0,2.0,-2.0,-1.0,0.0,1.0],[2.0,-2.0,-1.0,0.0,1.0,2.0,-2.0,-1.0],[0.0,1.0,2.0,-2.0,-1.0,0.0,1.0,2.0],[-2.0,-1.0,0.0,1.0,2.0,-2.0,-1.0,0.0],[1.0,2.0,-2.0,-1.0,0.0,1.0,2.0,-2.0],[-1.0,0.0,1.0,2.0,-2.0,-1.0,0.0,1.0]]
Output: [[-2.0,1.0,-1.0,2.0,0.0,-2.0,1.0,-1.0],[-2.0,1.0,-1.0,2.0,0.0,-2.0,1.0,-1.0],[-2.0,1.0,-1.0,2.0,0.0,-2.0,1.0,-1.0],[-2.0,1.0,-1.0,2.0,0.0,-2.0,1.0,-1.0],[-2.0,1.0,-1.0,2.0,0.0,-2.0,1.0,-1.0],[-2.0,1.0,-1.0,2.0,0.0,-2.0,1.0,-1.0],[-2.0,1.0,-1.0,2.0,0.0,-2.0,1.0,-1.0],[-2.0,1.0,-1.0,2.0,0.0,-2.0,1.0,-1.0]]
Input: x = [[7.5,-3.5,-3.0,2.0,6.5,0.0,-0.5,-4.5,-3.0,-1.0,-5.0,-8.0,8.0,-4.0,-5.0,-6.5],[-2.5,-3.5,-5.0,-7.5,0.5,1.0,-6.0,4.0,6.0,8.0,5.5,-0.5,3.0,5.0,-1.0,5.5],[8.0,-7.0,4.5,6.5,2.5,-4.0,-5.0,7.0,8.0,0.5,-2.0,7.0,-2.5,-1.5,-3.0,-4.5],[-0.5,-4.0,-3.5,-7.5,-2.0,2.0,5.0,4.0,1.5,3.0,-5.5,-3.0,-5.5,2.0,-3.0,1.0],[-2.5,2.0,8.0,5.5,-4.0,0.0,0.5,7.5,1.5,7.5,-7.0,6.5,-7.5,-2.5,-1.0,-3.0],[-7.5,-1.5,4.5,-3.0,1.5,-0.5,-3.5,1.0,6.5,5.5,3.5,-4.0,0.0,-2.5,2.0,-2.0],[-4.0,4.0,-2.0,-4.0,-1.0,-1.5,-8.0,7.5,2.5,-2.5,3.5,6.5,3.5,-4.5,3.0,-1.5],[5.5,-5.5,3.5,1.5,4.5,-5.5,8.0,2.0,3.0,8.0,3.0,-6.0,-5.0,4.0,5.0,-1.0]]
Output: [[4.0,-19.0,7.0,-6.5,8.5,-8.5,-9.5,28.5,26.0,29.0,-4.0,-1.5,-6.0,-4.0,-3.0,-12.0],[4.0,-19.0,7.0,-6.5,8.5,-8.5,-9.5,28.5,26.0,29.0,-4.0,-1.5,-6.0,-4.0,-3.0,-12.0],[4.0,-19.0,7.0,-6.5,8.5,-8.5,-9.5,28.5,26.0,29.0,-4.0,-1.5,-6.0,-4.0,-3.0,-12.0],[4.0,-19.0,7.0,-6.5,8.5,-8.5,-9.5,28.5,26.0,29.0,-4.0,-1.5,-6.0,-4.0,-3.0,-12.0],[4.0,-19.0,7.0,-6.5,8.5,-8.5,-9.5,28.5,26.0,29.0,-4.0,-1.5,-6.0,-4.0,-3.0,-12.0],[4.0,-19.0,7.0,-6.5,8.5,-8.5,-9.5,28.5,26.0,29.0,-4.0,-1.5,-6.0,-4.0,-3.0,-12.0],[4.0,-19.0,7.0,-6.5,8.5,-8.5,-9.5,28.5,26.0,29.0,-4.0,-1.5,-6.0,-4.0,-3.0,-12.0],[4.0,-19.0,7.0,-6.5,8.5,-8.5,-9.5,28.5,26.0,29.0,-4.0,-1.5,-6.0,-4.0,-3.0,-12.0]]
me = jax.lax.axis_index("ring") is each device's position as a traced int.
jax.lax.ppermute(x, "ring", [(i, (i + 1) % 8) for i in range(8)]) sends to the next device.
chunks[i] and chunks.at[i].set(v) both accept a traced i.
Accepts: array
