code: def mean_on_masked(loss, mask, epsilon=1e-8, axis=None):
      
      mask = tf.cast(tf.cast(mask, tf.bool), tf.float32)
      active_pixels = tf.reduce_sum(mask)
      active_pixels = tf_if(tf.equal(active_pixels, 0), epsilon, active_pixels)
      return tf.reduce_sum(loss, axis=axis) / active_pixels
docs: Average a loss correctly when it was masked.
  
      :param loss: A float tensor of shape [batch_size, ...] representing the (already masked) loss to be averaged.
      :param mask: A float tensor of shape [batch_size, ...] representing the mask.
      :param epsilon: Offset of log for numerical stability.
      :param axis: The dimensions to reduce. If None (the default), reduces all dimensions.
                   Must be in the range [-rank(input_tensor), rank(input_tensor)).