KerasでTypeError: You are passing KerasTensor...が出た時の対処法(になるかも)

概要

最近TensorFlow/Kerasで機械学習を始めようとAuto Encoderを試してみたのですがfitの際に下の感じのエラーが出てその解決法が探してもなかなか見つからなかったため記事にしました。 著者の環境はMac OS Monterey 12.5 Python 3.8.3 tensorflow 2.9.1です。

   raise TypeError(
        f'You are passing KerasTensor(...), an intermediate Keras symbolic input/output, '
        'to a TF API that does not allow registering custom dispatchers, such '
        'as `tf.cond`, `tf.function`, gradient tapes, or `tf.map_fn`. '
        'Keras Functional model construction only supports '
        'TF API calls that *do* support dispatching, such as `tf.math.add` or '
        '`tf.reshape`. '
        'Other APIs cannot be called directly on symbolic Keras'
        'inputs/outputs. You can work around '
        'this limitation by putting the operation in a custom Keras layer '
        '`call` and calling that layer '
        'on this symbolic input/output.')

　原因と解決

これの原因はvae_lossで外部のTensorを呼び出したのが原因らしいみたいです。

x = Input(shape=(784,))
z_mean = encode(x)[0]
z_log_var = encode(x)[1]
def vae_loss(x, predict):
    return kl_div(z_mean, z_log_var) + metrics.mse(x,predict)
model.compile(optimizer="adam", loss=vae_loss)

#ここでTypeError(...)が発生
model.fit(train,train)

#Tensorを関数内部のxで取り直す
def vae_loss(x, predict):
    z_mean = encode(x)[0]
    z_log_var = encode(x)[1]
    return K.mean(kl_div(z_mean, z_log_var) + original_dim * metrics.mse(x,predict))

model.compile(optimizer="adam", loss=vae_loss)
#今度は成功しました。
model.fit(train,train)

　追記

上記の方法は教師なし学習でtrainとcorrectが等しくないと使えません。
一応解決した後にkerasの公式ブログでのやり方を知りましたのでこちらの使用を推奨します。

reconstruct_loss = original_dim * losses.mse(x,output)
kl_loss = kl_div(z_mean,z_log_var)
vae_loss = K.mean(reconstruct_loss + kl_loss)
vae.add_loss(vae_loss)
vae.compile(optimizer="adam")