Vmap#
vmap()#
vmap is another of the key components of jax. What it does is that it lets you pass a batch of inputs to your functions in a single go instead of passing them one by one. Think of it like this. You have 10 arrays which you need to pass to a function. In a typical python program, you would be writing a loop. With vmap, you can send them as a batch, take advantage of vector and matrix operations supported by your accelerator (GPU, TPU etc.) and make your computation way faster.
Let’s consider a simple affine transformation over some arrays.
import jax
import jax.numpy as jnp
key = jax.random.key(42)
key, *subkeys = jax.random.split(key, 3) # main key + 2 keys more
@jax.jit
def affine(x, w, b):
return x @ w.T + b
# say for this affine transformation, we are taking
# a 10 dim vector
# and making it a 2 dim vector
in_dim = 10
out_dim = 2
x = jax.random.normal(key, shape=(in_dim, ))
w = jax.random.normal(subkeys[0], shape=(out_dim, in_dim))
b = jax.random.normal(subkeys[1], shape=(out_dim, ))
affine(x, w, b)
Array([-0.98225653, 1.9490494 ], dtype=float32)
But that was over one x, w and b. What if we want to pass a batch as intended?
batch_size = 10
xs = jax.random.normal(key, shape=(batch_size, 10, ))
ws = jax.random.normal(subkeys[0], shape=(batch_size, out_dim, in_dim))
bs = jax.random.normal(subkeys[1], shape=(batch_size, out_dim, ))
If you want to add \(b\) and make it a full fledged affine transformation ….
# won't run due to shape mismatch
affine(xs, ws, bs)
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
Cell In[7], line 1
----> 1 affine(xs, ws, bs)
[... skipping hidden 13 frame]
Cell In[3], line 3, in affine(x, w, b)
1 @jax.jit
2 def affine(x, w, b):
----> 3 return x @ w.T + b
File ~/projects/annotated-jax/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:1141, in _forward_operator_to_aval.<locals>.op(self, *args)
1140 def op(self, *args):
-> 1141 return getattr(self.aval, f"_{name}")(self, *args)
File ~/projects/annotated-jax/.venv/lib/python3.12/site-packages/jax/_src/numpy/array_methods.py:604, in _defer_to_unrecognized_arg.<locals>.deferring_binary_op(self, other)
602 args = (other, self) if swap else (self, other)
603 if isinstance(other, _accepted_binop_types):
--> 604 return binary_op(*args)
605 # Note: don't use isinstance here, because we don't want to raise for
606 # subclasses, e.g. NamedTuple objects that may override operators.
607 if type(other) in _rejected_binop_types:
[... skipping hidden 14 frame]
File ~/projects/annotated-jax/.venv/lib/python3.12/site-packages/jax/_src/numpy/tensor_contractions.py:254, in matmul(a, b, precision, preferred_element_type, out_sharding)
252 a = lax.squeeze(a, tuple(a_squeeze))
253 b = lax.squeeze(b, tuple(b_squeeze))
--> 254 out = lax.dot_general(
255 a, b, (((np.ndim(a) - 1,), (np.ndim(b) - 1 - b_is_mat,)), (a_batch, b_batch)),
256 precision=precision, preferred_element_type=preferred_element_type,
257 out_sharding=out_sharding)
258 result = lax.transpose(out, perm)
259 return lax._convert_element_type(result, preferred_element_type, output_weak_type)
[... skipping hidden 13 frame]
File ~/projects/annotated-jax/.venv/lib/python3.12/site-packages/jax/_src/lax/lax.py:5303, in _dot_general_shape_rule(lhs, rhs, dimension_numbers, precision, preferred_element_type, out_sharding)
5300 if not core.definitely_equal_shape(lhs_contracting_shape, rhs_contracting_shape):
5301 msg = ("dot_general requires contracting dimensions to have the same "
5302 "shape, got {} and {}.")
-> 5303 raise TypeError(msg.format(lhs_contracting_shape, rhs_contracting_shape))
5305 return _dot_general_shape_computation(lhs.shape, rhs.shape, dimension_numbers)
TypeError: dot_general requires contracting dimensions to have the same shape, got (10,) and (2,).
Enter vmap:#
vmapped_affine = jax.vmap(affine)
vmapped_affine(xs, ws, bs)
Array([[-9.8225653e-01, 1.9490494e+00],
[-3.1628523e+00, -6.4753151e-01],
[ 3.6174815e+00, -2.1167760e+00],
[-5.5663157e+00, -6.2317848e-03],
[-1.7511344e+00, -3.6932726e+00],
[ 2.4414740e+00, -9.9153006e-01],
[-4.6595273e+00, -1.0493579e+01],
[-3.0498884e+00, 1.7762201e+00],
[-6.9430470e-03, 1.6245110e+00],
[ 1.3217969e+00, -1.9404722e+00]], dtype=float32)
It worked!
What if I want to vmap on one axis?#
Now if you might be yelling at me for having separate w and b for each x, especially if you’re someone who is used to affine transformation (actually linear but all linear transformations are affine anyway) from machine learning or pytorch’s nn.Linear. Fret not. We can have a single w and b for all xs. Actually let’s write a simple linear discriminator (or classification) model, that takes a vector / batch of vector as inputs and gives a dim 2 vector as output as class probabilities.
(Certainly not going to train this!, just the inference part. Besides I don’t have any labels :P )
# using the same w, b values from before
params = (w, b)
@jax.jit
def model(params, x):
w, b = params
logits = x @ w.T + b
return jax.nn.softmax(logits, axis=-1) # convert to probability distribution
Okay but we have to keep only one value of params for all x, right? You can tell vmap to ignore params. How? using in_axes
vmapped_model = jax.vmap(model, in_axes=(None, 0))
vmapped_model(params, xs)
Array([[0.0507002 , 0.94929975],
[0.02904104, 0.970959 ],
[0.03516598, 0.964834 ],
[0.150932 , 0.84906805],
[0.09174331, 0.9082567 ],
[0.2507983 , 0.74920166],
[0.88580555, 0.11419442],
[0.9345587 , 0.06544128],
[0.95485044, 0.0451496 ],
[0.04433183, 0.9556682 ]], dtype=float32)
in_axes lets you mention which params of the function you would like to be “vmapped”. If you wanted to pass as w, b, x instead of params, x; you could’ve written
vmapped_model = jax.vmap(model, in_axes=(None, None, 0))
A value of None will tell vmap to ignore that particular parameter. (Maintain order, you can randomly assign None !). For everything else, you can mention on which dimension should the array be vectorised. I’ve used 0 here. You can use something else for more complex tasks.