Нет градиента в пользовательской функции потерь

Question

Я пытался создать собственную функцию потерь, используя корреляцию Спирмена в тензорном потоке 2.1.0-rc1. Однако я столкнулся с ошибкой «нет градиента». У меня есть несколько целей. Чтобы копейщик не дал nan для ярлыков с одним уникальным значением, я просто пропустил и ничего не сделал. Вот мой код:

import tensorflow as tf
from scipy.stats import spearmanr

def spear_func(a, b):
    return spearmanr(a, b).correlation.astype('float32')

def custom_loss(num_targets):
    def custom_rhos(y_true, y_pred):
        rhos = tf.constant(0, dtype='float32')
        for ind in range(num_targets):
            a = tf.slice(y_true, [0, ind], [-1, 1])
            a = tf.reshape(a, [-1])
            b = tf.slice(y_pred, [0, ind], [-1, 1])
            b = tf.reshape(b, [-1])
            rhos = tf.cond(tf.equal(tf.argmax(a), tf.argmin(a)),
                    lambda: rhos, # <-- Here is the problem
                    lambda: tf.subtract(rhos,
                        tf.compat.v1.py_func(spear_func, [a, b], Tout=tf.float32)))
        return rhos
    return custom_rhos

Сообщение об ошибке:

/tensorflow-2.1.0/python3.6/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py in _filter_grads(grads_and_vars)
   1037   if not filtered:
   1038     raise ValueError("No gradients provided for any variable: %s." %
-> 1039                      ([v.name for _, v in grads_and_vars],))
   1040   if vars_with_empty_grads:
   1041     logging.warning(

ValueError: No gradients provided for any variable:

Я попытался написать свою собственную функцию копья. Но это все равно дает ту же ошибку.

from scipy.stats import rankdata
BATCH_SIZE = 4

def custom_loss(num_targets):
    def custom_rhos(y_true, y_pred):
        rhos = []
        for ind in range(num_targets):
            a = tf.squeeze(tf.slice(y_true, [0, ind], [-1, 1]))
            # a = tf.reshape(a, [-1])
            b = tf.squeeze(tf.slice(y_pred, [0, ind], [-1, 1]))
            # b = tf.reshape(b, [-1])
            rank_a, rank_b = tf.numpy_function(rankdata, [a],  Tout=tf.float32), tf.numpy_function(rankdata, [b],  Tout=tf.float32)
            rho = 1 - 6 * tf.reduce_sum((rank_a - rank_b) ** 2) / (BATCH_SIZE ** 3 - BATCH_SIZE)
            rhos.append(rho)
        return -tf.reduce_sum(rhos)
    return custom_rhos