Как манипулировать tf.Tensor «столбец за столбцом» (для функции потерь Tensorflow Keras)?

Question

Существует функция потерь, например:

def r2(y_true, y_pred):
        te = tf.reduce_sum(tf.square(tf.math.subtract(y_true, tf.reduce_mean(y_true))))
        ue = tf.reduce_sum(tf.square(tf.math.subtract(y_true, y_pred)))
        R_squared = tf.math.subtract(1., tf.math.truediv(ue, te))
        return R_squared

y_true и y_pred являются двумерными tf.Tensors . Он предназначен для расчета R² метрик. Но для многоколоночных тензоров я хочу, чтобы это вычислялось для каждого столбца отдельно, а затем усреднялись результаты для каждого столбца. Первое измерение - размер пакета.

Я изо всех сил пытаюсь манипулировать столбцами в тензорах. : (

Как обновить эту функцию потерь, чтобы она работала с тензорными столбцами вместо целых тензоров?