ドキュメントによれば、tf.py_funcを使用して自分の操作を定義できます。 Numpy配列を受け取り、Numpy配列を返す必要があります。それは文書に明示的に書かれていませんが、私はそれがCPUデバイスに束縛されると結論づけますか？CPU操作のpy_func？
これをGPU上で動作する可能性がある他のopsと組み合わせると、TFは貪欲に多くの計算をGPUに移して、自動的にGPUとCPUの間でメモリを転送しますか？tf.py_func op？ （ちょうどTheanoのように）
tf.py_funcのようなGPU操作を定義するようなものはありますか？
py_funcは、奇妙なものです - それはopを作成するのに使用されたインタプリタと同じPythonインタプリタでPythonコードを実行します。だからあなたの周りのオペレーションがGPUの場合、GPU < - > CPU転送があります。さらに、Pythonのアドレス空間とTensorFlowのアドレス空間（つまり、memcpy here）の間でデータを移動するためのコピーがあります。
GPUでPythonコードを実行する方法を尋ねる場合、Numbaがあります。あなたがTensorFlow GPUに「機能」のようなものを持つ方法を尋ねるなら、Defunがあります。一般に、TensorFlowはnumpyを使用する必要がないように十分なnumpyフィーチャを定義する方向に移動していますが、ネイティブTFプリミティブで関数を実装できます。
TensorFlowエンジンに頼っている間に、Python-landにとどまる可能性があります。 Python数値型interfaceをエミュレートするPythonラッパーを作成できますが、実際の作業を基になるTensorFlowエンジンに委譲します。どのようにnumpyがそのインターフェイスを実装し、BLASライブラリの基礎となる作業を委任するかのようなものです。
このトリックでは、データが常にGPUに残ります。概念実証の概要を概説しますhere
貴重なフィードバックをありがとう。 'Defun'について：関数内で引数はどのような型ですか？ whileループなどの命令的ロジックを実装することはできますか？それとも、テスト例のようにスカラーだけで動作するのでしょうか？そしてそれをGPUオペレーションにどのように変換しますか？場合によっては、CUDAコードを直接記述したいと考えています。 TheanoのようなシンプルなPythonインターフェイスがあるのですか、それともC++ APIを使用する必要がありますか？ – Albert
Defunでは、TensorFlowグラフを複数のオブジェクトにグループ化できます。したがって、TFグラフ操作だけがそこで実行されることができます。 TFグラフは、制限付きwhileループ（副作用やブレーク/継続を許さない）など、いくつかのフロー制御をサポートしています。カスタムCUDAカーネルを作成するには、C APIを使用する必要があります。例は次のとおりです。https://github.com/MycChiu/fast-LayerNorm-TF –