vmap function-valued functions

[pharringtonp19]

Aim

I would like to apply vmap to a function-valued function.

Motivation

I have a functional and a parameterized class of functions. I would like to be able to evaluate the functional in parallel across this class of functions.

Current Understanding

This isn't feasible because function is not a jax type

Hope

Is there a workaround?

Toy Example

import jax 
import jax.numpy as jnp 

def func(alpha):
  
  def f(x):
    return 1-x**alpha 
  
  return f 

jax.vmap(func)(jnp.linspace(1., 4., 10))

TypeError: Output from batched function <function func.<locals>.f at 0x7f39930dfcb0> with type <class 'function'> is not a valid JAX type

[jakevdp]

Can you say more about what you'd like the output to be? How would you interact with the resulting array of functions?

[pharringtonp19]

@jakevdp thanks for the response.

As you suggest, I would like the output to be an array of functions. That way, I could apply vmap to the functional defined below and pass in the array of functions.

def integrate_grad(f, n=2000):
  xs = jnp.linspace(0., 1., n)
  return 0.5* (1/n)*jnp.sum(jax.vmap(jax.grad(f))(xs)**2)

In other words, I would like to be able to do something like the following

jax.vmap(integrate_grad)(jax.vmap(func)(jnp.linspace(1., 4., 10)))

[jakevdp]

OK, thanks for the clarification. As you mention, function is not a valid JAX type, and I don't imagine this changing any time soon.

As a workaround, you might consider parameterizing your functions and passing around arrays of parameters, rather than passing around arrays of function objects. For example, you might define your function like this:

def integrate_grad(f, *args, n=2000):
  ...

and pass f as a static, unmapped function parametrized by the content of args, which is mappable.

[pharringtonp19]

@jakevdp I appreciate your insight on this -- that works

def f(alpha, x):
  return 1-x**alpha 

def integrate_grad(f, alpha, n=2000):
  f = partial(f, alpha)
  xs = jnp.linspace(0., 1., n)
  return 0.5* (1/n)*jnp.sum(jax.vmap(jax.grad(f))(xs)**2)

jax.vmap(integrate_grad, in_axes=(None, 0))(f, jnp.linspace(1., 2., 10))