vmap over a callable[[float], Array]

[gautierronan]

Let's assume I have a function f given by the user (i.e. I am allowed to wrap it, but not allowed to modify it or to know its internals). This function has signature f(t: float) -> Array. I would like to jax.vmap over the returned array of this function.

In other words, I would like to achieve this kind of behavior.

import jax
import jax.numpy as jnp
from jax import Array

def f(t: float) -> Array:
    return jnp.array([t, t**2])

def g(t: float) -> Array:
    return jnp.array([-t**2, -t])

def foo(f: callable[[float], Array], g: callable[[float], Array]) -> Array:
    return f(1.0) + g(2.0)

jax.vmap(foo, in_axes=(0, 0))(f, g)
jax.vmap(foo, in_axes=(0, None))(f, g)
# ValueError: vmap was requested to map its argument along axis 0,
# which implies that its rank should be at least 1, but is only 0 (its shape is ())

For standard arrays, it would work fine.

x = jnp.array([1.0, 1.0])
y = jnp.array([-4.0, 2.0])

def bar(x: Array, y: Array) -> Array:
    return x + y

jax.vmap(bar, in_axes=(0, 0))(x, y)
jax.vmap(bar, in_axes=(0, None))(x, y)
# Array([-3.,  3.], dtype=float32)
# Array([[-3.,  3.],
#        [-3.,  3.]], dtype=float32)

I'm guessing I should be able to wrap f and g in a certain way to allow this kind of vmap? How could I achieve something like that? I know it should involve jax.eval_shape in some way, but would love any kind of pointers to get me started.

Thanks for the help.

[jakevdp]

The arguments to a vmapped function must be arrays, not functions. jax.vmap(foo, in_axes=(0, 0))(f, g) will not work, because f and g are functions, not arrays.

The fix would be to do something like your second solution, where you are passing arrays rather than functions as arguments to the vmapped function.

I'm not entirely sure what else to suggest, because it's not clear to me what your desired outcome is when passing a function in place of an array.

[gautierronan]

Yes exactly, I'm thinking of something with the following form:

x = 3.0
def f(t):
    return x + t

def foo(f):
    return f(2)

In the case I'm looking at, the function foo has f as input and f has a fixed output shape (which can be evaluted using jax.eval_shape). So I'm guessing that in principle, vmap has all the knowledge required to vectorize over the output dimension of f. It doesn't actually need to know about x, if that makes sense.

[jakevdp]

There's no way to do that kind of action-at-a-distance in Python, so there's no way to do that kind of action-at-a-distance in JAX. You could maybe hack things by temporarily setting global variables within the outer function, but that seems brittle and confusing at the very best.

I'd suggest defining your functions to take explicit arguments rather than relying on global state.

[gautierronan]

I don't quite agree. What I'm looking for is not action-at-a-distance. The following jax syntax is very much valid (i.e. fully explicit) and cannot be parametrized by an additional explicit argument in any simple way. Yet, it should be vmappable in the sense that it's just the lowered code executing things at once instead of sequentially.

def f(t: float) -> Array:
    return jnp.array([t, 3 * t, t**2])

def foo(f: callable[[float], Array]) -> Array:
    return f(1.0) + f(2.0)

This example in particular will run without any need for a vmap. But my very first example wouldn't.

[jakevdp]

OK, I guess I didn't understand what you wanted to do after all.

If what you want is for 1.0 and 2.0 to be arrays, you'll have to create those arrays and pass them as arguments to f and g. There's no way to implicitly convert constant scalars to arrays using vmap.

[jakevdp]

For example, this will work:

def f(t: float) -> Array:
    return jnp.array([t, 3 * t, t**2])

def foo(x, y, f: Callable[[float], Array]) -> Array:
    return f(x) + f(y)

jax.vmap(foo, in_axes=(0, 0, None))(jnp.ones(3), jnp.zeros(3), f)

Or, if you want to map the inner function rather than the outer, you can do this:

foo(jnp.ones(3), jnp.zeros(3), jax.vmap(f))