Do we still need cuda programming from the perspective of jax/flax?

[wztdream]

Hi,
This is a general question.

Sometimes we need some customized operation which is not provided by framework lib, such as deformable convolution. In pytorch or tensorflow this can be done by cuda programing.

Considering jax/flax themself have the mechanism for parallel acceleration (XLA), so arise the quesiton: does it means we do not need cuda programming in jax/flax? How is the performance compered with cuda programming theoretically?

[avital]

I'm not sure if the repository is up-to-date but this might help: https://github.com/dfm/extending-jax

[avital]

(But in general, some things in JAX don't require custom kernels because XLA knows how to optimize many operations out-of-the-box.)

[wztdream]

Thank you, I take it that for most case we do not need cuda programming XLA from jax can handle it properly, but for some edge cases we may need.

[wztdream]

I hope one day we can focus only on logic instead of worrying about performance too much.