Run one gradient descent step for a linear model with the batch split across 8 devices. Use jax.shard_map so that every device differentiates the loss of its own shard of the batch, and use jax.lax.pmean to average the per-device results.
The model predicts xnW+b for example n, where xn is row n of x. The loss is the squared error summed over the K outputs and averaged over all N examples.
L(W,b)=N1n=1∑Nk=1∑K((xnW+b)k−yn,k)2Here W is w with shape (D, K), b is b with shape (K,), x has shape (N, D) and y has shape (N, K). Build a one-axis mesh named data over all 8 devices and split x and y along their first dimension, so each device receives N/8 consecutive examples. Each device evaluates the same loss formula on its own N/8 examples and differentiates it with respect to w and b. Every shard has the same size, so the average of the 8 shard losses equals L and the average of the 8 shard gradients equals the gradient of L. The step then applies
W←W−lr∇WL,b←b−lr∇bL.Return a dict with three keys. w is the updated float32 JAX array of shape (D, K), b is the updated float32 JAX array of shape (K,), and loss is a float32 JAX array of shape () holding L at the weights before the update. w and b are returned replicated on all 8 devices.
Input: w = [[0.0]], b = [0.0], x = [[1.0],[2.0],[3.0],[4.0],[5.0],[6.0],[7.0],[8.0]], y = [[3.0],[5.0],[7.0],[9.0],[11.0],[13.0],[15.0],[17.0]], lr = 0.01
Output: {"w":[[1.11]],"b":[0.2],"loss":121.0}
Explanation: Every device receives one example, so device d holds x = d + 1 and y = 2d + 3. With zero weights every prediction is 0, so the shard losses are 9, 25, 49, 81, 121, 169, 225 and 289, and their average is 121. The shard gradients with respect to w are -2xy, which average to -111, so the step gives w = 0 + 0.01 times 111 = 1.11. The bias gradients are -2y, which average to -20, so b becomes 0.2.
Input: w = [[0.1],[0.3]], b = [0.0], x = [[0.2,0.1],[0.7,0.8],[0.2,-0.7],[-1.3,-3.2],[-0.9,-0.8],[0.1,0.3],[0.3,-0.3],[-0.4,0.2],[0.1,-0.5],[0.4,0.2],[-0.6,0.5],[0.1,-0.2],[0.5,0.0],[1.0,1.0],[-0.1,0.4],[-0.9,-0.4]], y = [[-0.8],[-0.4],[-0.3],[0.1],[-0.6],[-0.3],[-0.6],[-0.7],[0.1],[-0.2],[-0.4],[-0.4],[-0.7],[0.2],[-0.1],[-0.3]], lr = 0.1
Output: {"w":[[0.072288],[0.239088]],"b":[-0.057],"loss":0.303913}
Input: w = [[0.1,1.0],[-1.3,-0.4]], b = [-0.2,-0.3], x = [[-0.5,0.8],[0.4,-0.4],[-0.9,-0.1],[0.8,0.7],[2.0,0.0],[-0.7,-0.4],[-0.2,-0.4],[-0.9,0.1]], y = [[-0.7,1.1],[0.6,-0.6],[0.1,1.3],[-1.0,-0.2],[0.2,-1.1],[0.0,1.0],[-0.3,0.6],[-0.5,1.0]], lr = 0.05
Output: {"w":[[0.104475,0.8363],[-1.291212,-0.39135]],"b":[-0.195125,-0.22975],"loss":3.745738}
jax.make_mesh((8,), ("data",), axis_types=(AxisType.Auto,)) lets shard_map reshard plain inputs.
jax.shard_map(f, mesh=mesh, in_specs=(P(), P(), P("data"), P("data")), out_specs=...).
jax.lax.pcast((w, b), "data", to="varying") marks shared params as per-device; jax.lax.pmean averages across devices.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: number
Run one gradient descent step for a linear model with the batch split across 8 devices. Use jax.shard_map so that every device differentiates the loss of its own shard of the batch, and use jax.lax.pmean to average the per-device results.
The model predicts xnW+b for example n, where xn is row n of x. The loss is the squared error summed over the K outputs and averaged over all N examples.
L(W,b)=N1n=1∑Nk=1∑K((xnW+b)k−yn,k)2Here W is w with shape (D, K), b is b with shape (K,), x has shape (N, D) and y has shape (N, K). Build a one-axis mesh named data over all 8 devices and split x and y along their first dimension, so each device receives N/8 consecutive examples. Each device evaluates the same loss formula on its own N/8 examples and differentiates it with respect to w and b. Every shard has the same size, so the average of the 8 shard losses equals L and the average of the 8 shard gradients equals the gradient of L. The step then applies
W←W−lr∇WL,b←b−lr∇bL.Return a dict with three keys. w is the updated float32 JAX array of shape (D, K), b is the updated float32 JAX array of shape (K,), and loss is a float32 JAX array of shape () holding L at the weights before the update. w and b are returned replicated on all 8 devices.
Input: w = [[0.0]], b = [0.0], x = [[1.0],[2.0],[3.0],[4.0],[5.0],[6.0],[7.0],[8.0]], y = [[3.0],[5.0],[7.0],[9.0],[11.0],[13.0],[15.0],[17.0]], lr = 0.01
Output: {"w":[[1.11]],"b":[0.2],"loss":121.0}
Explanation: Every device receives one example, so device d holds x = d + 1 and y = 2d + 3. With zero weights every prediction is 0, so the shard losses are 9, 25, 49, 81, 121, 169, 225 and 289, and their average is 121. The shard gradients with respect to w are -2xy, which average to -111, so the step gives w = 0 + 0.01 times 111 = 1.11. The bias gradients are -2y, which average to -20, so b becomes 0.2.
Input: w = [[0.1],[0.3]], b = [0.0], x = [[0.2,0.1],[0.7,0.8],[0.2,-0.7],[-1.3,-3.2],[-0.9,-0.8],[0.1,0.3],[0.3,-0.3],[-0.4,0.2],[0.1,-0.5],[0.4,0.2],[-0.6,0.5],[0.1,-0.2],[0.5,0.0],[1.0,1.0],[-0.1,0.4],[-0.9,-0.4]], y = [[-0.8],[-0.4],[-0.3],[0.1],[-0.6],[-0.3],[-0.6],[-0.7],[0.1],[-0.2],[-0.4],[-0.4],[-0.7],[0.2],[-0.1],[-0.3]], lr = 0.1
Output: {"w":[[0.072288],[0.239088]],"b":[-0.057],"loss":0.303913}
Input: w = [[0.1,1.0],[-1.3,-0.4]], b = [-0.2,-0.3], x = [[-0.5,0.8],[0.4,-0.4],[-0.9,-0.1],[0.8,0.7],[2.0,0.0],[-0.7,-0.4],[-0.2,-0.4],[-0.9,0.1]], y = [[-0.7,1.1],[0.6,-0.6],[0.1,1.3],[-1.0,-0.2],[0.2,-1.1],[0.0,1.0],[-0.3,0.6],[-0.5,1.0]], lr = 0.05
Output: {"w":[[0.104475,0.8363],[-1.291212,-0.39135]],"b":[-0.195125,-0.22975],"loss":3.745738}
jax.make_mesh((8,), ("data",), axis_types=(AxisType.Auto,)) lets shard_map reshard plain inputs.
jax.shard_map(f, mesh=mesh, in_specs=(P(), P(), P("data"), P("data")), out_specs=...).
jax.lax.pcast((w, b), "data", to="varying") marks shared params as per-device; jax.lax.pmean averages across devices.
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: number
