How to code "if" statement when using "vmap"

[shixinxing]

I encounter a problem when I use vmap for a simple function containing if statement.
I wonder if this is the common problem since I am a new user to JAX.

Here is the code:

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

def k(u):
  if u ==1.:
    return u
  elif u==-1.:
    return 2*u
  else:
    return 3*u

k_map = vmap(k,(0))
u = jnp.linspace(-1.,1.,5)
k_map(u)

Then it shows error:

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected.

The problem arose with the bool function.

Thanks a lot for your help!

[jeremiecoullon]

Check out the control flow section of "The Sharp Bits" in the documentation; it discusses this!
Issue #196 also explains this (in @mattjj's main reply).

If you want to use if statements in Jax, use lax.cond or lax.switch.
Hope that helps!

[mattjj]

You can also often use jnp.where in cases like the example in the OP.