在Tensorflow中實現新操作時使用簡單的數學函數

Question

我想通過編寫一個.cc文件（CPU和GPU版本）來實現C++中的新操作。這個操作使用簡單的數學函數，如cos或sin。

什麼是正確的方式來調用CPU/GPU浮點數的COS或罪（或類似）。我們是否必須包含外部庫？它在Eigen嗎？我們是否必須爲CPU和GPU使用不同的庫？ 如何申報？

除了zero_out.cc以外，還有一些其他的文檔/例子嗎？

Answer 1

我用std::cos和std::sin它工作正常ON CPU！不知道它是如何，但我會保持這個答案，直到有人反駁它。

編輯：好吧，我說得太快了：GPU版本需要CUDA功能，而不是C++功能（所謂的主機功能）。