Frank Chen
Frank Chen
While it is technically possible to write an XLA HLO op and get it to run on TPUs, we currently don't expose any way to load arbitrary user written HLO...
Yup, working on it. Thanks for bringing that to our attention!
@orsharir Would it be possible to just encapsulate the op you want in a Python function, and then just switch between your custom op implementation and the TF default op...
@XMaster96 Unfortunately not because the underlying TPU ISA and associated tools for you to be able to write a XLA op isn't exposed even with the TPU VM preview.
You can load custom ops that someone else (or you) has written so long as they are CPU custom ops (or build them into a custom TF build). I don't...