pmap.py

x = np.arange(5)
w = np.array([2., 3., 4.])

def convolve(x, w):
  output = []
  for i in range(1, len(x)-1):
    output.append(jnp.dot(x[i-1:i+2], w))
  return jnp.array(output)

convolve(x, w)

n_devices = jax.local_device_count() 
xs = np.arange(5 * n_devices).reshape(-1, 5)
ws = np.stack([w] * n_devices)
jax.pmap(convolve)(xs, ws)
# ShardedDeviceArray([[ 11.,  20.,  29.],
#                    .................
#                     [326., 335., 344.]], dtype=float32)