How can we calculate multiple quantiles simultaneously?

[yyang97]

x_test = jax.random.normal(key,shape = (100,))
ott.tools.soft_sort.quantile(x_test, level=0.2)

It is good to calculate the 0.2 quantile value of x_test.

However, how can I simultaneously calculate 0.2 and 0.8 quantile values of x_test?
like

x_test = jax.random.normal(key,shape = (100,))
ott.tools.soft_sort.quantile(x_test, level=jnp.array([.2,.8]))

However, it reports some errors:
ValueError: All input arrays must have the same shape.

[michalk8]

Hi @yyang97 ,

yes, you can vmap over the levels, e.g.:

import jax
from ott.tools import soft_sort

key = jax.random.PRNGKey(0)
x_test = jax.random.normal(key, shape=(100,))
levels = jnp.array([0.2, 0.8])

jax.vmap(lambda level: soft_sort.quantile(x_test, level=level))(levels)
# Array([[-0.7245259],
#        [ 0.93384  ]], dtype=float32)

[yyang97]

Thank you very much for your help!
May I ask one more question?
How do I calculate jax.vmap(lambda level: soft_sort.quantile(x_test[x_test<1], level=level))(levels), but in a jitted-compilable way?
So I want to calculate the quantiles of the parts lower than 1...

[marcocuturi]

another approach would be to have t+2 targets where t is the number of quantiles, and then "craft" more intelligently a target measure. This is what I was referring to in this TODO, and which was done in a previous implementation

https://github.com/google-research/google-research/blob/fa49cd0231954458e2b02278d95528b70245f06d/soft_sort/ops.py#L210

If you need this for an application, can you raise an issue? I will try to implement this shortly.

[marcocuturi]

@yyang97 you can take a look at #382

[yyang97]

Thank you Marco for doing this part!

May I ask one more question?
How do I calculate jax.vmap(lambda level: soft_sort.quantile(x_test[x_test<1], level=level))(levels), but in a jitted-compilable way?
So I want to calculate the quantiles of the parts lower than 1...

[marcocuturi]

The selection step x_test[x_test<1] is going to return an array of variable size. JIT won't like that indeed. Maybe an alternative is only to keep a subset of a predefined size. Another option would be to add padding, to always get the same size, and then set the weights parametera to 0 inside soft-sorting for those elements you''d like to discard