在Tensorflow中实现新的op时使用简单的数学函数

Question

我试图通过编写.cc文件（CPU和GPU版本）来实现C ++中的新操作。这个操作使用简单的数学函数，如cos或sin。

在CPU / GPU上调用浮点数时调用cos或sin（或类似函数）的正确方法是什么。我们是否必须包含外部库？它在Eigen吗？我们是否必须为CPU和GPU使用不同的库？ 如何申报？

除了zero_out.cc之外，您还知道更多的文档/示例吗？

Answer 1

我使用了std::cos和std::sin，它可以正常运行 ON CPU ！不知道它应该是怎么回事，我会保留这个答案，直到有人反驳它为止。

编辑：好吧，我说得太快了：GPU版本需要CUDA功能而不是C ++功能（所谓的主机功能）。