TPUがサポートしていないOPをTensorboardで調べる

TPUを使うと良く謎のエラーに悩まされます。

RuntimeError: Compilation failed: Compilation failure: Detected unsupported operations when trying to compile graph cluster_666512355060360340[] on XLA_TPU_JIT: Placeholder

どのoperationをTPUがサポートしていて、そうでないのかが良く分からず悩んでました。 Tensorboardを使うと、どのoperationがTPUをサポートしているのかを教えてくれます。

やり方は簡単で、tensorboardのgraphを表示させた状態で、TPU compatibilityをクリックするだけです。

この例ですと6のoperationがサポート対象外です。サポート対象外であってもCPUで動かせる時もあるので全部TPU compatibilityにする必要は無いのですが、エラーの解決やパフォーマンス向上には役に立ちますね。