Используйте tf.autodiff.ForwardAccumulator внутри функции, украшенной tf.function (RNN). ⇐ Python

[Anonymous]

В Tensorflow мне нужно вычислить произведение вектора Якобиана с помощью функции tf.autodiff.ForwardAccumulator(). Все работает нормально, если я не украшу функцию @tf.function. В этом случае я получаю эту ошибку:
TypeError: Exception encountered when calling GRU.call().

`dtype` 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, и это отлично работает с декоратором и без него. Я предполагаю, что проблема как-то связана с преобразованием модели в граф.
Вот код, воспроизводящий ошибку в Tensorflow v2.16.1:
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)


Подробнее здесь: https://stackoverflow.com/questions/790 ... f-function