Computes negative log probability for the discretized mixture of logistics.

  The distribution of a whole pixel is a mixture of 3-dimensional discretized
  logistic distributions. The 3-D discretized logistic factorizes as 3 1-D
  discretized logistic distributions, one for each channel. It defines

  ```none
  P(X = x)
  = sum_{k=1}^K probs[k] * P(X = x | locs[k], scales[k])
  = sum_{k=1}^K probs[k] * [
      prod_{c=1}^3 DiscretizedLogistic(X[c] = x[c] | means[k][c], scales[k]) ]
  ```

  The means tensor is a linear combination of location parameters and previous
  channels. The discretized logistic distribution assigns probability mass to an
  event P(X=x) via logistic CDFs: P(X <= x + 0.5) - P(X < x - 0.5) for 1 < x <
  254; P(X <= 0.5) for x = 0; and 1 - P(X < 245.5) for x = 255. Instead of
  8-bit inputs, this implementation assumes the events are rescaled to [-1, 1].

  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 true pixel intensities
      rescaled to [-1, 1]. The computation assumes channels is 3.

  Returns:
    A [batch, height, width] tensor of the negative log conditional probability
    of each pixel given all previous pixels.

def discretized_mix_logistic_loss(pred, labels):
  """Computes negative log probability for the discretized mixture of logistics.

  The distribution of a whole pixel is a mixture of 3-dimensional discretized
  logistic distributions. The 3-D discretized logistic factorizes as 3 1-D
  discretized logistic distributions, one for each channel. It defines

  ```none
  P(X = x)
  = sum_{k=1}^K probs[k] * P(X = x | locs[k], scales[k])
  = sum_{k=1}^K probs[k] * [
      prod_{c=1}^3 DiscretizedLogistic(X[c] = x[c] | means[k][c], scales[k]) ]
  ```

  The means tensor is a linear combination of location parameters and previous
  channels. The discretized logistic distribution assigns probability mass to an
  event P(X=x) via logistic CDFs: P(X <= x + 0.5) - P(X < x - 0.5) for 1 < x <
  254; P(X <= 0.5) for x = 0; and 1 - P(X < 245.5) for x = 255. Instead of
  8-bit inputs, this implementation assumes the events are rescaled to [-1, 1].

  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 true pixel intensities
      rescaled to [-1, 1]. The computation assumes channels is 3.

  Returns:
    A [batch, height, width] tensor of the negative log conditional probability
    of each pixel given all previous pixels.
  """

  logits, locs, log_scales, coeffs = split_to_discretized_mix_logistic_params(
      pred)

  # Tile labels to broadcast compute across the mixture dimension.
  batch, height, width, num_mixtures = shape_list(logits)  # pylint: disable=unbalanced-tuple-unpacking
  labels = tf.tile(
      tf.reshape(labels, [batch, height, width, 1, 3]),
      [1, 1, 1, num_mixtures, 1])

  # p(x) = sigmoid((x - means_i + 1/255.)/scale_i) -
  #        sigmoid((x - means_i - 1/255.)/scale_i)
  # for each channel i. The means are linearly parameterized.
  means_0 = locs[..., 0]
  means_1 = locs[..., 1] + coeffs[..., 0] * labels[..., 0]
  means_2 = (
      locs[..., 2] + coeffs[..., 1] * labels[..., 0] +
      coeffs[..., 2] * labels[..., 1])
  means = tf.stack([means_0, means_1, means_2], axis=-1)
  centered_labels = labels - means
  inv_stdv = tf.exp(-log_scales)
  plus_in = inv_stdv * (centered_labels + 1. / 255.)
  min_in = inv_stdv * (centered_labels - 1. / 255.)
  cdf_plus = tf.nn.sigmoid(plus_in)
  cdf_min = tf.nn.sigmoid(min_in)

  # Compute log probability for edge case of 0 (before scaling), 255 (before
  # scaling), and all other cases respectively.
  log_prob_0 = plus_in - tf.nn.softplus(plus_in)
  log_prob_255 = -tf.nn.softplus(min_in)
  prob_event = tf.maximum(cdf_plus - cdf_min, 1e-12)
  log_prob_event = tf.log(prob_event)

  # Robustly select log-prob based on numerical edge-cases: (a) [-1, -1+eps);
  # (b) (1-eps, 1]; (c) NaNs during `tf.gradients` of `tf.select`, which may
  # cause `tf.log(0.)`; (d) p(x) < 1e-5.
  mid_in = inv_stdv * centered_labels
  log_prob_event_approx = (
      mid_in - log_scales - 2. * tf.nn.softplus(mid_in) - np.log(127.5))
  log_probs = tf.where(
      labels < -0.999, log_prob_0,
      tf.where(
          labels > 0.999, log_prob_255,
          tf.where(prob_event > 1e-5, log_prob_event, log_prob_event_approx)))

  # Sum over channels and compute log-probability of each mixture.
  log_probs = tf.reduce_sum(log_probs, -1) + tf.nn.log_softmax(logits, axis=-1)
  output = -tf.reduce_logsumexp(log_probs, axis=-1)
  return output