Understanding vmap

[ajithmoola]

Lets say a vmapped function can be defined as jax.vmap(func, in_axes=(0, None, None))(mapped_array, array1, array2, indexing_arr). The mapped dimension of func has a size of 100000. Is jax.vmap going to make 100000 copies of array1 and array2?

[jakevdp]

No, JAX will not in general make copies of unmapped arrays, although its possible for particular primitives to do things this way in their batching rule. If you want to be sure, you can use jax.make_jaxpr and inspect the size of the arrays created by the transformed function. For example:

import jax
import jax.numpy as jnp

mapped_array = jnp.zeros((100000, 3))
array1 = jnp.arange(3)
array2 = jnp.ones(3)

def func(arr, arr1, arr2):
  return arr + arr1 + arr2

jax.make_jaxpr(jax.vmap(func, in_axes=(0, None, None)))(mapped_array, array1, array2)

{ lambda ; a:f32[100000,3] b:i32[3] c:f32[3]. let
    d:f32[3] = convert_element_type[new_dtype=float32 weak_type=False] b
    e:f32[1,3] = broadcast_in_dim[broadcast_dimensions=(1,) shape=(1, 3)] d
    f:f32[100000,3] = add a e
    g:f32[1,3] = broadcast_in_dim[broadcast_dimensions=(1,) shape=(1, 3)] c
    h:f32[100000,3] = add f g
  in (h,) }

You can see that the unmapped input arrays (called b and c in jaxpr variables) are never explicitly duplicated across the batch dimension.