multi condtions/switch with predicate functions

[jecampagne]

Hello
Is there a JAX version with a multi-switch version of

def f(x):
   if (x < 0): 
        x = -x

    if (x == 0) : 
        return 1; 

    if (x <= 4): #            x in (0, 4]
        value = ...
    elif (x <= 8.0):             # x in (4, 8]
       value = ...
    else :  
        value = ...
    return value

Thanks

[jakevdp]

One way to do this is using jnp.select, which runs through a list of conditions and returns the associated value. For example:

import jax.numpy as jnp
import jax

@jax.jit
def f(x):
  x = abs(x)
  return jnp.select(
      [x == 0, x <= 4, x <= 8],
      [x, x - 1, x - 2],
      default = x)
  
print(f(0))
# 0
print(f(-2))
# 1
print(f(7))
# 5
print(f(10))
# 10

[jecampagne]

Ha... zut I have not seen this select function...
Suggestion: may be in the JAX doc when one searches "condition" this function would be selected... (nb Piecewie is selected but very deep in the search).

[jakevdp]

I added some additional docs in #14272