В 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
