Cant jit the function using soft_sort.quantile() and jnp.where()

[yyang97]

import jax
import ott
import jax
import jax.random
import jax.numpy as jnp

N = 200
x = jax.random.normal(key = jax.random.PRNGKey(0), shape=(N,))
a = 0.5
quant = jnp.array([.2,.8])
def quantile_cal(x,quant,a):
   return( jax.vmap(lambda quant: ott.tools.soft_sort.quantile(jnp.where(x < a, x, a), 
                                                                    q=quant))(quant))


quantile_cal(x,quant,a)

# Array([[-0.81860286],
#        [ 0.49906105]], dtype=float32)

If I don't jit the function, it looks fine.

But if I jit it, it fails:

@jax.jit
def quantile_cal(x,quant,a):
   return( jax.vmap(lambda quant: ott.tools.soft_sort.quantile(jnp.where(x < a, x, a), 
                                                                    q=quant))(quant))


quantile_cal(x,quant,a)


# NonConcreteBooleanIndexError: Array boolean indices must be concrete; got ShapedArray(bool[3])

# See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.NonConcreteBooleanIndexError

However, I don't have any array boolean things...

Further, it seems that ott package about the quantile function is updated... The same code works before...

[michalk8]

Hi @yyang97 , this should be fixed in the latest version (0.4.3) by this #408 !