在@tf.function装饰的RNN函数中使用tf.autodiff.ForwardAccumulator报错

在TensorFlow 2.16.1环境中，使用tf.autodiff.ForwardAccumulator计算雅可比向量积（JVP）时，未用@tf.function装饰函数可正常运行，但添加装饰器后会触发GRU层的类型不兼容错误：

TypeError: Exception encountered when calling GRU.call().

dtype <dtype: 'variant'> is not compatible with 1 of dtype int64.

Arguments received by GRU.call():
• sequences=tf.Tensor(shape=(128, 64, 1), dtype=float64)
• initial_state=None
• mask=None
• training=False

将GRU替换为Dense层后，无论是否使用@tf.function都能正常运行，问题根源在于图转换过程中，GRU层内部的variant类型状态张量与tf.autodiff.ForwardAccumulator的前向微分机制存在兼容性冲突。

复现代码 #

import tensorflow as tf

x = tf.random.normal([128, 64, 1], dtype=tf.float64)
layer1 = tf.keras.layers.GRU(32, return_sequences=True, activation=tf.nn.relu, dtype=tf.float64, use_cudnn=False)
layer2 = tf.keras.layers.GRU(32, return_sequences=True, activation=tf.nn.relu, dtype=tf.float64, use_cudnn=False)
layer3 = tf.keras.layers.Dense(1, activation=tf.nn.relu, dtype=tf.float64)

tangent = tf.ones((128, 64, 1), dtype=tf.float64)

@tf.function
def jac_vec_prod(inp, tangent):
    with tf.autodiff.ForwardAccumulator(primals=inp, tangents=tangent) as acc:
        feature = layer1(inp)
        feature = layer2(feature)
        out = layer3(feature)    
    jvp = acc.jvp(out)
    return jvp

jvp = jac_vec_prod(inp=x, tangent=tangent)

解决方案 #

方法一：改用反向模式计算JVP（推荐）

利用tf.GradientTape的反向模式模拟JVP，这种方式对GRU层的兼容性更好，且无需修改原有模型结构：

import tensorflow as tf

x = tf.random.normal([128, 64, 1], dtype=tf.float64)
layer1 = tf.keras.layers.GRU(32, return_sequences=True, activation=tf.nn.relu, dtype=tf.float64, use_cudnn=False)
layer2 = tf.keras.layers.GRU(32, return_sequences=True, activation=tf.nn.relu, dtype=tf.float64, use_cudnn=False)
layer3 = tf.keras.layers.Dense(1, activation=tf.nn.relu, dtype=tf.float64)

tangent = tf.ones((128, 64, 1), dtype=tf.float64)

@tf.function
def jac_vec_prod(inp, tangent):
    with tf.GradientTape() as tape:
        tape.watch(inp)
        feature = layer1(inp)
        feature = layer2(feature)
        out = layer3(feature)
    # 反向模式下，JVP等价于梯度与tangent的点积
    grad = tape.gradient(out, inp, output_gradients=tangent)
    return grad

jvp = jac_vec_prod(inp=x, tangent=tangent)

方法二：自定义GRU层显式处理状态（适用于必须用前向微分的场景）

如果必须使用tf.autodiff.ForwardAccumulator的前向模式，可以自定义GRU层，显式展开循环并管理状态，避免内置GRU的variant类型张量：

import tensorflow as tf

class CustomGRU(tf.keras.layers.Layer):
    def __init__(self, units, **kwargs):
        super().__init__(**kwargs)
        self.units = units
        self.gru_cell = tf.keras.layers.GRUCell(units, activation=tf.nn.relu, dtype=self.dtype)
    
    def call(self, inputs, initial_state=None):
        if initial_state is None:
            initial_state = self.gru_cell.get_initial_state(inputs=inputs)
        # 显式展开循环，规避内置GRU的variant状态追踪
        outputs = []
        state = initial_state
        for step in tf.range(tf.shape(inputs)[1]):
            x_step = inputs[:, step, :]
            output, state = self.gru_cell(x_step, state)
            outputs.append(output)
        return tf.stack(outputs, axis=1)

x = tf.random.normal([128, 64, 1], dtype=tf.float64)
layer1 = CustomGRU(32, dtype=tf.float64)
layer2 = CustomGRU(32, dtype=tf.float64)
layer3 = tf.keras.layers.Dense(1, activation=tf.nn.relu, dtype=tf.float64)

tangent = tf.ones((128, 64, 1), dtype=tf.float64)

@tf.function
def jac_vec_prod(inp, tangent):
    with tf.autodiff.ForwardAccumulator(primals=inp, tangents=tangent) as acc:
        feature = layer1(inp)
        feature = layer2(feature)
        out = layer3(feature)    
    jvp = acc.jvp(out)
    return jvp

jvp = jac_vec_prod(inp=x, tangent=tangent)

说明 #

• 方法一的反向模式JVP在绝大多数场景下性能足够，且实现简单，是优先选择的方案。
• 方法二需手动处理循环逻辑，会带来一定的性能开销，仅适合必须使用前向自动微分的特殊场景。

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