def mean_on_masked(loss, mask, epsilon=1e-8, axis=None):
    """
    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)).
    """
    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