在Tensorflow中的损失函数中将自定义函数批量应用于每个向量

Question

我正在编写用于神经网络训练的自定义损失函数。我目前正在使用带有tf.keras的Tensorflow 2.0-rc来定义和训练模型。问题是我需要对每个向量批量进行一些计算，并且执行这些操作的代码无法处理张量，而且过于复杂，无法使用Tensorflow进行重写。

是否有一种方法可以对批处理中的每个向量应用一个函数，以将数据提取为一个Numpy数组，然后使用该函数生成另一个数组？

我知道我不能简单地获取批处理大小或将损失函数的张量参数转换为数组。我也这样尝试过tf.map_fn：

def some_non_tf_function(t):
    return np.mean(t)

@tf.custom_gradient
def my_loss(y_true, y_pred):
    def grad(dy):
        # just a dummy
        return tf.identity(y_true), dy*y_true

    return tf.reduce_mean(tf.map_fn(some_non_tf_function, y_pred)), grad

但是它抱怨说(TensorArrayV2Read/TensorListGetItem:0) to a numpy array. 有没有什么方法可以在不编写自定义培训例程的情况下使它起作用？