def diagonal_normal(outputs, name="normal"):
  """Split outputs into mu and log_sigma and return z."""
  with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
    loc, log_scale = tf.split(outputs, 2, axis=-1)
    scale = tf.exp(log_scale)
    dist = tfp.distributions.Normal(
        loc=loc,
        scale=scale + tf.keras.backend.epsilon(),
        allow_nan_stats=False)
    return dist
Split outputs into mu and log_sigma and return z.
tensor2tensor/layers/transformer_glow_layers_ops.py