在自定義loss函數中使用tf.while_loop的正確方法是什么？

我打算使用以下功能作為培訓的損失：

import tensorflow as tf

def wrap(dist): 
    return tf.while_loop(
        cond=lambda X: tf.math.abs(X) > 0.5,
        body=lambda X: tf.math.subtract(X, 1.0),
        loop_vars=(dist))


# PBC-aware MSE, period = 1.0 ([0, 1.0])
def custom_loss(y_true, y_pred):
    diff = tf.math.abs(y_true - y_pred)
    diff = tf.nest.flatten(diff)
    diff = tf.vectorized_map(wrap, diff)
    return tf.math.reduce_mean(tf.math.square(diff))

# ...other code for loading data and defining the model

model.compile(optimizer=tf.keras.optimizers.SGD(momentum=0.1),
              loss=custom_loss)

但是我遇到了一堆錯誤信息。由于日志太長，我把它們放在一個要點中：https://gist.github.com/HanatoK/f75fddd82372f499c37279f1128cad7a

上面代碼的等效numpy版本應該是

def wrap_diff2(x, y, period=1.0):
    diff = np.abs(x - y)
    while diff > 0.5 * period:
        diff -= period
    return diff * diff

def custom_loss_numpy(y_true, y_pred):
    diff2 = np.vectorize(wrap_diff2)(y_true, y_pred)
    return np.mean(diff2)

有什么想法嗎？完整的代碼示例在googlecolab上共享：https://colab.research.google.com/drive/1ExVHgyKHQfGcpXvo5ZsuBBmzmHzxUekC？usp=共享