JAX vmap behaviour

user_3952424 0
Viewed 475

I'm trying to understand the behaviour of JAX vmap, so I wrote the following code:

import jax.numpy as jnp
from jax import vmap

def what(a,b,c):
  z = jnp.dot(a,b)
  return z + c

v_what = vmap(what, in_axes=(None,0,None))

a = jnp.array([1,1,3])
b = jnp.array([2,2])
c = 1.0

v_what(a,b,c)

And the output is:

DeviceArray([[3., 3., 7.],
             [3., 3., 7.]], dtype=float32)

I understand that the only input that is being altered is b, but Can someone shed some light on why this is the result? And how the dot product behaves after I vectorized the function?

jax vectorization python
1 Answers

You have specified that the transformed function should map over the first axis of b, and not map over any axis of a or c. So roughly, you've created a mapped function that does this:

def v_what(a, b, c):
  return jnp.stack([what(a, b_i, c) for b_i in b], axis=0)

For your inputs, within each row the dot product looks like jnp.dot(a, 2), and the result is equivalent to a * 2.

user_29378310
Related