<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def mean_on_masked(loss, mask, epsilon=1e-8, axis=None):
 """ Average a loss correctly when it was masked. :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