vmap with non vmappable parameter

[aniquetahir]

Is there a way to create a vmap inside a jitted function where one of the parameters is broadcast across the vmapped function?

e.g.

def inner_function(x, y):
    return x*y

@jit
def outer_function(X, y_):
    vmapped_function = vmap(inner_function)
    return vmapped_function(X, y_)  # here y, I don't want y to be vmapped

@jit
def outer_outer_function(XX, Y):
    vmapped_function = vmap(outer_function)
    return vmapped_function(XX, Y) # here Y is vmapped 

Dim XX: (10, 10, 10, 1)
Dim X: (10, 10, 1)
Dim Y: (10, 10, 1)
Dim x: (10, 1)
Dim y: (10, 1)
Dim y_: (10, 1)

Basically, here I don't want y_ to be vmapped. But by default it will be.

[jakevdp]

I think in_axes is what you're looking for:

def inner_function(x, y):
    return x*y

@jit
def outer_function(X, y_):
    vmapped_function = vmap(inner_function, in_axes=(0, None))
    return vmapped_function(X, y_)

[aniquetahir]

Thanks!