Relationship with XLA consumers

[bhack]

We have a quite old topic about the tension between custom-op vs XLA composability

I really want to know what is the strategy to interact with the XLA "consumers" projects when they hit the limit of expressing their computation with XLA/*HLO.

Very recently we had again a proposal to introduce C++ native OPS for the first time in historically pure python library like Keras starting from the limit to express IOU3D.

But as this project is not related exclusively to Keras or TF it seems that also Pytorch 3d needed to rely on c++/cuda custom ops so I suppose that currently we will have problem also with Pytorch-XLA/Torch-MLIR.

Then, to add also another framework, JAX is recently introducing something like custom-ops API in the framework

JAX uses XLA to compile staged-out Python programs, which are represented with MHLO. MHLO offers common operations used in machine learning, such as the dot-product, exponential function, and so on. Unfortunately, sometimes these operations (and compositions of these operations) are not sufficiently expressive or performant and users look to external libraries (Triton, Numba) or hand-written code (CUDA, C/C++) to implement operations for their desired program.

In general I could understand that we want a fast "exit strategy" to call well historical performing native implementation without contributing/interacting too much with the compiler stack/team but it seems to me that the strategy to stress test XLA on a more rich distribution of Tensor Programs it is really not too much clear IMHO.

When a "consumer" framework hit the barrier of expressing a Tensor Program for a specific operation how they are going to interact with the OpenXLA team?

I think that just introducing every time the "custom-ops" shortcut also in "relatively side" python projects will not help to positively stress test the compiler technology and its standard golden set of operations (?)HLO (/cc @burmako)

[bhack]

Just to make an example, I'am experimenting a top-down approach with:
tensorflow/tensorflow#58166

I have many doubts that it will work but let see what happens and if it will open a clear path for the contribution ((/cc @stellaraccident).

[bhack]

Another top-down case is 3d deformable convolution proposed recently in Keras-cv:
keras-team/keras-cv#1140

/cc @SimonBiggs @axeldavy