in_axes and out_axes: Controlling Vectorization

The in_axes parameter tells vmap which axis of each input to vectorize over. The out_axes parameter controls where the batch dimension appears in the output.

import jax
import jax.numpy as jnp

def dot_product(a, b):
    """Dot product of two vectors."""
    return jnp.sum(a * b)

# in_axes=(0, 0): vectorize over axis 0 of both a and b
batched_dot = jax.vmap(dot_product, in_axes=(0, 0))

A = jnp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])  # (3, 2)
B = jnp.array([[0.5, 0.5], [1.0, 0.0], [0.0, 1.0]])  # (3, 2)

# Computes dot product for each pair of rows
print(batched_dot(A, B))  # [1.5, 3.0, 6.0]

None: Broadcast (Don't Vectorize) an Argument

import jax
import jax.numpy as jnp

def scale_and_add(x, scale, bias):
    return x * scale + bias

# Vectorize over x (axis 0), but broadcast scale and bias
batched = jax.vmap(scale_and_add, in_axes=(0, None, None))

X = jnp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
scale = jnp.array(2.0)
bias = jnp.array([0.1, 0.2])

result = batched(X, scale, bias)
print(result)
# [[ 2.1  4.2]
#  [ 6.1  8.2]
#  [10.1 12.2]]

Vectorizing Over Different Axes

import jax
import jax.numpy as jnp

def matrix_vector_product(matrix, vector):
    return matrix @ vector

# matrix: vectorize over axis 0 (batch of matrices)
# vector: vectorize over axis 0 (batch of vectors)
batched_mv = jax.vmap(matrix_vector_product, in_axes=(0, 0))

matrices = jax.random.normal(jax.random.PRNGKey(0), (5, 3, 4))  # 5 matrices of 3x4
vectors = jax.random.normal(jax.random.PRNGKey(1), (5, 4))       # 5 vectors of length 4

result = batched_mv(matrices, vectors)
print(result.shape)  # (5, 3) — 5 results of length 3

out_axes: Controlling Output Batch Position

import jax
import jax.numpy as jnp

def my_fn(x):
    return x ** 2

# Default: batch dimension at position 0 in output
default = jax.vmap(my_fn)(jnp.ones((5, 3)))
print(default.shape)  # (5, 3) — batch dim at 0

# Put batch dimension at position 1 in output
swapped = jax.vmap(my_fn, out_axes=1)(jnp.ones((5, 3)))
print(swapped.shape)  # (3, 5) — batch dim at 1