How to efficiently do functional evaluation with a condition based on an array with integers.

[stefanvuckovic1]

I want to efficiently perform conditional function evaluation based on an array of integers and other arrays with real numbers serving as input for those functions. I hope to find a JAX-based solution that provides significant performance improvements over a for-loop approach that I describe below:

import jax
from jax import vmap;
import jax.numpy as jnp
import jax.random as random

def g_0(x, y, z, u):
    return x + y + z + u

def g_1(x, y, z, u):
    return x * y * z * u

def g_2(x, y, z, u):
    return x - y + z - u

def g_3(x, y, z, u):
    return x / y / z / u

g_i = [g_0, g_1, g_2, g_3]
g_i_jit = [jax.jit(func) for func in g_i]

def g_git(i, x, y, z, u):
    return g_i_jit[i](x=x, y=y, z=z, u=u)

def g(i, x, y, z, u):
    return g_i[i](x=x, y=y, z=z, u=u)


len_xyz = 3000
x_ar = random.uniform(random.PRNGKey(0), shape=(len_xyz,))
y_ar = random.uniform(random.PRNGKey(1), shape=(len_xyz,))
z_ar = random.uniform(random.PRNGKey(2), shape=(len_xyz,))

len_u = 1000
u_0 = random.uniform(random.PRNGKey(3), shape=(len_u,))
u_1 = jnp.repeat(u_0, len_xyz)
u_ar = u_1.reshape(len_u, len_xyz)


len_i = 50
i_ar = random.randint(random.PRNGKey(5), shape=(len_i,), minval=0, maxval= len(g_i)) #related to g_range-1


total = jnp.zeros((len_u, len_xyz))

for i in range(len_i):
    total= total + g_git(i_ar[i], x_ar, y_ar, z_ar, u_ar)

The role of "i_ar" is to act as an index that selects one of the four functions from the list g_i. "i_ar" is an array of integers, with each integer representing an index in the g_i list. On the other hand, x_ar, y_ar, z_ar, and u_ar are arrays of real numbers that are input to the functions selected by i_ar.

I suspect that this difference in nature between i_ar and x_ar, y_ar, z_ar, and u_ar is what could be difficult to find a JAX way that would be more efficient replacement of the for loop above. Any ideas how to use JAX (or something else) to replace the foor loop to obtain 'total' more efficiently? I have tried naively using vmap/pmap, but failed.

[soraros]

You might use lax.switch or jnp.choose for g (check out all the different ways to branch in JAX here), and jnp.sum along the first dimension at the very end. The performance is probably not that large for your example, except when running on accelerator.

def g(i, x, y, z, u):
  return lax.switch(i, g_i, x, y, z, u)
  # or
  return jnp.choose(i, [func(x, y, z, u) for func in g_i], mode='wrap')

@jit
def total(i, x, y, z, u):
  res = vmap(g, in_axes=(0, None, None, None, None))(i, x, y, z, u)
  return jnp.sum(res, axis=0)

total(i_ar, x_ar, y_ar, z_ar, u_ar)

this difference in nature between i_ar and x_ar, y_ar, z_ar, and u_ar

IIUC, this is actually a different question, and is exactly where vmap comes in handy. This answer about nested vmaps might be relevant. Or even better, try jnp.vectorize.

Example code here

def f(x, y, z, u):
  # let's pretend it's some complex pointwise function where broadcasting doesn't work
  assert all(a.shape == () for a in [x, y, z, u])
  return jnp.array(0)

x_ar = y_ar = z_ar = jnp.zeros(10)
u_ar = jnp.zeros(7)

# nested vmaps saves the day!
# batched_f deals with the said 'difference in nature' between xyz and u
batched_f = vmap(vmap(f, in_axes=(0, 0, 0, None)), in_axes=(None, None, None, 0))
assert batched_f(x_ar, y_ar, z_ar, u_ar).shape == (len(u_ar), len(x_ar))

# g works with arrays (x, y, z, u), and a single i
def g(i, x, y, z, u):
  g_i = [batched_f] * 3  # mock 3 branches
  return lax.switch(i, g_i, x, y, z, u)

# so we vmap again
batched_g = vmap(g, in_axes=(0, None, None, None, None))

i_ar = jnp.arange(11) % len(g_i)
assert batched_g(i_ar, x_ar, y_ar, z_ar, u_ar).shape == (len(i_ar), len(u_ar), len(x_ar))

[stefanvuckovic1]

@soraros this worked great! Thanks!