We're trying to implement a piecewise function, basically around 100 polynomials with different coefficients depending on the value of x.
This will be implemented in TensorFlow or jax with JIT and be optimized for arrays of data. The question is what is probably the best way to achieve this?
One could use one-hundred wheres, but that is not really optimal. Or use the tf.switch_case with tf.vectorize_map (or similar).
Are there any ideas?
If I understand correctly, I think that jax.lax.switch provides the kind of functionality you're interested in. For example:
import jax.numpy as jnp
from jax import vmap, lax
import matplotlib.pyplot as plt
def f1(x):
return 0.0 * x
def f2(x):
return (x - 1.0) ** 2
def f3(x):
return 2 * x - 3
branches = (f1, f2, f3)
bounds = jnp.array([1, 2]) # boundaries between branches
x = jnp.linspace(0, 3)
index = jnp.searchsorted(bounds, x) # index in branches for each value in x
result = vmap(lambda i, x: lax.switch(i, branches, x))(index, x)
plt.plot(x, result)
