vmapping lambda functions

[aadityacs]

I was wondering how to vmap functions of functions.

Lets suppose I have a simple function

def fn(x, y, z):
  return x + y + z

Then

sum = jax.vmap(fn)(x, y, z)

with

x = jnp.array([1, 2, 3])
y = jnp.array([4, 5, 6])
z = jnp.array([7, 8, 9])

returns as expected

Array([12, 15, 18], dtype=int32)

Now instead lets define a lambda

gn = lambda x: fn(x, y, z)

Then,

sum = jax.vmap(gn)(x)

returns

Array([[12, 14, 16],
       [13, 15, 17],
       [14, 16, 18]], dtype=int32)

Could someone kindly let me know how to get the same result for both?

[ianwilliamson]

Your lambda function, gn takes only a single argument and is being vmap-ed over only the elements of x. Inside of fn where you are doing x + y + z, each element of x gets broadcast with the values of y and z, which each have shape (3,). These get stacked by the vmap into the (3, 3) result.

[aadityacs]

Thank you! I now see what is happening!

[jakevdp]

I think the larger question here is "is it possible to vmap closed-over variables in a function". The answer to that is no: semantically, vmap may only map over explicit function arguments. With that in mind, there's no way to get your first output (which maps the function across all arguments) using the lambda function that closes over some arguments.

[aadityacs]

Thank you! I think I can use your pointer to redesign my approach. Really appreciate your time!