在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++功能（所谓的主机功能）。