Branching behaviour depending on device used.

[AdrienCorenflos]

Hi,

What is the preferred way for a user to branch behaviour based on which kind of device the input is located?
To ground the question, say I am implementing the following trinket function:

@jax.jit
def inverse_cdf(cs, us):
    return jnp.searchsorted(cs, us)

for which I know the values cs and us are sorted.

As per the documentation of searchsorted

  Note:
    The ``method`` argument controls the algorithm used to compute the insertion indices.

    - ``'scan'`` (the default) tends to be more performant on CPU, particularly when ``a`` is
      very large.
    - ``'scan_unrolled'`` is more performant on GPU at the expense of additional compile time.
    - ``'sort'`` is often more performant on accelerator backends like GPU and TPU, particularly
      when ``v`` is very large.
    - ``'compare_all'`` tends to be the most performant when ``a`` is very small.

So I'd like to automatically use jnp.searchsorted(..., method="scan") on CPU and jnp.searchsorted(..., method="sort") otherwise.

I am seeing (although not understanding yet how I would use it) that JAX uses mlir primitive CPU/GPU lowering in the library's lower levels.
Is this the preferred end-user interface too or are there simpler alternatives?

Thanks