in_axes tells vmap() which input axes to vectorize over and which arguments to leave unchanged. It lets a single-element function run over batches without rewriting the function itself. section 2.5
The value can be an integer, None, or a matching Python container such as a tuple, list, or dict. A None entry means that argument is broadcast rather than split for mapping. section 6.2.1 section 7.2.1
Common patterns include:
- Use 0 for ordinary batch dimensions on positional arguments. section 6.2.1
- Use a different integer when the batch axis is not first, such as transposed data. section 6.2.1
- Use None for scalars or shared values that should be copied unchanged across mapped calls. section 6.2.1
- Mirror nested input structure with dicts or other standard containers. section 6.2.1
Mixed-axis vectorization
A scaled dot product can batch one vector on axis 0, another on axis 1, and keep a coefficient fixed with in_axes=(0,1,None). The same idea works when the vectors arrive inside a dict. section 6.2.1 section 6.2.1
Keyword arguments are always mapped over their leading axis, so they do not follow the positional in_axes tuple. When a keyword value should stay fixed, wrap or partially apply it first. section 6.2.3 section 6.2.3
