When a transformed function returns a mapped result, out_axes specifies where that mapped dimension appears in the output. The default leading position is not always convenient: setting out_axes=1 places the batch dimension at index 1, producing a transposed layout without an extra transformation. section 6.2.2 section 6.4
Place the mapped vector axis at index 1
def scale(v, koeff): return koeff * v scale_batched = jax.vmap( scale, in_axes=(0, None), out_axes=1 ) res = scale_batched(v1s, 2.0) With input shape (8, 3), the result has shape (3, 8): vector components occupy axis 0 and the eight scaled vectors occupy axis 1. The coefficient is broadcast rather than mapped.
The same control applies to pmap(). In data-parallel training, out_axes=(None, 0) selects replicated model parameters from the first device while retaining one loss value per device. Use None only when every device returns the same value. section 7.2.1 section 7.3.2
Listing 6.13 Using the out_axes parameter
def scale(v, koeff): #1
return koeff*v
scale_batched = jax.vmap(scale,
in_axes=(0,None), #2
out_axes=(1)) #3
scale_batched(v1s, 2.0)
>>> Array([[-1.4672383 , -1.6510035 , 3.5308602 , -2.2189112 , 0.3024418 ,
>>> 0.7649379 , -4.028754 , -3.0968533 , 0.34476107, -2.9087348 ],
>>> [-1.5357308 , -0.7061183 , 4.0082793 , -0.69232166, -3.2186441 ,
>>> 2.0812016 , 3.585087 , 0.15288436, 2.0001278 , 2.0246687 ],
>>> [-1.6228952 , 1.5497094 , -3.2013843 , 0.50127184, -0.2000112 ,
>>> 1.6244346 , 0.17156784, -1.3113307 , -2.532448 , -1.60574 ]],
>>> dtype=float32)
#1 A simple function to scale a vector
#2 Specifies batch dimensions for the input
#3 Specifies batch dimension for the output
Runnable listing showing how in_axes and out_axes jointly transform scaled vectors. section 6.2.2
