use raw ops

[fsx950223]

Fix #1779.
cc @WindQAQ
Related google/automl#480

[bot-of-gabrieldemarmiesse]

@mels630

You are owner of some files modified in this pull request.
Would you kindly review the changes whenever you have the time to?
Thank you very much.

[fsx950223]

cc @gabrieldemarmiesse.

[seanpmorgan]

Related: #1779

[WindQAQ]

Thank you @fsx950223! I believe this is correct, but the serialization/deserialization might break the saved models (for whom also save the input pipeline) due to different registration name.

Hi @seanpmorgan @gabrieldemarmiesse, should we throw some warning to tell users the implementation has changed? Thank you!

[fsx950223]

Thank you @fsx950223! I believe this is correct, but the serialization/deserialization might break the saved models (for whom also save the input pipeline) due to different registration name.

Hi @seanpmorgan @gabrieldemarmiesse, should we throw some warning to tell users the implementation has changed? Thank you!

It won't influence saved model. Since saved model needs to load custom kernel separately.

[gabrieldemarmiesse]

We have a function to load all SO, so as long as we keep the SO somewhere, we won't break existing models. This PR doesn't break existing saved models but it will remove the gradient for the op though

[fsx950223]

Because the PR fix AA on TPUs, I believe the above issues shouldn't block it.

[WindQAQ]

Thanks :-)