code: def dml_loss(pred, labels, weights_fn=_weights_one_third, reduce_sum=True):
    
    real_labels = convert_rgb_to_symmetric_real(labels)
    dml_loss_value = discretized_mix_logistic_loss(pred=pred, labels=real_labels)
    weights = weights_fn(labels)
    loss_num = weights * dml_loss_value
    loss_den = weights_nonzero(weights)
    if reduce_sum:
      loss_num = tf.reduce_sum(loss_num)
      loss_den = tf.reduce_sum(loss_den)
    return loss_num, loss_den
docs: Discretized mixture of logistics loss.
  
    Args:
      pred: A [batch, height, width, num_mixtures*10] tensor of floats
        comprising one unconstrained mixture probability, three means
        (one per channel), three standard deviations (one per channel),
        and three coefficients which linearly parameterize dependence across
        channels.
      labels: A [batch, height, width, channels] tensor of 8-bit pixel
        intensities. The computation assumes channels is 3.
      weights_fn: A function of labels, returning a Tensor of shape
        [batch, height, width] which weights each loss term. Default is to scale
        each loss term by 1/3 so that they capture the average across channels.
      reduce_sum: A boolean, to return scalar loss instead of per position.
  
    Returns:
      Tuple of loss tensors for numerator and denominator, each a scalar if
      reduce_sum else of shape [batch, height, width]. The sum of their divisions
      is the number of nats for each pixel in labels.