import tensorflow as tf



def log_sum_exp(x):
  """numerically stable log_sum_exp implementation that prevents overflow."""
  axis = len(x.get_shape()) - 1
  m = tf.reduce_max(x, axis)
  m2 = tf.reduce_max(x, axis, keep_dims=True)
  return m + tf.log(tf.reduce_sum(tf.exp(x - m2), axis))


def log_prob_from_logits(x):
  """numerically stable log_softmax implementation that prevents overflow."""
