Conditional substract in tensorflow

I'm trying to calculate a custom MSE in tf.keras, such as:

def custom_mse(y_true, y_pred):
    return tf.reduce_mean(tf.square(y_true - y_pred), axis=-1) 

But I want to calculate the difference y_true - y_pred, except when values of y_true are equal to -99.0.

How could I do this?

What do you want the function to return if y_true == -99.0?

Can you not do this?

def custom_mse(y_true, y_pred):
    #check if the batch of y_true has -99
    if -99 in y_true:
          #do whatever you like with the batch
          return #whatever
    return tf.reduce_mean(tf.square(y_true - y_pred), axis=-1) 

If you meant to average only the errors where y_true is not -99, I guess you can simply do this:

def custom_mse(y_true, y_pred):
        return tf.reduce_mean(tf.square(y_true - y_pred)[y_true != -99], axis=-1) 

Cheers,

Keivan

Another way.

def custom_mse(y_true, y_pred):
    return tf.cond(y_true !=  -99.0, lambda: tf.reduce_mean(tf.square(y_true - y_pred), axis=-1),
                                     lambda: #Add logic)