如何解决TensorFlow中tf.function开启jit_compile时与tf.py_function冲突报错

错误原因 #

XLA（jit_compile）编译要求计算图中所有算子都为TensorFlow原生支持编译的算子，而tf.py_function封装的是Python侧的numpy/scipy自定义逻辑，XLA无法识别该类算子，也没有对应的XLA_GPU_JIT设备内核实现，因此触发报错。tf.py_function本身也仅支持CPU运行，即使强行塞入JIT域也无法获得加速收益，还会产生额外的GPU-CPU数据拷贝开销。

可行解决方案 #

方案1：拆分JIT编译域（最推荐，改造成本最低） #

不要将tf.py_function纳入JIT编译的作用范围，把需要XLA加速的纯TensorFlow算子逻辑抽离为独立的tf.function(jit_compile=True)函数，和numpy/scipy逻辑分层调用：

import tensorflow as tf
import numpy as np

def my_numpy_func(x):
    # 原有numpy/scipy逻辑不需要修改
    return np.sinh(x)

# 抽离需要JIT加速的纯TF运算逻辑
@tf.function(jit_compile=True)
def jit_post_process(y):
    return y * y

# 主流程不使用全局JIT编译
def main_flow(input):
    # 若存在前置需要JIT加速的逻辑，也可单独抽离为JIT函数调用
    y = tf.py_function(my_numpy_func, [input], tf.float32)
    # 调用JIT加速的后置逻辑
    return jit_post_process(y)

a = main_flow(tf.constant(1.))
print(a)

该方案可以保证纯TF运算部分获得XLA加速的同时，兼容自定义numpy/scipy逻辑，改造成本极低。

方案2：替换自定义逻辑为TensorFlow原生实现 #

如果你的scipy运算存在TensorFlow原生对应的实现，优先替换为原生算子，即可将整个计算图纳入JIT编译范围，获得最优性能。比如示例中的np.sinh可以直接替换为tf.sinh，不需要再使用tf.py_function包裹。

方案3：自定义XLA算子（仅适合性能要求极高的场景） #

如果自定义numpy/scipy逻辑不可替换，且需要全程跑在GPU上无数据拷贝开销，可以基于TensorFlow自定义Op接口开发支持XLA编译的算子，需要用C++实现算子逻辑并注册对应的XLA设备内核，开发成本较高，适合高频复用的核心自定义运算场景。

注意事项 #

• 不要尝试用tf.numpy_function替代tf.py_function解决该问题，两类封装Python逻辑的算子均不支持XLA编译。
• tf.py_function运行时会自动将GPU张量拷贝到CPU执行，再将结果拷回GPU，若自定义逻辑运算量很小，数据拷贝开销会远大于JIT加速收益，可以考虑直接在CPU侧运行该部分逻辑。

内容的提问来源于stack exchange，提问作者Qiqin Zhan