TensorflowOp Module+Plugin to handle TensorFlow Custom Ops

[ghost]

Request description

My company is interested in experimenting and potentially adopting OpenXLA/IREE in our model inference stack. But one blocker is that all our major models have at least one TF custom op, so we can't experiment IREE easily. I hope a TensorFlowOpModule can be formally supported and a TF model with custom ops can be compiled and executed with IREE out of the box. Below I list out all the impacts of this project from my perspective:

• Unblock quick experiments. This is very critical for new users to onboard IREE. If our company get good results from experiments and decides to onboard, there will be a whole team building our stack on top of IREE, and we will contribute a lot to the source code if needed. With this, the impact is actually much bigger than this feature alone.
• Easily land models that are not latency-critical. I understand using custom op may harm the performance. But there are always some small models that don't care much about latency. Since there are already custom ops in TF, people won't spend much time migrating the custom ops to native IREE ops.
• More friendly to ML engineers. In my company, only optimization team touches model compiler. However, models are developed and maintained by modeling teams. The ML engineers are more familiar with TF APIs and it's hard for them to onboard IREE custom modules or MLIR. It's not always the case that the optimization team has the bandwidth to immediately help them develop the IREE custom module, so it's very useful if they can use TF custom ops as a temporary workaround.
• Better TensorFlow integration. In classic XLA world, TF developed a feature "autoclustering" to handle ops that are not supported by XLA. However, that's not very ideal. A model is broken up into tens of sub XLA engines, thus it's not even possible to visualize the XLA graph or use XLA tools to debug the whole model. A TensorFlow module would be much better, if there is plan to replace classic XLA with OpenXLA.

What component(s) does this issue relate to?

No response

Additional context

No response

[aaron-schneider]

Thanks for the query. Since there's no one on the other side of this issue, we're closing for now. Feel free to open a new issue if you'd like to bring it up again.

[stellaraccident]

We can leave this open for now. I am trying to connect it to someone internally.

(But as mentioned on Discord, the IREE team is not planning to do work on this in 2023)

[ShuaiShao93]

Sorry I was using a wrong account so I closed that account. This account is the POC for this FR.