Delegate to specialized function on CPU

[AdrienCorenflos]

Hi,

I have a function func(a) for which I have a specialized cython implementation when on CPU. What is the recommended recipe to branch evaluation between CPU and GPU for this?
I can see that the base code is using "lowering" to either CPU/GPU/TPU but I am not sure how to do this for a simple function.

For the sake of the example, we can assume that I have two implementations func_cpu(a) and func_gpu(a) and I'm looking to branch from one to the other automatically based on where a is.

Note: I do not care about gradients here.

Thanks!

[jakevdp]

Unfortunately we don't have a public API for this, though we've discussed in the past the possibility of providing a custom_lowering mechanism similar to custom_jvp (cc/ @dfm).

The mechanism we do have for this is to define a primitive, and register device-specific lowering rules for that primitive. This is a bit involved and requires some non-public APIs, but a simple example in JAX's codebase is the spsolve_p primitive:

jax/jax/experimental/sparse/linalg.py

Lines 592 to 598 in 7c4cc95

	spsolve_p = core.Primitive('spsolve')
	spsolve_p.def_impl(functools.partial(xla.apply_primitive, spsolve_p))
	spsolve_p.def_abstract_eval(_spsolve_abstract_eval)
	ad.defjvp(spsolve_p, _spsolve_jvp_lhs, None, None, _spsolve_jvp_rhs)
	ad.primitive_transposes[spsolve_p] = _spsolve_transpose
	mlir.register_lowering(spsolve_p, _spsolve_gpu_lowering, platform='cuda')
	mlir.register_lowering(spsolve_p, _spsolve_cpu_lowering, platform='cpu')

Note that it has different lowerings on CPU and on GPU. mlir.lower_fun can be useful for converting Python implementations into mlir lowerings that you can register.

You can find more discussion of implementing custom primitives at https://jax.readthedocs.io/en/latest/jax-primitives.html

[dfm]

It sounds to me like jax.lax.platform_dependent might actually do what you need here. For example, the example from the original post might look like the following:

def func(a):
  return jax.lax.platform_dependent(a, cpu=func_cpu, cuda=func_gpu)

[jakevdp]

Wow, I'd somehow never encountered that. +1

[AdrienCorenflos]

Neat! Thanks @dfm you saved me a lot of pain :D

[AdrienCorenflos]

FYI, this question was related to branching jnp.searchsorted automatically as per this note:

(not exactly this because I have a very specific case where both my input array a and the values I'm searching therein v are sorted).