Splits input tensor into parameters of discretized mixture logistic.

  Args:
    inputs: 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.

  Returns:
    Tuple of unconstrained mixture probabilities, locations, scales, and
    coefficient parameters of the distribution. The mixture probability has
    shape [batch, height, width, num_mixtures]. Other parameters have shape
    [batch, height, width, num_mixtures, 3].

def split_to_discretized_mix_logistic_params(inputs):
  """Splits input tensor into parameters of discretized mixture logistic.

  Args:
    inputs: 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.

  Returns:
    Tuple of unconstrained mixture probabilities, locations, scales, and
    coefficient parameters of the distribution. The mixture probability has
    shape [batch, height, width, num_mixtures]. Other parameters have shape
    [batch, height, width, num_mixtures, 3].
  """
  batch, height, width, output_dim = shape_list(inputs)  # pylint: disable=unbalanced-tuple-unpacking
  num_mixtures = output_dim // 10
  logits, locs, log_scales, coeffs = tf.split(
      inputs,
      num_or_size_splits=[
          num_mixtures, num_mixtures * 3, num_mixtures * 3, num_mixtures * 3
      ],
      axis=-1)
  split_shape = [batch, height, width, num_mixtures, 3]
  locs = tf.reshape(locs, split_shape)
  log_scales = tf.reshape(log_scales, split_shape)
  log_scales = tf.maximum(log_scales, -7.)
  coeffs = tf.reshape(coeffs, split_shape)
  coeffs = tf.tanh(coeffs)
  return logits, locs, log_scales, coeffs